기술 요약: 마스크드 증류(Masked Distillation) – 언어 모델의 사고 사슬(Chain-of-Thought) 내재화

문제 정의

대규모 추론 모델(LRM)은 복잡한 작업(수학, 코드, 계획 등)을 수행할 때 최종 정답을 내놓기 전, 길고 명시적인 중간 토큰의 사고 사슬(Chain-of-Thought, CoT)을 생성함으로써 높은 성능을 달enc합니다. 이러한 추적 과정은 정확도를 높여주지만, 추론 지연 시간, 메모리 사용량(KV-캐시 점유), 서빙 에너지 비용을 크게 증가시키며, 이는 추적 길이에 따라 선형적으로 증가합니다.

최근 연구 결과에 따르면, 이러한 중간 추적 과정의 정확성이 최종 정답의 정확성과 인과적으로 연결되어 있지 않으며, 추적의 길이가 문제의 복잡도와 반드시 상관관계를 갖는 것도 아니라는 사실이 밝혀졌습니다. 이는 다음과 같은 중요한 질문을 제기합니다: 이러한 중간 토큰에 표현된 계산 과정을 언어 모델의 파라미터 내부로 내재화할 수 있는가? 즉, 추론 시점에 명시적인 추론 단계를 생성하지 않고도(또는 훨씬 짧은 추적을 통해) 직접 정답을 낼 수 있게 할 수 있는가 하는 점입니다.

방법론: 마스크드 증류 (Masked Distillation)

저자들은 "생각하는" 교사 모델(Teacher Model)의 추론 과정을 "생각하지 않는" 학생 모델(Student Model)로 내재화하기 위해 설계된 지식 증류 프레임워크인 **마스크드 증류(Masked Distillation)**를 제안합니다.

핵심 프레임워크

이 프레임워크에서 학생 모델은 입력된 질문에 조건화되어 오직 솔루션 토큰만을 예측하도록 훈련됩니다. 반면, 교사 모델은 질문과 그 자체로 생성된 CoT 추적 모두에 조건화되어 피드백을 제공합니다. 목표는 학생의 분포(입력에 조건화됨)와 교사의 분포(입력 및 중간 추론 과정에 조건화됨) 사이의 발산(divergence)을 최소화하는 것입니다.

훈련 목적 함수는 **온-폴리시 역-KL 발산(on-policy reverse-KL divergence)**을 사용합니다:
$$L_{MD} = D_{KL}(\pi_S(\cdot | x, y_{<t}) \parallel \pi_T(\cdot | x, \text{ITs}, y_{<t}))$$
여기서 $\pi_S$는 학생 정책, $\pi_T$는 교사 정책, $x$는 입력, $y_{<t}$는 이전에 생성된 토큰, $\text{ITs}$는 교사의 중간 토큰을 의미합니다.

실험 설정

저자들은 이 프레임워크를 두 가지 설정에서 구현했습니다:

1. 자기 증류 (Self-Distillation): 동일한 모델이 교사(사고 모드)이자 학생(비사고 모드) 역할을 수행합니다.
2. 이중 모델 설정 (Dual-Model Setting): 더 큰 규모의 추론 모델이 교사로서 별도의 더 작은 비사고 학생 모델을 감독합니다.

$\alpha$-접미사 스캐폴드 ($\alpha$-Suffix Scaffold)

전체 내재화와 추론 비용 사이의 트레이드오프를 탐구하기 위해, 저자들은 $\alpha$-접미사 마스크드 증류 변형을 도입했습니다.

• $\alpha = 0$ (완전 마스킹): 학생 모델이 전체 중간 추적 과정을 내재화하여 오직 솔루션만을 출력합니다.
• $\alpha = 1$ (비마스킹): 학생 모델이 정답을 내놓기 전 전체 추적 과정을 재현합니다 (표준 CoT 생성과 동일).
• $0 < \alpha < 1$: 학생 모델은 교사 추적의 $(1-\alpha)$ 부분 접두사(prefix)를 내재화하고, 추론 시에는 나머지 $\alpha$-분율의 접미사(suffix)를 솔루션과 함께 출력하도록 훈련됩니다.

이를 통해 저자들은 중간 토큰을 일부 유지하거나 완전히 내재화할 수 있는 "스캐폴드(scaffold, 비계)"로 취급하여 조절할 수 있습니다.

주요 기여

1. 프레임워크 제안: CoT 추론을 모델 파라미터에 내재화하여 직접 정답 생성을 가능하게 하는 마스크드 증류를 도입했습니다.
2. 스캐폴드 분석: $\alpha$-접미사 파라미터를 정의하고 평가하여, 정확도와 추론 비용 간의 트레이드오프를 조절할 수 있는 통제된 축을 확립했습니다.
3. 포괄적 평가: 다음 두 영역에 대해 광범위한 실험을 수행했습니다:
   • 수학: GSM8K (In-Distribution), MATH-500, AIME-25 (Out-of-Distribution).
   • Countdown: 타겟 범위 및 탐색 깊이 변화(OOD 분할 포함)가 있는 숫자 퍼즐 탐색 작업.
4. 어블레이션 연구 (Ablation Studies): 교사 추적에 대한 지도 미세 조정(SFT) 및 명시적인 "교사 토큰(teaching tokens)"을 스캐폴드로 사용하는 방식과 마스크드 증류(역-KL 사용)를 비교했습니다.

결과

1. 완전 내재화의 작업 의존성

추론을 완전히 내재화($\alpha=0$)할 수 있는 능력은 작업의 종류와 학생 모델이 사전 학습 단계에서 해당 도메인에 노출된 정도에 따라 크게 달라집니다:

• GSM8K (수학): 완전 내재화가 성공적입니다. 자기 증류 설정에서 학생 모델은 **76.0%**의 정확도(교사의 87.9% 대비)를 달성하면서 추론 토큰을 6.5배 줄였습니다. 이중 모델 설정에서는 정확도가 15.69%(베이스라인)에서 45.26%로 향상되었으며, 토큰은 5배 감소했습니다.
• Countdown: 스캐폴드 없이는 완전 내재화가 실패합니다. 자기 증류 설정에서 완전 마스킹된 학생 모델은 **41.7%**의 정확도로 하락했습니다(베이스 모델의 56.2%보다 낮음). 이중 모델 설정에서는 0%에서 34.1%로 개선되었으나, 교사(87.3%) 및 비마스킹 변형(81.1%)에 비해서는 여전히 훨씬 낮았습니다.

2. $\alpha$-접미사 스캐폴드의 효능

적은 양의 스캐폴드(추적의 일부를 출력함)를 제공하는 것은 특히 완전 내재화가 실패하는 작업에서 성능을 크게 회복시킵니다:

• Countdown (자기 증류): $\alpha = 0.3$으로 설정하면 정확도가 41.7%에서 **86.2%**로 상승하여, 교사(87.3%)와 거의 일치하는 성능을 보였으며, 비마스킹 변형 대비 추론 토큰을 약 1.3배 절감했습니다.
• Countdown (이중 모델): $\alpha = 0.5$ 변형은 **84.2%**의 정확도를 달성하여 비마스킹 베이스라인(81.1%)을 능가했으며, 파레토 최적점(Pareto-optimal point)을 형성했습니다.
• GSM8K: 작은 스캐폴드($\alpha=0.3$)는 정확도 격차 대부분을 회복하였고(83.2% vs 87.9%), 비마스킹 변형의 토큰 비용의 약 절반만을 사용했습니다.

3. 분포 외(Out-of-Distribution, OOD) 일반화

• 타겟 범위 변화 (Target-Range Shift): 다양한 스캐폴드 패러다임으로 훈련된 모델들은 타겟 숫자 범위가 변경된 OOD 작업에서도 깔끔하게 전이되었습니다.
• 탐색 깊이 변화 (Search-Depth Shift): Countdown의 자기 증류 설정에서, 스캐폴드 훈련을 받은 학생 모델(특히 $\alpha=0.5$)은 탐색 깊이가 증가했을 때 비마스킹 변형보다 더 나은 일반화 성능을 보였습니다 (32.8% vs 25.7%).
• 이중 모델의 한계: 더 어려운 수학 벤치마크(MATH-500, AIME-25)의 이중 모델 설정에서는 교사와 학생 사이의 용량 차이(capacity gap)가 너무 컸습니다. 이 경우 완전 내재화나 스캐폴딩 모두 OOD 성능을 회복하지 못했으며, 모든 변형이 베이스라인 학생 모델의 수준으로 붕괴되었습니다.

4. 훈련 목적 함수

• 역-KL vs. SFT: 본 논문은 교사 추적에 대한 **지도 미세 조정(SFT)**이 스캐폴딩을 효과적으로 활용하지 못한다는 것을 발견했습니다. Countdown에서 동일한 스캐폴드를 사용했을 때 SFT는 약 38%의 정확도를 보인 반면, 역-KL 마스크드 증류는 약 67%에 도달했습니다.
• 교사 토큰 (Teaching Tokens): 스캐폴드로 명시적인 "교사 토큰"을 제공하는 것이 완전 마스킹보다는 성능을 개선했지만, $\alpha$-접미사 방식보다는 효과가 떨어졌습니다.

의의 및 주장

본 논문은 **접미사 스캐폴딩(suffix scaffolding)**이 정확도와 추론 효율성 사이의 균형을 맞출 수 있는 제어 가능한 메커즘을 제공한다고 주장합니다. 최적의 작동 지점은 특정 작업과 교사-학생 간의 용량 차이에 의해 결정됩니다.

• 내재화는 가능하지만 보편적이지는 않다: 모델은 자신이 익숙한 작업(GSM8K 등)에 대해서는 추론을 내재화할 수 있지만, 새로운 탐색 작업(Countdown 등)에서는 명시적인 추적을 유지하지 않고는 어려움을 겪습니다.
• 비용-성능 트레이드오프: 적은 양의 명시적 추론(예: $\alpha=0.3$)을 제공함으로써, 전체 CoT 생성의 성능을 거의 그대로 유지하면서도 상당한 추론 비용 절감(토큰 1.3배~2배 감소)을 얻을 수 있습니다.
• 일반화: 스캐폴딩은 특정 OOD 시나리오(예: 자기 증류에서의 탐색 깊이 변화)에서 일반화 성능을 향상시킬 수 있으며, 이는 부분적인 추적을 유지하는 것이 특정 맥락에서 완전히 내재화된 모델이나 전체 추적 모델보다 복잡한 탐색 공간을 항해하는 데 도움이 될 수 있음을 시사합니다.

저자들은 완전한 내재화가 효율성을 위한 매력적인 목표이긴 하지만, $\alpha$-접미사 접근 방식이 높은 성능을 유지하면서도 지연 시간을 줄일 수 있는 실용적이고 조절 가능한 대안을 제공한다고 결론짓습니다.