기술 요약: 자기 증류 궤적 인식 볼츠만 모델링 (TABOM)

1. 문제 제기

확산 언어 모델 (DLM) 은 autoregressive 모델에 비해 우수한 전역 인식 능력과 매우 병렬적인 토큰 생성과 같은 장점을 제공합니다. 그러나 표준 지도 미세 조정 (SFT) 을 통한 DLM 의 사후 학습은 중요한 학습 - 추론 불일치에 직면해 있습니다:

• 학습: 표준 SFT 는 균일 무작위 마스킹을 사용하는 음의 증거 하한 (NELBO) 목적 함수를 사용하며, 모델이 모든 마스킹된 토큰을 단일 단계에서 동시에 재구성해야 합니다. 이는 모든 토큰을 본질적인 예측 가능성과 관계없이 동등하게 어렵게 취급하는 "균일 유도 편향"을 심어줍니다.
• 추론: DLM 추론은 신뢰도 기반의 다단계 "쉬운 것에서 어려운 것"으로의 궤적을 따릅니다. 토큰은 예측 엔트로피가 가장 낮은 것 (가장 확실한 것) 에서 시작하여 더 불확실한 토큰으로 진행되며 반복적으로 마스킹이 해제됩니다.

딜레마:
최근 자기 증류 궤적(모델이 자신의 디코딩 흔적에서 학습하는 것) 을 사용하여 이 격차를 해소하려는 시도는 제한적인 결과만 낳았습니다.

1. 정답에 대한 표준 SFT(SFT-GT): 도메인 내 성능은 향상시키지만, 분포 외 (OOD) 능력에 대한 심각한 파괴적 망각을 초래합니다.
2. 자기 증류 데이터에 대한 표준 SFT(SFT-SD): OOD 능력을 보존하여 망각을 방지하지만, 도메인 내 이득은 매우 미미합니다.
   본 논문은 자기 증류 궤적이 더 낮은 최적화 장벽을 제공하지만, 표준 NELBO 목적 함수를 단순히 적용하면 궤적 내에 내재된 구조화된 "쉬운 것에서 어려운 것" 정보를 활용하지 못해 하위 최적 성능으로 이어진다고 주장합니다.

2. 방법론: TABOM

저자들은 추론 시간 디코딩 역학과 학습 목적 함수를 정렬하는 사후 학습 프레임워크인 **볼츠만 모델링을 통한 궤적 정렬 최적화 (TABOM)**를 제안합니다.

2.1 이론적 기반: 볼츠만 마스킹 해제 분포

저자들은 이론적으로 쉬운 것에서 어려운 것으로의 디코딩 스케줄 하의 추론 마스킹 해제 분포가 토큰의 이상적인 예측 엔트로피에 대한 볼츠만 분포로 모델링될 수 있음을 증명합니다.

• 명제: 마스킹 해제된 집합 $U$의 확률은 $\exp(-\beta \sum_{r \in U} H_{ideal}(x_r))$에 비례하며, 여기서 $H_{ideal}$은 예측 엔트로피입니다.
• 함의: 이 공식은 모델이 "쉬운"(엔트로피가 낮은) 토큰을 먼저 마스킹 해제하는 것을 선호한다는 유도 편향을 명시적으로 포착합니다. 목표는 모델이 예측한 분포와 이러한 목표 볼츠만 분포 간의 쿨백 - 라이블러 (KL) 발산을 최소화하는 것입니다.

2.2 계산 가능한 대리 모델: 쌍별 순위

전역 샘플링과 모든 가능한 마스킹 해제 상태에 대한 분할 함수 (partition function) 계산이 필요하기 때문에 KL 발산을 직접 최적화하는 것은 계산 불가능합니다. TABOM 은 **쌍별 순위 (Pairwise Ranking)**에 기반한 계산 가능한 대리 목적 함수를 도입합니다:

• 메커니즘: 전역적으로 가능도를 최대화하는 대신, 이 방법은 마스킹 해제된 상태 쌍 간의 상대적 순위를 강제합니다. 궤적에서 토큰 $r$이 토큰 $s$보다 먼저 디코딩된다면 (즉, $r$이 "더 쉽다"는 의미), 모델이 $r$에 대해 예측한 엔트로피는 $s$에 대한 것보다 낮아야 합니다.
• 손실 함수: 로컬 디코딩 윈도우에 대해 쌍별 힌지 손실 (pairwise hinge loss) 이 적용됩니다. 새로 마스킹 해제된 토큰 집합 $\Delta$에 대해, 손실은 "더 쉬운" 토큰의 엔트로피가 "더 어려운" 토큰의 엔트로피보다 충분히 낮지 않은 경우를 패널티로 부과합니다:
  $$L_{rank} = \max(0, h_\theta(r) - h_\theta(s) + \gamma)$$
• 로컬 윈도우: 매우 다른 난이도를 가진 토큰 (예: 첫 번째 단계 vs 마지막 단계) 을 비교함으로써 발생하는 노이즈를 피하기 위해, 순위 매기기는 고정된 로컬 윈도우 크기 $W$(예: $W=32$) 내에서 엄격하게 적용됩니다.

2.3 최종 목적 함수

총 TABOM 손실은 다음을 결합합니다:

1. 궤적 인식 재구성: 이전에 마스킹 해제된 토큰의 맥락을 기반으로 새로 마스킹 해제된 토큰 $\Delta$를 예측하는 표준 교차 엔트로피 손실 (균일 마스킹 대체).
2. 쌍별 순위 정규화: 엔트로피 지형도를 쉬운 것에서 어려운 것으로의 궤적과 정렬시키기 위한 위의 힌지 손실.

$$L_{TABOM} = L_{reconstruction} + \lambda L_{rank}$$

3. 주요 기여

1. 이론적 통찰: 본 논문은 추론 마스킹 해제 분포를 예측 엔트로피에 대한 볼츠만 분포로 공식적으로 정의하고, 표준 NELBO 의 "균일 유도 편향"을 학습 - 추론 불일치의 근본 원인으로 규명합니다.
2. 방법론적 혁신: TABOM 은 계산 불가능한 분할 함수 계산 없이 모델의 확실성 순서를 관찰된 디코딩 궤적과 정렬시키는 쌍별 순위 목적 함수를 도입합니다.
3. 경험적 해결: 이 프레임워크는 "망각 대 미미한 이득" 딜레마를 해결합니다. 파괴적 망각을 효과적으로 완화하면서 분포 외 능력을 보존하는 동시에 상당한 도메인 내 성능 개선을 달성합니다.
4. 진단 지표: 저자들은 디코딩 궤적 along 토큰 수준 불확실성 차이가 보존되는지 정량화하기 위해 **궤적 구별 점수 (Trajectory Discrimination Score, TDS)**를 도입합니다.

4. 실험 결과

실험은 코드 생성(HumanEval, MBPP) 및 수학적 추론(GSM8K, MATH500) 도메인 전반에 걸쳐 두 가지 DLM(Dream-7B 및 LLaDA-8B) 에서 수행되었습니다.

• 성능 향상: TABOM 은 일관되게 기준선 (SFT-GT, SFT-SD, dInfer, T3D) 을 능가했습니다.
  • Dream-7B(코드) 에서 TABOM 은 HumanEval 에서 +7.70, MBPP 에서 +2.60의 향상을 기록한 반면, SFT-SD 는 각각 +1.00, +1.20 의 미미한 향상을 보였습니다.
  • 결정적으로 TABOM 은 OOD 성능을 유지했습니다 (예: GSM8K 는 81.73 으로 높게 유지된 반면, SFT-GT 는 52.33 으로 급감했습니다).
• 견고성: 2 토큰 병렬 디코딩 하에서 TABOM 은 높은 성능을 유지한 반면, SFT-GT 는 심각한 성능 저하를 겪었으며, 이는 TABOM 이 모델의 자연스러운 추론 분포와 정렬되어 있음을 강조합니다.
• 어블레이션 연구:
  • 순위 손실을 제거하면 이득이 제한적이어서 엔트로피 지형도를 정렬하는 데 그 필요성이 확인되었습니다.
  • 로컬 윈도우 없이 전역적으로 순위 매기기를 적용하면 노이즈가 발생하고 성능이 저하되어 로컬 윈도우 제약의 필요성이 검증되었습니다.
• 궤적 구별: TDS 분석은 TABOM 이 디코딩 단계 전반에 걸쳐 예측 엔트로피의 분산을 크게 증가시켰음을 보여주었습니다 (예: HumanEval TDS 는 SFT-SD 의 0.035 에서 TABOM 의 0.711 로 상승). 이는 표준 SFT 가 토큰을 균일하게 취급하는 것과 달리 모델이 성공적으로 쉬운 토큰과 어려운 토큰을 구별하도록 학습했음을 증명합니다.

5. 중요성과 주장

본 논문은 TABOM 이 DLM 을 위한 자기 증류 궤적 활용 방식에 근본적인 변화를 가져왔다고 주장합니다:

• 효율성 이상: 이전 궤적 기반 방법들이 샘플링 단계 압축(가속) 에 초점을 맞췄다면, TABOM 은 이러한 궤적을 진정한 지식 획득 및 성능 향상을 위해 사용합니다.
• 단축이 아닌 정렬: 추론 행동을 볼츠만 분포로 모델링하고 엔트로피 기반 순위를 강제함으로써, TABOM 은 모델의 최적화 지형도를 실제 생성 역학과 정렬시킵니다.
• 파괴적 망각 완화: 이 방법은 표준 SFT 의 일반적인 실패 양상인 사전 학습 능력을 희생하지 않으면서 DLM 의 유효 지식 경계를 새로운 도메인으로 확장하는 데 성공합니다.

저자들은 TABOM 이 확산 언어 모델의 학습 - 추론 불일치에 대한 강력한 해결책을 제공하며, 생성 과정의 본질적인 "쉬운 것에서 어려운 것" 구조를 존중하는 더 효과적이고 효율적인 사후 학습 전략으로 가는 길을 제시한다고 결론지었습니다.