UniSD: 대형 언어 모델을 위한 통합 자기 증류 프레임워크

강력한 외부 교사 모델에 의존하지 않고 LLM이 생성한 데이터만으로 모델을 개선하는 Self-Distillation의 한계를 극복했습니다. 데이터의 신뢰성, 표현 정렬, 학습 안정성을 동시에 해결하여 자원 제약 환경에서도 효율적인 모델 고도화가 가능함을 입증했습니다.

기존의 자기 증류(Self-Distillation)는 모델이 스스로 내뱉은 답을 정답으로 삼아 다시 배우는 방식이다. 하지만 LLM은 자유로운 문장 생성을 하기 때문에 스스로 만든 답이 부분적으로 틀리거나 논리적 비약이 섞여 있을 수 있으며, 이를 그대로 학습하면 오류가 강화되는 문제가 발생한다.

UniSD는 이 문제를 해결하기 위해 '내부의 지혜'를 활용한다. 모델이 생성한 하나의 답변에 대해 여러 가지 프롬프트(예: 예시 추가, 지침 변경)를 적용한 여러 명의 '가상 교사'들이 검토하게 한다. 만약 여러 교사가 공통적으로 높은 확률을 부여하는 토큰이라면 신뢰할 수 있는 정보로 판단하고 학습 비중을 높이며, 의견이 갈리는 부분은 신뢰도가 낮은 것으로 간주해 학습에서 제외한다.

또한, 단순히 겉으로 보이는 텍스트 확률 분포만 맞추는 것이 아니라 모델 내부의 특징값(Feature)까지 일치시키고, 학습 과정에서 교사 모델이 너무 빨리 변하지 않도록 천천히 업데이트되는 EMA 방식을 도입했다. 결과적으로 모델은 외부의 도움 없이도 자신의 출력물 중 알짜배기 정보만 골라 학습하며 스스로 진화하게 된다.

UniSD는 세 가지 핵심 축을 기반으로 설계되었다. 첫째는 감독 신호의 신뢰성(Supervision Reliability)으로, Multi-Teacher Agreement를 통해 여러 보조 교사 뷰의 일관성을 측정한다. 각 교사 k가 토큰 t에 대해 부여하는 로그 확률 ℓ_t^k를 입력으로 하여 분산이나 범위를 계산하는 변동성 통계량 A({ℓ_t^k})를 수행해 신뢰도 가중치 w_t를 얻는다. 이 값이 작을수록 교사들 간의 의견이 일치함을 의미하며 학습 시 더 큰 비중을 갖는다.

둘째는 표현 정렬(Representation Alignment)이다. 출력 확률 분포뿐만 아니라 모델의 마지막 레이어 은닉 상태(Hidden State)를 직접 비교하는 Feature Matching 손실 함수 L_feat를 추가했다. 학생 모델의 특징 f_t^θ와 교사 모델의 특징 f_t^* 사이의 L2 거리 제곱을 계산하여 내부 구조의 일관성을 강제한다.

셋째는 학습 안정성(Training Stability)이다. 매 단계마다 학생 모델의 파라미터 θ_n을 지수 이동 평균하여 교사 모델 ¯θ_n = β¯θ_{n-1} + (1-β)θ_n을 갱신한다. 또한 토큰별 발산 값 D_t가 특정 임계값 κ를 넘지 않도록 min(D_t, κ) 연산을 수행하는 Divergence Clipping을 적용하여, 스타일 차이로 인한 비정상적인 그래디언트 발생을 억제한다.

Qwen2.5-7B 모델을 사용한 실험에서 UniSD* 통합 파이프라인은 ScienceQA, MBPP, ToolAlpaca 등 6개 벤치마크 평균 73.3점을 기록하여 원본 모델(67.9점) 대비 5.4점 향상되었다. 이는 기존의 강력한 베이스라인인 GKD(70.5점)보다도 2.8점 높은 수치이다.

Ablation Study 결과, EMA Teacher와 Multi-Teacher Agreement가 단일 컴포넌트로서 가장 큰 성능 향상을 견인하는 것으로 나타났다. 특히 ToolAlpaca와 같은 도구 사용 작업에서는 EMA Teacher가 원본 대비 +16.1점이라는 압도적인 개선을 보였다.

모델 크기별 확장성 실험에서도 0.5B부터 7B까지 모든 규모에서 일관된 성능 향상을 확인했다. 특히 UniSD*는 학습 후에도 원본 모델의 일반적인 생성 능력을 유지하는 지표인 Base-distribution Retention에서 SFT 대비 33.9% 낮은 perplexity를 기록하며 카타스트로픽 포게팅(Catastrophic Forgetting) 현상을 효과적으로 방지함을 증명했다.

UniSD는 온폴리시(On-policy) 궤적 위에서 동작하는 신뢰성 기반 자기 교정 프로세스로 정의된다. 학생 정책 π_θ가 생성한 샘플 ŷ에 대해 1차 교사 π_*가 감독을 제공하고, 다수의 보조 교사가 신뢰도를 평가하는 구조이다. Multi-Teacher Agreement는 토큰 레벨과 시퀀스 레벨의 두 가지 입도로 구현되어 국소적 노이즈와 전역적 불안정성을 동시에 제어한다.

Token-level Contrastive Learning은 긍정적 감독 신호 y+와 부정적 대안 y- 사이의 마진 기반 손실을 사용하여, 모델이 정답에 가까워지는 동시에 그럴듯하지만 틀린 오답과는 멀어지도록 유도한다. 이는 특히 코드 생성과 같이 구조는 비슷하지만 로직이 다른 작업에서 효과적이다.

구현 측면에서 모든 보조 교사 뷰는 동일한 모델 가중치를 공유하며 컨텍스트(Few-shot 예시 등)만 다르게 구성하여 배치 처리함으로써 GPU 메모리 점유율과 지연 시간을 최적화했다. 실험 결과 UniSD*는 단일 교사 방식 대비 약 21~28%의 메모리 오버헤드만으로 훨씬 높은 신뢰성을 확보했다.

본 연구는 주로 싱글 턴(Single-turn) 시나리오에 집중하고 있어, 여러 단계의 의사결정이 필요한 롱 호라이즌(Long-horizon) 에이전트 작업에서의 효과는 아직 충분히 검증되지 않았습니다. 또한 최종 답변의 정오표기 방식의 평가에 의존하고 있어, 중간 추론 단계의 논리적 타당성을 직접적으로 평가하는 세밀한 지표 도입이 향후 과제로 남아있습니다.