어디에 증류할지 학습하여 결정한다: BiToK-SD가 Top-K 토큰 선택을 이중 레벨 최적화로 해결하는 메커니즘

TL;DR

on-policy 자체 증류(OPSD 계열)는 외부의 대규모 교사 모델 없이도 소형 모델의 추론 능력을 향상시킬 수 있습니다. 하지만 손실은 생성된 토큰 전체에 균일하게 적용됩니다. 탐색을 위한 토큰이든, 정정을 위한 토큰이든, 답 도출 과정이든 모두 같은 강도로 처리되는 식입니다. 고정 휴리스틱(엔트로피 또는 교사-학생 간의 KL)으로 상위 토큰을 선택하는 TIP과 같은 방법도 있지만, 선택 기준이 훈련 전반에 걸쳐 고정되어 있고 게다가 선택된 위치는 모두 동일한 가중치를 갖습니다.

BiToK-SD는 '어디에, 얼마나 강하게 증류할지' 자체를 학습 대상으로 삼습니다. Top-K 선택을 미분 가능한 임계값 완화(threshold relaxation)로 정식화하고, 이중 레벨 최적화(bi-level optimization)를 통해 하위 문제에서 적응형 임계값을 찾습니다.

Qwen3-1.7B/4B/8B × AIME24/25・HMMT25 (Avg@64)에서 BiToK-SD는 모든 스케일에서 평균 최고 성능을 보였습니다. OPSD 대비 +0.63/+0.58/+0.21의 향상을 기록했습니다. 반면, 동일한 Soft-OR 점수를 사용하더라도 하드 마스크로 선택하는 TIP은 9개 셀 모두에서 OPSD보다 낮은 성능을 보였는데, 이는 '선택' 그 자체가 함정이며, '학습된 소프트 선택'이 효과적임을 의미합니다.

학습된 셀렉터는 (1) 학생의 엔트로피를 높게 유지하여 탐색(exploration)을 보존하고, (2) wait/hmm/maybe와 같은 반사 마커 위치에 평균 1.63배의 가중치를 할당하며, (3) 토큰 수는 오히려 1.6%~4.8% 감소시킵니다.

핵심 통찰: 토큰 선택은 훈련 전에 결정하는 고정 하이퍼파라미터가 아니라, 증류 목적 함수와 함께 최적화해야 할 학습 대상입니다.

배경: 균일한 증류의 비효율성과 고정 규칙의 구식화

on-policy 자체 증류의 부상

대규모 LLM의 추론 능력을 소형 모델로 옮기는 지식 증류(Knowledge Distillation, KD)는 주로 off-policy와 on-policy 두 계통으로 나뉩니다. off-policy는 고정된 교사 출력이나 수집된 사고 궤적에 의존하지만, 학생의 분포와 궤적의 분포가 어긋나면 성능이 정체되기 쉽습니다.

on-policy KD는 학생 스스로 생성한 궤적에서 교사 신호를 만들기 때문에, 학생의 진화하는 분포를 따라갈 수 있습니다. 더 발전된 형태가 **자체 증류(self-distillation)**로, 외부의 거대한 교사 모델을 필요로 하지 않고 모델 자체의 생성물로부터 학습합니다. OPSD (Self-Distilled Reasoner)는 이러한 흐름의 대표적인 예시로, 피드백 조건부의 자기 지도 분포를 사용하여 생성 궤적의 각 토큰 위치에서 KL(Kullback–Leibler) 다이버전스를 최소화합니다.

하지만 OPSD 계열은 모든 위치에 동일한 손실을 적용합니다. CoT(Chain-of-Thought) 과정에서는 시도와 오류 정정, 반사, 답 도출 등 토큰의 역할이 균일하지 않습니다. 관련 없는 위치까지 증류하면 노이즈가 되고, 반대로 중요한 위치를 놓치면 기회 손실이 됩니다.

고정 휴리스틱의 한계

여기서 선택적 토큰 증류(selective token distillation)가 등장합니다. TIP은 Soft-OR 점수—학생 엔트로피와 교사-학생 간의 역 KL을 결합한 지표—를 사용하여 상위 50% 위치를 하드 마스크로 지정하여 증류합니다.

하지만 이 접근 방식에는 두 가지 문제가 있습니다.

• 고정 기준: Soft-OR 점수는 정의되는 시점에 고정되며, 훈련 중 학생의 분포 변화에 적응하지 못합니다. 같은 위치라도 훈련 초기와 후기에는 유용도가 다를 것입니다.
• 이진 동일 가중치: 선택된 위치에는 모두 동일한 가중치가 부여됩니다. 중요도의 차이를 무시하고 있습니다.

결과적으로, TIP은 OPSD보다 낮은 성능을 보입니다. 본 논문의 핵심은 토큰 선택 그 자체를 학생과 함께 최적화함으로써 이 두 가지 문제를 해결하는 데 있습니다.

![BiToK-SD の全体構造](/images/bitok-sd-bilevel-topk-self-distillation/fig1.png)

방법 상세: 이중 레벨 최적화로서의 토큰 선택

전제: on-policy 자체 증류의 위치별 손실

OPSD와 마찬가지로, BiToK-SD는 학생 [생략된 내용] 위치별 손실을 사용합니다.

위의 내용을 모든 위치에 동일한 가중치로 합산하는 것이 OPSD입니다. BiToK-SD는 여기에 위치별 가중치를 추가합니다.

유틸리티 점수: 무엇을 '증류하기 쉬운지' 보는가

먼저, 각 위치가 증류에 얼마나 적합한지를 추정합니다. 기본적으로 Soft-OR 점수를 사용합니다.

여기서 [생략된 내용]

하지만 이 점수는 어디까지나 입력 특징입니다. 후술하듯이, 학생 엔트로피, 교사 엔트로피, 순 KL, 역 KL 중 무엇을 사용해도 프레임워크는 성립됩니다.

미분 가능한 Top-K 임계값 완화

BiToK-SD는 이 임계값을 [생략된 내용]로 설정합니다.

첫 번째 항은 [생략된 내용]

마스크가 바뀌면 학생의 손실도 바뀐다. (Mask changes, student loss changes)

소프트 마스크와 가중치 증류 손실

최종적인 손실은 정규화된 가중치 합이 됩니다.

이를 이중 레벨 최적화로 표현하면 다음과 같습니다.

강조되는 위치와 그 강도가 훈련에 따라 동적으로 변화합니다. 이것이 고정 휴리스틱과의 결정적인 차이점입니다.

실험 결과: 고정 선택은 부정적, 학습된 소프트 선택은 긍정적

실험 설정

모델은 Qwen3-1.7B/4B/8B를 사용한다. 데이터는 OpenThoughts의 수학 추론 서브셋(약 30K 문제)이다. 모든 기법을 동일한 체크포인트에서 200 스텝 동안 LoRA (r=64, α=128)로 학습시킨다. 평가는 AIME24, AIME25, HMMT25를 thinking 모드로 진행하며, 각 문제당 64 샘플의 평균 정확도(Avg@64)를 최종 체크포인트에서 보고한다.

비교 기법은 SFT, GRPO, OPSD, TIP 네 가지이다. OPSD는 모든 위치에 동일한 가중치로 증류한다. TIP는 Soft-OR 상위 50%를 하드 마스크하여 역KL(reverse KL)을 증류한다.

주요 결과

SFT는 모든 스케일에서 베이스보다 성능이 저하된다. 고정된 정답 궤적을 그대로 모방하는 것은 on-policy 추론 설정에서는 효과가 없음을 보여준다. GRPO는 1.7B에서 약간 높지만, 4B와 8B에서는 베이스 수준이거나 약간 낮아, 아웃컴 보상만으로는 거친 교사 신호에 불과하다.

OPSD가 가장 강력한 베이스라인이지만, BiToK-SD는 모든 스케일에서 이를 능가한다. OPSD 대비 1.7B+0.63, 4B+0.58, 8B+0.21점의 개선이다.

대조적으로, 동일한 Soft-OR 점수를 사용하는 TIP는 9개의 스케일×벤치마크 셀 전체에서 OPSD보다 낮은 성능을 보인다. 이는 매우 시사적이다. 선택을 '해야 하는 것'이 아니라, 어떻게 선택할지가 성능을 가른다.

분석: 증류가 변화시키는 행동

탐색의 보존: BiToK-SD는 학습을 거치면서 학생(student)의 엔트로피를 높게 유지하며, 교사(teacher)와의 엔트로피 갭도 크다. 이는 학생 분포가 너무 뾰족해지지 않고 탐색을 유지하고 있음을 의미한다. OPSD는 모든 위치를 교사에 가깝게 만들려고 하면서 분포가 뾰족해지기 쉬우며, 이것이 차이로 이어진다.

반사 행동의 유도: 'wait', 'hmm', 'maybe' 등의 인식적 마커(epistemic markers)를 경량 지표로 계산했을 때, BiToK-SD는 학습 후반부에 OPSD나 TIP보다 더 많은 반사 토큰을 생성한다. 셀렉터는 이러한 마커가 포함된 위치에 다른 위치의 평균 1.63배 가중치를 할당하고 있다.

길어짐이 아니다: AIME24의 64 샘플 평가에서, BiToK-SD는 OPSD 대비 1.7B-3.7%, 4B-1.6%, 8B-4.8%로 토큰 수를 줄인다. 4-gram 반복률과 중단율도 모든 스케일에서 감소한다. 반사 마커가 증가하는 것은 단순히 길어졌기 때문이 아니라, 추론 과정의 질이 변화했기 때문이다.

마스크 밀도의 유지: 연속 가중치임에도 불구하고 전체 밀도는 목표

어베레이션 (Ablation)

Top-K 비율: 55.99→54.95→54.53으로 역U자 형태를 보이며, 중간의 50%에서 정점을 찍는다. 극단적으로 너무 적거나 많으면 효과가 작아진다.

utility 점수: k=50%로 고정하여 점수를 교체했을 때, Soft-OR이 가장 높고(55.99), 순KL(forward KL)이 그 다음(55.78), Uniform(55.10), 학생 엔트로피(student entropy)(54.79), 교사 엔트로피(teacher entropy)(54.69), 역KL(reverse KL)(53.85) 순이다. 순KL 계열이 유효하며, 역KL이나 엔트로피 단독으로는 오히려 손해를 본다.

마스크 온도: [내용 없음]

임계값의 동태: [내용 없음]

off-policy로의 일반화 (Generalization to Off-Policy)

BiToK-SD 프레임워크는 on-policy에 국한되지 않는다. Qwen3-8B→1.7B, FineWeb 80M 토큰, 교사 분포를 top-64로 캐시하는 off-policy 설정에서는 Full-KD(모든 위치), SE-KD(학생 엔트로피 선택), SE-KD3X(위치・클래스・샘플 축으로 선택) 모두를 능가했다.

특히 PPL 개선(6.94 vs 7.30)이 크며, 교사 분포의 일부만 캐시할 수 있는 설정에서도 적절하게 증류 대상을 선택함으로써 왜곡을 억제할 수 있음을 보여준다.

고찰: '학습된 선택'이 승리하는 이유

하드 선택이 실패하는 이유

TIP의 실패는 OPSD와 비교했을 때 5.14점(1.7B)에 달한다. 이는 선택 규칙(하드 마스크)과 발산 정도(역 KL, inverse KL)가 모두 영향을 미치지만, 핵심은 다음 두 가지다.

첫째, 고정된 Soft-OR 점수는 학생 분포의 변화를 따라가지 못한다. 훈련 초기에 유용해 보였던 위치가 후기에는 이미 학생에게 자명해져 증류 효과가 약해진다. 반대로, 초기에는 점수가 낮았던 위치가 후기에는 중요한 오차를 포함하게 될 수도 있다.

둘째, 이진 마스크는 선택된 위치에 동일한 가중치를 부여한다. 예를 들어, 교사와 학생의 의견이 거의 일치하는 위치와 크게 차이가 나는 위치를 같은 강도로 증류하는 것은 비효율적이다. BiToK-SD의 소프트 마스크는 유틸리티(utility)가 임계값에서 얼마나 벗어났는지에 따라 가중치를 변경할 수 있다.

소프트 선택이 탐색을 보존한다

모든 위치에 일률적으로 증류하면, 학생은 교사 분포에 가까워지려고 엔트로피를 낮추기 쉽다. 특히 피드백 조건부 자기 지도 학습(feed-back conditional self-supervised learning)은 학생 자신의 오류를 수정한 '더 나은 분포'이지만, 그럼에도 완전한 모방은 탐색의 여지를 줄인다.

BiToK-SD는 엔트로피가 높고 교사와의 차이가 큰 위치(즉, 학생이 아직 불확실하고 배울 여지가 있는 위치)를 우선적으로 증류한다. 그 결과, 학생은 필요한 위치에서 교사로부터 배우면서도 전체적으로 넓은 분포를 유지한다. 이는 탐색과 활용의 균형을 맞추는 강화학습(RL)의 발상과 공통점이 있다.

반성 행동으로 재분배

추론 모델에서 wait나 hmm 같은 인식적 마커(cognitive marker)는 '자신의 생각을 재검토하고 있다'는 상태를 나타낸다. BiToK-SD는 이러한 위치를 더 무겁게 증류함으로써, 모델에 자기 교정 습관을 강하게 새겨 넣는다.

흥미로운 점은 그 결과로 토큰 수가 늘어나는 것이 아니라 줄어든다는 것이다. 무분별한 추론의 반복이나 탈선이 줄어드는 대신, 질 높은 반성이 증가한다. 이는 단순한 지식 복사가 아니라, 추론 스타일의 재분배라고 할 수 있다.

이중 레벨 완화의 계보

Top-K나 랭킹(ranking)의 미분 가능 완화는 추천이나 검색에서 연구되어 왔다(PiRank, NeuralNDCG, X-risk 최소화 등). BiToK-SD는 이를 LLM의 토큰 수준 증류에 전용했다. 앞으로 중요도의 척도를 보상(reward)이나 불필요한 위치 제거로 확장할 여지가 크다.

한계

개선은 견실하지만, 1.7B에서 0.63점, 8B에서 0.21점으로 작다. 평가 역시 수학 추론에 국한되어 있다. LoRA 200 스텝이라는 단기 설정이며, 풀 파인튜닝(full fine-tuning)이나 장기 훈련에서 같은 경향이 나타날지는 불분명하다. 임계값 업데이트는 매 스텝의 엄밀 해가 아닌 근사해이므로, 더 고도화된 내점 솔버(interior point solver)를 사용하면 더욱 달라질 가능성도 있다.

관련 연구

OPSD (Zhao et al., 2026): on-policy 자기 증류의 기반. BiToK-SD는 Uniform weighting 설정을 통해 OPSD와 일치한다.

TIP (Xu et al., 2026): Soft-OR 점수로 상위 50%를 하드 마스크 처리한다. BiToK-SD와 같은 중요도 점수를 사용하지만, 선택이 고정적이고 이진적이다. 본 논문의 결과는 TIP의 선택 방식이 아닌 '학습 가능한 가중치'의 중요성을 뒷받침한다.

Tavor et al. (2026): 선택적 KD를 재고하여 위치 축 외에 클래스 축, 샘플 축에서도 선택하는 SE-KD3X를 제안했다. off-policy 실험에서 BiToK-SD가 이를 능가한다.

X-risk 최소화・ranking 완화 (Yang, 2022; Swezey et al., 2021; Pobrotyn & Białobrzeski, 2021; Qiu et al., 2022): Top-K나 NDCG의 미분 가능 완화, bilevel 정식화의 계보. BiToK-SD는 이를 토큰 수준 증류 가중치 학습에 응용했다.

on-policy KD 일반 (Agarwal et al., 2024; Shao et al., 2024; Ko et al., 2026): 학생 분포에 맞춘 증류의 중요성을 보여준다. BiToK-SD는 그 연장선상에서, 증류 신호를 어느 위치에 배치할지도 분포에 따라 최적화한다.

본 연재의 관련 기사: UECR-GRPO(증류와 GRPO의 통합), AMPLE-Math(OPSD의 증류 개선과 특권 정보 분리), RISE(RLVR 궤적에서 순차적 교사 구축) 등과도 연결된다. 증류 연구의 흐름은 데이터 → 교사 신호 → 위치 선택으로 단계적으로 정교화되고 있다.

요약 (Summary)

BiToK-SD가 주장하는 바는 명확하다: 증류해야 할 토큰은 고정된 규칙으로 자르는 것이 아니라, 학습을 통해 결정한다.

문제: on-policy 자체 증류의 손실은 모든 위치에 균일하다. 고정 휴리스틱 기반의 하드 Top-K는 학생 분포의 진화에 따르지 못할 뿐만 아니라, 이진(binary) 등가 가중치로 정보를 손상시킨다.
해결책: Top-K 임계값을 미분 가능한 완화 문제로 만들고, 이중 레벨 최적화(bilevel optimization)를 통해 $ heta$와 $
ho$를 동시에 진행한다. 소프트 마스크(soft mask)를 사용하여 '어디에', '얼마나 강하게' 증류할지 연속적으로 학습한다.
결과: Qwen3-1.7B/4B/8B에서 모든 스케일 최고 성능을 달성했다. 동일한 중요도 점수에서도 하드 선택은 열등하고, 학습된 소프트 선택은 우세하다.
행동 양식: 탐색(exploration)을 보존하고, 반사 마커 위치를 1.63배 더 무겁게 증류하여 장문화 없이 추론 스타일을 변경한다.

물론 평가 도메인이나 개선 폭에는 한계가 있다. 하지만 '선택하는 법을 학습한다'는 방향성은 증류뿐만 아니라 RL(강화학습), active learning, 데이터 선택 등 다양한 분야에서 활용될 것이다. 다음으로 모델을 작게 만들 때, 무엇을 모방할지 모델 스스로에게 맡기는 것—이러한 발상이 추론 모델의 증류를 다음 단계로 끌어올릴 것이다.

참고 (References)

• Liang, H., Zhang, X., & Gao, H. (2026). Learning What to Distill: Bilevel Top-K Token Selection for Self-Distillation in Large Language Models. arXiv:2610.07247.
• Zhao, S. et al. (2026). Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models. arXiv:2601.18734.
• Xu, Y. et al. (2026). TIP: Token Importance in On-Policy Distillation. arXiv:2604.14084.
• Tavor, A. et al. (2026). Rethinking Selective Knowledge Distillation. arXiv:2602.01395.
• Swezey, R. et al. (2021). PiRank: Scalable Learning to Rank via Differentiable Sorting. NeurIPS 34.
• Pobrotyn, P., & Białobrzeski, R. (2021). NeuralNDCG: Direct Optimisation of a Ranking Metric via Differentiable Relaxation of Sorting. arXiv:2102.07831.
• Yang, T. (2022). Algorithmic Foundations of Empirical X-Risk Minimization. arXiv:2206.00439.
• Qiu, Z. et al. (2022). Large-Scale Stochastic Optimization of NDCG Surrogates for Deep Learning with Provable Convergence. arXiv:2202.12183.
• Agarwal, R. et al. (2024). On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes. ICLR.
• Shao, Z. et al. (2024). DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models. arXiv:2402.03300.

토론 (Discussion)