기술 요약: STaR-Quant

문제 정의

확산 대규모 언어 모델(Diffusion Large Language Models, DLLMs)은 양방향 컨텍스트를 활용한 반복적인 마스크 디노이징(iterative masked denoising)을 통해 텍스트를 생성하며, 자기회귀(autoregressive) LLM의 유망한 대안으로 부상했습니다. 그러나 이러한 모델들은 거대한 모델 크기와 다단계 반복 디노이징 과정에서 발생하는 상당한 메모리 및 연산 오버헤드로 인해 배포에 어려움을 겪고 있습니다. 사후 학습 양자화(Post-Training Quantization, PTQ)가 자기회귀 LLM의 압축에는 성공적으로 적용되어 왔으나, 기존의 PTQ 방식을 DLLM에 직접 적용할 경우, 특히 공격적인 저비트 가중치-활성화 양자화 상황에서 심각한 성능 저하가 발생합니다.

본 논문은 이러한 실패를 야기하는 DLLM 특유의 두 가지 구체적인 과제를 식별합니다:

1. 상태 의존적 활성화 불일치(State-Dependent Activation Disparity): 단일 디노이징 단계 내에서 시퀀스는 마스크된 토큰(예측될 미확정 위치)과 마스크되지 않은 토큰(가시적 컨텍스트)을 모두 포함합니다. 이 두 상태는 서로 다른 활성화 분포와 아웃라이어 패턴을 보입니다. 두 상태 모두에 대해 단일한 활성화 변환이나 양자화 공간을 적용하는 것은 두 상태의 분포를 효과적으로 평활화(smoothing)하는 데 실패합니다.
2. 시간적 오차 누적(Temporal Error Accumulation): DLLM 생성은 한 단계의 출력이 다음 단계의 입력 역할을 하는 반복적인 단계를 통해 진행됩니다. 양자화 오차, 특히 소프트맥스 어텐션 맵(softmax attention map)과 밸류(value) 상태 사이의 행렬 곱셈에서 발생하는 오차는 단일 단계에 국한되지 않고 디노이징 궤적을 따라 누적되며, 이는 전체 정밀도(full-precision) 경로로부터의 상당한 편차를 초래합니다.

방법론: STaR-Quant

이러한 과제를 해결하기 위해 저자들은 DLLM을 위해 특별히 설계된 상태-시간 일관성 PTQ 프레임워크인 STaR-Quant를 제안합니다. 이 프레임워크는 두 가지 차원에서 작동합니다:

1. 상태 가이드 활성화 변환 (State-Guided Activation Transformation, SGAT)

SGAT는 마스크된 토큰과 마스크되지 않은 토큰 사이의 활성화 불일치를 해결합니다.

• 메커니즘: 모든 토큰에 대해 단일 변환을 사용하는 대신, SGAT는 은닉 차원(hidden dimension)을 공유 부분 공간과 상태별 부분 공간으로 분해합니다.
  • 공유 부분 공간 ($U_{sh}$): 공통적인 활성화 성분을 포착합니다.
  • 상태별 부분 공간 ($U_m, U_u$): 각각 마스크된 토큰과 마스크되지 않은 토큰에 특화된 아웃라이어 패턴을 흡수합니다.
• 구현: 이진 부분 공간 게이트(binary subspace gates)가 마스크된 토큰과 마스크되지 않은 토큰을 각각의 변환 공간($P_m$ 및 $P_u$)으로 라우팅합니다. 결정적으로, 가중치 측은 가중치에 한 번 적용되는 통합 정적 변환($P_e$)을 유지합니다. 이는 추론 시 변환된 가중치를 중복 생성하지 않으면서도 활성화 측이 토큰 상태에 적응할 수 있도록 합니다.
• 최적화: 변환 파라미터는 전체 정밀도 출력과 상태 가이드 양자화 투영(state-guided quantized projection) 사이의 레이어별 양자화 오차를 최소화함으로써 학습됩니다.

2. 시간적 어텐션 보상 (Temporal Attention Compensation, TAC)

TAC는 반복적인 디노이징 단계 전반에 걸친 양자화 오차의 누적을 해결합니다.

• 대상: 이 방법은 저비트 양자화 하에서 오차가 발생하기 쉬운 연산인 어텐션 메커니즘의 소프트맥스-밸류 곱셈 출력을 타겟으로 합니다.
• 메커니즘: TAC는 어텐션 출력 투영(attention output projection) 이전에 가벼운 블록 대각 선형 매핑(block-diagonal affine mapping)을 삽eli(삽입)합니다. 이는 어텐션 계산 자체를 변경하지는 않지만, 양자화된 어텐션 표현($\hat{Z}_t$)을 전체 정밀도 대응물($Z^{FP}_t$)과 일치하도록 교정합니다.
• 추정: 보상 파라미터(블록 단위 행렬 $M_b$ 및 편향 $d_b$)는 전체 정밀도와 양자화된 표현의 1차 및 2차 통계량(평균 및 공분산)을 일치시킴으로써 폐쇄형(closed form)으로 추정됩니다. 제한된 캘리브레이션 데이터를 고려하여 안정성을 보장하기 위해 댐핑 항(damping term)과 특이값 평활화(singular value smoothing)가 적용됩니다.

주요 기여

1. 문제 식별: 저자들은 상태 의존적 활성화 불일치와 시간적 오차 누적을 저비트 DLLM 양자화의 주요 병목 현상으로 공식적으로 식별했습니다.
2. 상태 가이드 활성화 변환 (SGAT): 가중치 중복 없이 상태별 아웃라이어를 효과적으로 평활화하면서, 마스크된 활성화와 마스크되지 않은 활성화를 공유 및 상태별 변환 공간으로 매핑하는 새로운 방법을 제시합니다.
3. 시간적 어텐션 보상 (TAC): 양자화된 어텐션 표현을 출력 투영 전에 전체 정밀도 대응물과 정렬시키는 폐쇄형 블록 단위 아핀 보상법을 통해 디노이징 궤적을 따른 오차 누적을 완화합니다.
4. 실증적 검증: 광범한 실험을 통해 STaR-Quant가 여러 DLLM 및 태스크에 걸쳐 강력한 PTQ 베이스라인(RTN, AWQ, QuaRot, DLLMQuant 포함)을 일관되게 능가함을 입증했습니다.

실험 결과

저자들은 세 가지 대표적인 DLLM인 LLaDA-8B, LLaDA-1.5-8B, Dream-7B에 대해 4비트 가중치 및 활성화(W4A4) 양자화를 사용하여 STaR-Quant를 평가했습니다.

• 성능: STaR-Quant는 9개 벤치마크(일반 지식, 추론, 코드 생성 포함)에서 가장 우수한 평균 성능을 달성했습니다.
  • LLaDA-8B에서, 가장 강력한 베이스라인(DLLMQuant++)보다 평균 점수를 2.78 포인트 향상시켰습니다.
  • Dream-7B에서, 평균 점수를 1.69 포인트 향상시켰습니다.
  • 중간 표현 오차에 민가한 추론(GSM8K) 및 코드 생성(HumanEval) 태스크에서 주목할 만한 이득이 관찰되었습니다.
• 효율성:
  • 속도 향상: STaR-Quant는 FP16 배포 대비 1.65배에서 1.69배의 추론 속도 향상을 달了했습니다.
  • 메모리: 피크 GPU 메모리 사용량을 3.05배에서 3.14배 감소시켰습니다 (예: LLaDA의 메모리를 약 15.9GB에서 5.2GB로 감소).
• 절제 연구 (Ablation Studies): TAC를 제거하면 성능이 크게 떨어졌으며(시간적 오차 교정의 중요성 입증), SGAT를 제거했을 때도 결과가 저하되어 두 요소의 상호 보완적인 성격을 확인했습니다.

의의 및 주장

본 논문은 STaR-Quant가 확산 기반 언어 모델을 위한 양자화 전략의 필수적인 전환을 제공한다고 주장합니다. 상태(마스크됨 vs 마스크되지 않음)와 시간(반복적 누적) 차원을 명시적으로 모델링함으로써, 이 프레임워크는 기존 PTQ의 가정과 확산 추론 프로세스 사이의 불일치를 해결합니다.

저자들은 STaR-Quant를 자원이 제한된 장치에서 DLLM을 배포하기 위한 실질적인 솔루션으로 포지셔닝하며, 높은 압축률(W4A4), 상당한 메모리 절감 및 속도 향상을 제공하는 동시에 전체 정밀도 모델에 근접한 정확도를 유지합니다. 이 연구는 "상태-시간 일관성"이 효율적인 DLLM 양자화를 위한 효과적인 원칙임을 시사하지만, 저자들은 다양한 아키텍처와 더 공격적인 비트 폭(예: 2비트 또는 3비트)에 대한 추가 검증이 열린 과제로 남아 있다고 언급했습니다.