기술 요약: 확산형 LLM의 사후 학습 양자화를 위한 FAIR-Calib

1. 문제 정의

확산형 대규모 언어 모델(dLLM)은 자기회귀적 디코딩에서 벗어나 반복적인 마스크 정제(iterative masked refinement)와 양방향 어텐션(bidirectional attention)을 활용하여 텍스트를 생성하는 패러다임의 전환을 나타냅니다. 이러한 메커니즘은 유연한 인필링(infilling) 및 전역적 수정 능력을 제공하지만, 사후 학습 양자화(PTQ)에 있어 독특한 과제를 안겨줍니다.

식별된 핵심 문제는 dLLM 디코딩에 내재된 **"안정성 지연(stability lag)"**입니다. dLLM에서 토큰은 모델의 후험적 믿음(posterior belief)이 이후의 확산 단계에서도 계속 진화하더라도, 일단 컨텍스트에 기록되면 되돌릴 수 없이 확정(commit)됩니다. 저자들은 다음을 관찰했습니다:

• 취약한 확정 상태(Fragile Commit States): 많은 위치에서 토큰이 되돌릴 수 없이 확정되는 순간과 모델의 Top-1 예측이 안정화되는 순간 사이에 상당한 지연이 발생합니다.
• 오류 증폭(Amplification of Errors): 표준 PTQ는 양자화 노이즈를 도입하여 "쓰기 프런티어(write frontier)"(확정의 순간)에서의 경계선상의 결정을 뒤집을 수 있습니다. 이러한 결정은 되돌릴 수 없기 때문에, 단 하나의 잘못된 토큰 확정은 이후 모든 단계의 컨디셔닝 컨텍스트가 됩니다.
• 순차적 발산(Sequential Divergence): 오류를 수정할 수 있는 자기회귀 모델과 달리, dLLM에서의 잘못된 확정은 모델이 잘못된 컨텍스트를 바탕으로 후속 토큰을 정제하도록 강제합니다. 이는 오류의 점진적이고 단계적인 증폭을 초래하여 생성 품질을 심각하게 저하시킵니다.

자기회귀 모델을 위해 설계된 기존 PTQ 방식은 이러한 비가역적 확정 메커니즘을 고려하지 못하며, 이로 인해 특히 추론 작업에서 성능 저하를 야기합니다.

2. 방법론: FAIR-Calib

저자들은 비용이 많이 드는 엔드 투 엔드 확산 롤아웃(diffusion rollouts)을 수행하지 않고도 취약한 프런티어 상태를 보호하도록 설계된 2단계 PTQ 프레임워크인 FAIR-Calib(Frontier-Aware Instability-Reweighted Calibration)를 제안합니다.

1단계: 교사 프로빙 (Teacher Probing, 사전 추정)

첫 번째 단계는 풀 프리시전(FP) 교사 모델을 프로빙하여 위치 인식 중요도 사전($\bar{w}$)을 추정합니다.

• 메커니즘: dLLM의 훈련 역학에 맞추고 마스크 상태 공간에 대한 정책 불가지론적(policy-agnostic) 커버리지를 보장하기 위해, 교사 모델을 무작위 확정(stochastic masking/unmasking) 상태로 실행합니다.
• 가중치 구축: 사전 $\bar{w}$는 확산 단계($t$)와 토큰 위치($i$)에 따른 두 가지 신호의 가산적 누적입니다:
  1. 프런티어 히트 신호(Frontier-Hit Signal): 위치가 확정되는 순간($i \in \hat{C}_t$)에 해당 위치의 가중치를 높입니다. 이는 쓰기 작업의 구조적 비가역성을 포착합니다.
  2. 마스크 단계 신뢰도(Masked-Stage Reliability): 위치가 여전히 마스킹된 상태일 때 교사의 신뢰도(예: 토큰 확률 또는 음의 엔트로피)에 기반한 항입니다. 이는 교사가 본질적으로 모호한 위치의 가중치를 낮추는 신뢰도 게이트 역할을 하여 사전 추정의 노이즈를 줄입니다.
• 시간 재가중치(Time Reweighting): 더 많은 후속 정제 단계에 영향을 미치는 초기 확정을 강조하기 위해 "얼리 부스트(early-boost)" 스케줄($\lambda_0(t)$)이 적용됩니다.

2단계: 오프-폴리시 정적 가중 캘리브레이션 (Off-Policy Static Weighted Calibration)

두 번째 단계는 고정된 사전 $\bar{w}$를 사용하여 양자화된 모델에 대해 효율적인 레이어별 캘리브레이션을 수행합니다.

• 오프-폴리시 대리물(Off-Policy Surrogate): 계산 비용이 매우 높은 전체 확산 궤적을 양자화된 모델로 롤아웃하는 대신, 완전히 관측된 실제 토큰(마스크 없음)에 대한 교사 강제(teacher-forcing) 방식을 사용합니다.
• 목적 함수: 캘리브레이션은 양자화된 모델과 FP 교사 간의 **위치 가중 은닉 상태 평균 제곱 오차(MSE)**를 최소화합니다:
  $$\min_{\theta_\ell} \mathbb{E}_{(x,y)} \left[ \sum_{i=1}^N \bar{w}_i \| h^q_{\ell,i}(x, y) - h^\star_{\ell,i}(x, y) \|_2^2 \right]$$
  여기서 $\bar{w}_i$는 1단계에서 식별된 취약한 프런티어 상태의 보호를 우선시합니다.
• 아핀 평탄화(Affine Flattening): 저비트 설정(예: W4A4)에서 흔히 발생하는 아웃라이어 및 분포 불일치를 해결하기 위해, 양자화 전 가중치와 활성화 분포를 평탄화하는 학습 가능한 아핀 변환(FlatQuant와 유사)을 포함합니다.

3. 이론적 근거

본 논문은 가중 은닉 상태 MSE가 출력 쿨백-라이블러(KL) 발산($KL(\mu^\star \| \mu^q)$)을 최소화하기 위한 대리물임을 정당화하는 이론적 유도를 제공합니다.

• 궤적 분해(Trajectory Decomposition): 모델 독립적인 무작위 확정 가정 하에, 출력 KL 발산은 단계별 커널 발산의 합으로 분해되는 궤적 KL 발산에 의해 상한이 결정됩니다.
• 희소 기여(Sparse Contribution): 결정적으로, 단계별 발산은 해당 단계에서의 확정된 위치에만 의존함을 보여줍니다.
• 대리 상한(Surrogate Bound): 소프트맥스 함수의 매끄러움과 서픽스 네트워크(suffix network)의 립시츠 연속성을 활용하여, 토큰 수준의 KL 발산은 제곱 로짓 오차에 의해 상한이 결정되며, 이는 다시 제곱 은닉 상태 오차에 의해 상한이 결정됩니다.
• 결론: 확정 빈도 및 타이밍에 해당하는 가중치를 사용하여 가중 은닉 상태 MSE를 최소화하는 것은 출력 발산에 대한 원칙적인 상한을 제공합니다.

4. 실험 결과

저자들은 LLaDA 및 Dream 모델 제품군에 대해 다양한 벤치마크 세트(PIQA, BoolQ, WinoGrande, ARC, MMLU, HumanEval, GSM8K 등)를 사용하여 W4A4(4비트 가중치, 4비트 활성화) 양자화를 통해 FAIR-Calib을 평가했습니다.

• 성능 향상: FAIR-Calib은 최신 베이스라인(RTN, QuaRot, FlatQuant)을 일관되게 능가합니다.
  • LLaDA-Base에서 FAIR-Calib은 평균 정확도 **61.09%**를 달사하여 FlatQuant(59.37%)와 RTN(49.13%)을 앞질렀습니다.
  • Dream-Base에서 **64.64%**를 달성하여 FlatQuant(62.08%)와 RTN(42.00%)을 크게 상회했습니다.
• 메커니즘 진단:
  • 감소된 플립(Reduced Flips): FAIR-Calib은 "교사 강제 확정-단계 플립(teacher-forced commit-step flips)"을 유의미하게 감소시켜, 시퀀스당 평균 플립 횟수를 2.9에서 1.9로 낮췄습니다.
  • 오류 억제(Error Suppression): 확정 후의 순차적 오류 증폭을 효과적으로 억제하여, 베이스라인에 비해 더 평탄한 확률-MSE 궤적을 생성합니다.
  • 확정 후 불일치(Post-Commit Mismatch): "평균 불일치(mean-disress)" 및 "결코 일치하지 않음(never-agree)" 비율을 모두 줄여, 확정된 토큰과 모델의 후속 예측 사이의 정렬을 개선했음을 보여줍니다.
• 절제 연구(Ablation Studies):
  • "프런티어 히트(frontier-hit)"와 "마스크 단계 신뢰도(masked-stage reliability)" 구성 요소 모두 성능에 기여하며, 두 조합이 최상의 결과를 냅니다.
  • "얼리 부스트(early-boost)" 시간 가중 스케줄은 균등 또는 레이트 부스트(late-boost) 스케줄보다 우수함이 입증되었습니다.
  • 프로빙된 사전 $\bar{w}$는 코퍼스 간 순위 일관성이 강함을 보여주며, 이는 이 방법이 코퍼스 특이적 통계가 아닌 고유한 디코딩 역학을 포착하고 있음을 시사합니다.
• 효율성: 이 방법은 더 긴 시퀀스나 온-폴리시 상태 수집이 필요한 베이스라인에 비해 캘리브레이션 GPU 메모리 풋프린트를 약 2.0배에서 2.7배 감소시켜, 소비자용 GPU에서도 실행 가능하게 만듭니다.

5. 의의 및 주장

본 논문은 FAIR-Calib이 확산 기반 언어 모델을 압축하는 데 있어 근본적인 격차를 해결한다고 주장합니다. **"확정이 곧 안정화는 아니다"**라는 점을 인식함으로써, 이 방법은 캘리브레이션의 초점을 균일한 레이어별 재구성이 아닌, 영향력이 큰 특정 "쓰기 프런티어" 상태를 보호하는 것으로 전환합니다.

저자들은 FAIR-Calib을 새로운 모델 능력을 도입하는 방법이 아니라, 자원이 제한된 하드웨어에서 dLLM의 실질적인 배포를 가능하게 하는 강력한 캘리브레이션 절차로 규정합니다. 이 연구는 캘리브레이션 과정에서 전체 확산 롤아웃을 수행하는 계산 비용을 들이지 않고도, 확산 디코딩의 비가역적 역학에 맞춰 양자화 목적 함수를 정렬함으로써 저비트 추론의 취약성을 완화할 수 있음을 보여줍니다. 결과는 메모리와 연산 효율이 중요한 프로덕션 환경으로 dLLM을 확장하는 데 이 접근 방식이 필수적임을 시사합니다.