S1 & Knowledge Distillation

최지원·2025년 2월 25일

📌 LLM-S1 : Simple Test-Time Scaling

등장배경

1) OpenAI O1모델의 비공개

O1은 테스트 시점에서 AI가 더 깊이 생각하도록 하는 구조 → 매우 뛰어난 추론 능력
하지만 O1 이 어떤 방식으로 구현되었는지 내부적인 테스트 시점 스케일링 기법 공개X
→ 많은 연구팀이 O1을 복제하려는 시도 (DeepSeek 등장 - 훈련비용 크고 복잡)

2) 더 단순한 방법 필요

복제 시도들은 강화학습과 엄청난 양의 데이터 필요
"꼭 그렇게 복잡한 방법이 아니더라도 간간한 방식으로 O1과 비슷한 성능을 낼 수 있지 않을까?"

3) LLM S1이 제안한 해결책 : 간단한 방법으로 O1 재현

S1 모델은 단 1,000개의 샘플만을 사용하여 지도미세조정을 수행하고 Budget Forcing을 통해 테스트 시점 연산량 조절 방식 제안
→ 많은 데이터 없이도 테스트 시점에서 사고 시간을 조절하면 성능을 높일 수 있다.

Budget Forcing & Test Time Scaling 개념 도입

소규모 데이터셋(S1K) 구축

원래는 59,000개의 질문을 모았는데 품질 + 난이도 + 다양성을 기준으로 1,000개의 고품질 데이터 생성
어려운 문제, 다양한 문제, 질 좋은 문제로 구성해서 AI가 더 좋은 학습을 할 수 있도록 설계
즉, S1K는 Budget Forcing과 Test-Time Scaling을 실험하기 위한 데이터 셋
좋은 데이터 1,000개만 잘 선택하면 추가적인 대량 학습 없이도 성능 상승

Test-Time Scaling

AI가 테스트 시점에서 사용하는 사고시간을 조절하여 성능 최적화

1) 병렬적 스케일링

여러 개의 답을 동시에 생성한 후 가장 좋은 답을 선택하는 방식
ex) 다수결, Best-of-N
테스트 과정에서 여러번의 독립적인 추론-수행
빠르게 여러 답을 만들 수 있지만 논리적으로 더 깊이있는 답을 만드는 과정이 부족할 가능성 존재

2) 순차적 스케일링

하나의 답을 생성하는 과정에서 이전 사고 과정을 활용하여 점진적으로 더 나은 답을 찾아가는 방식
단계적인 사고 과정이 필요한 문제에서 유용하게 적용
이를 최적화하기 위해 Budget Forcing 기법 도입

Budget Forcing(예산강제기법)

OpenAI의 O1 모델처럼 대규모 강화학습을 하지않고도 이러한 효과를 얻을 수 있는지 확인하고자 Budget Forcing이라는 단순한 기법을 제안하여 테스트 시점에서 모델의 사고과정을 강제적으로 조절
Test-Time Scaling을 효율적으로 조절하는 방법 = 실현하기 위한 도구

작동방식

1) 모델의 사고시간을 강제적으로 제한

모델이 설정된 최대 사고 토큰 수를 초과하면 End-Of-Thinking Token을 추가하여 강제로 사고를 중단하고 답변 생성
모델이 불필요하게 긴 추론을 하지 않도록 하면서도 테스트 시점에서 연산량 조절

2) 모델의 사고시간을 인위적으로 연장

모델이 더 깊이 사고하도록 "End-Of-Thinking Token"의 생성을 억제하고 "Wait"를 반복 추가하여 사고시간 연장
이를 통해 모델이 더 많은 경우를 검토하여 답 추출
어려운 문제일수록 충분한 사고 시간을 늘리는 것만으로도 성능이 개선될 수 있음을 실험적으로 입증
→ AI의 사고시간을 강제조절하여 더 깊이 고민하도록 만드는 기법

📌 Knowledge Distillation

업로드중..

지식 증류는 대규모의 복잡한 Teacher(선생님) 모델이 학습한 지식을 Student(학생) 모델에 전달하여 작은 모델이 큰 모델의 성능을 최대한 유지하면서도 효율적으로 동작할 수 있도록 하는 딥러닝 기법

📎 Teacher Model

크고 강력한 모델로 학생 모델을 지도하는 역할
대량의 데이터로 학습되었으며 높은 정확도를 갖고 있음
학생 모델이 학습할 수 있도록 다양한 정보 제공

📎 Student Model

작고 가벼운 모델로 선생님 모델의 성능을 최대한 모방하도록 학습
선생님 모델이 제공하는 지식을 기반으로 학습하여 성능 높임
연산량을 줄이고 실시간 동작이 가능하도록 최적화

필요성

대형 모델은 뛰어난 성능을 보이지만 연산 비용이 높고 실시간 응답이 어렵거나 디바이스에 배포하기 어려움
작은 모델(Student 모델)을 사용하면 속도와 효율성이 향상되지만 성능 저하가 발생
지식 증류를 통해 작은 모델이 큰 모델의 지식을 효과적으로 학습하면서 성능을 유지할 수 있음

종류

1) 전통적인 지식 증류

Teacher 모델이 생성한 Soft Targets을 사용하여 Student 모델 학습
온도 매개변수를 조절하여 Soft Targets을 부드럽게 만듦
학생 모델의 손실 함수는 Hard Target과 Soft Target을 모두 고려

2) 반전 증류
일반적인 지식증류에서는 Teacher 모델이 Student 모델을 지도하지만 반대로 Student 모델이 Teacher 모델을 개선하는 방식

3) 온라인 지식 증류
기존 지식증류는 Teacher 모델이 사전 학습된 상태에서 진행되지만 온라인 지식 증류에서는 Teacher와 Student 모델이 동시에 학습

4) 자체 증류
하나의 모델이 자신의 예측값을 활용하여 학습

📌 S1 모델의 학습 방식

교사 모델의 지식을 학생 모델에게 전달하는 지식 증류 방식 사용

1) 교사 모델(Teacher Model) → Gemini 2.0 Flash Thinking Ecperimental

Google DeepMind의 Gemini 2.0 Flash 모델을 사용하여 학생 모델(s1)을 학습시킴
강화학습을 사용한 강화학습 기반 피드백 방식이 아니라 교사 모델의 지식을 학생 모델로 전달하는 지식 증류 방식 채택

2) 지식증류 방식

Gemini 2.0 Flash의 응답을 학습 데이터로 활용

Gemini 2.0 Flash의 답변을 데이터 셋으로 만들어 S1 모델이 이를 따라 학습하도록 설게됨
단순한 정답을 학습하는 것이 아니라 출력의 패턴, 사고 방식, 문제 해결 방식까지 그대로 학습하도록 설계

추론 능력 전수

Budget Forcing은 제한된 리소스 내에서 더 깊은 사고를 하도록 유도하는 기법
이를 통해 S1 모델이 적은 연산량으로도 효과적으로 reasoning을 수행할 수 있도록 함
작은 모델에서도 단계별 사고 과정을 유지하도록 학습

학습 과정

1) 데이터 준비

Gemini 2.0 Flash Thinking Experimental 모델의 답변을 대량으로 생성하여 데이터셋 구축
특정 문제에 대한 답변을 저장하고 이를 S1 모델이 학습 데이터로 사용

2) 모델 학습(지식 증류 과정)

교사 모델(Teacher) : Gemini 2.0 Flash
학생 모델(Student) : S1
과정
- Gemini 2.0 Flash의 응답을 정제하여 데이터셋 구축
- 학생 모델(S1)이 이 데이터셋을 학습하면서 출력 패턴 모방
- Budget Forcing을 적용하여 연산량을 줄이면서도 reasoning을 강화

3) 평가 및 최적화