디퓨전 언어모델 양자화는 과제와 디노이징 단계에 따라 민감도가 다르다
Different Tasks, Different Timesteps: Temporal Task-Aware Quantization for Diffusion Language Models
무엇인가
이 논문은 디퓨전 언어모델(DLM)을 4비트 수준으로 양자화할 때 생기는 활성화 정밀도 문제를 다룬다. DLM은 양방향 문맥을 쓰는 반복적 디노이징으로 토큰을 병렬 예측하기 때문에 자기회귀 LLM보다 모델 평가가 여러 번 반복되고, 기존 자기회귀용 양자화 기법을 그대로 옮기기 어렵다. 기존 DLM 양자화 연구들은 활성화 이상치나 단계별 분포 변화, 오차 누적을 다뤘지만 디노이징 전 구간에 고정된 활성화 정밀도를 쓰는 경우가 많았다. 논문은 정밀도 민감도가 과제와 디노이징 시점에 따라 어떻게 달라지고, 이를 라벨 없이 예측해 예산을 배분할 수 있는지를 묻는다.
어떻게 동작하나
분석은 LLaDA-8B-Instruct에서 GSM8K, HumanEval, MMLU-Pro를 대상으로 했다. 512토큰 생성, 128 디노이징 스텝을 8스텝씩 16개 윈도로 나눠 보면 민감도는 균일하지 않았다. 균일하다면 윈도마다 6.25%씩이어야 하는데, GSM8K는 첫 윈도가 27%, 첫 4분의 1이 55%를 차지했고 HumanEval은 중간 구간인 32~71스텝의 다섯 윈도가 52%를 차지했다. MMLU-Pro는 가장 고르게 퍼졌지만 가장 민감한 윈도와 덜 민감한 윈도가 거의 2배 차이 났다. GSM8K는 초반 집중, MMLU-Pro는 초반이 넓게 퍼진 형태, HumanEval은 중반 피크라는 서로 다른 패턴을 보여 하나의 고정 스케줄로는 맞추기 어렵다.
무엇과 다른가
논문은 추가 정밀도가 어디에 유효한지 라벨 없이 예측하는 신호를 만든다. 각 윈도를 8비트로 올렸을 때의 과제 지표 변화와 별개로, 커밋된 위치에서 전체정밀도와 4비트 예측 분포의 KL 발산 K_w를 재고, 윈도 이후에도 마스크된 위치에 대한 전체정밀도 모델의 믿음 변화 I_w를 잰다. 두 값을 곱한 S_w로 윈도를 순위화하면 GSM8K와 MMLU-Pro에서 가장 이득인 윈도를 맞히며 전체정밀도 격차의 90%와 100%를 회복했고, 윈도별 이득과의 상관은 각각 r=0.98, 0.96이었다. HumanEval은 측정 잡음으로 상관이 r=0.63에 그쳤지만, 발산만 쓰면 마지막 윈도를 잘못 고르는 문제를 영향도 I_w가 보정한다. 라벨 없는 보정 프롬프트 16개만으로도 전체 프롬프트로 계산한 것과 r≥0.96으로 유사했다.
어떻게 쓰나
T2Quant는 이 신호를 과제별 스케줄 최적화로 확장한다. 디노이징 진행률을 K개 윈도로 나누고, 지원 정밀도 집합 B와 기준 정밀도 b0, 윈도별 비용 c_k(b), 기준 예산 C0를 정의한 뒤 각 윈도에 어떤 정밀도를 줄지 bτ를 고른다. 보정 프롬프트에서 전체정밀도 궤적을 저장하고 각 정밀도로 상태를 재생해 국소 발산 κ와 하류 영향 ι를 계산한 다음, gτ,k(b)=ῑτ,k[κτ,k(b0)-κτ,k(b)]로 정밀도 상향과 하향의 추정 이득을 매긴다. 이 이득을 합산하면서 비용 합이 C0가 되도록 동적 계획법으로 최적화한다. 논문은 비용이 음이 아닌 정수이고 B가 유한할 때 이 재귀가 대리 목적함수의 전역 최적해를 준다고 정리하지만, 이는 실제 과제 지표가 아니라 보정된 가산 대리 목적함수에 대한 보장이다.
전제와 한계
실험은 Dream-7B-Instruct, LLaDA-1.5-8B, LLaDA-8B-Instruct와 GSM8K, HumanEval, ARC-Challenge, MMLU-Pro, DROP, BBH에서 이뤄졌다. 비교 대상은 FP16, RTN, DLLMQuant, FAIR-Calib, STaR-Quant이며 각 양자화 기법에 T2Quant를 얹은 경우와 같은 예산의 랜덤 스케줄을 함께 봤다. 모든 구성은 4비트 가중치와 평균 활성화 4.00비트를 맞춘 W4A4다. T2Quant는 9개 베이스라인·모델 조합 모두에서 평균 정확도를 올렸고, 최대 상승폭은 Dream 6.7포인트, LLaDA-1.5 3.4포인트, LLaDA-8B-Instruct 3.5포인트였다. 54개 벤치마크 단위 비교 중 42개가 개선됐고 GSM8K와 ARC-Challenge는 모든 조합에서 개선됐다.
수치로 보면 Dream의 GSM8K는 DLLMQuant에서 35.6에서 53.5로, STaR-Quant에서 46.1에서 58.8로 올랐다. GSM8K 개선폭은 LLaDA-1.5에서 9.6포인트, LLaDA-8B-Instruct에서 10.4포인트에 이른다. T2Quant를 얹은 구성은 각 모델에서 가장 높은 양자화 평균 두 개를 차지했고, 가장 강한 단독 베이스라인보다 Dream 3.55, LLaDA-1.5 1.13, LLaDA-8B-Instruct 1.40포인트 앞섰다. FP16과의 격차는 각각 2.0, 2.5, 1.6포인트까지 줄었다. 같은 평균 비트 예산의 랜덤 스케줄은 Dream에서 원래 베이스라인보다 14~19포인트 낮아, 정밀도를 어디에 몰아주는지가 중요함을 보여준다.
절제 실험은 LLaDA-8B-Instruct와 GSM8K의 W4A4에서 균일 44.1, T2Quant 49.4, FP16 52.9를 기록했다. 보정 프롬프트 16개로 전체 이득의 약 75%를 회복했고, 스텝당 토큰 수를 4에서 2로 줄여도 약 5포인트 개선이 유지됐지만 8토큰에서는 전체 정확도가 나빠지며 이득이 사라졌다. 평균 3.75비트만 써도 48.4%로 균일 W4A4의 44.1%를 넘었고, 4비트를 넘겨 정밀도를 더 올리는 것은 거의 이득이 없었다. 개발자 입장에서는 기존 PTQ 위에 과제별 보정 프롬프트로 스케줄만 학습해 얹을 수 있고, 추론 시에는 윈도 인덱스로 정밀도를 조회하므로 추가 진단용 순전파가 필요 없다는 점이 실용적이다. 다만 보정 프롬프트가 대상 과제를 대표해야 하고, 스케줄은 과제마다 하나씩 나온다.
한계도 분명하다. T2Quant가 최적화하는 것은 실제 정확도가 아니라 국소 발산과 하류 영향을 결합한 대리 신호이며, 동적 계획법의 최적성도 가산 비용을 가정한 대리 목적함수에 한정된다. 논문은 하드웨어 효율, 즉 추론 속도와 에너지 소비, 그리고 양자화가 안전성이나 견고성에 미치는 영향을 평가하지 않았고 이를 향후 과제로 남긴다. 따라서 실제 배포에서는 대상 과제의 비라벨 보정 프롬프트 확보, 스텝당 토큰 수와 윈도 크기 같은 샘플러 설정, 기존 양자화 기법과의 결합 호환성, 안전성 회귀를 함께 확인해야 한다.