출력 헤드의 softmax 등가성으로 저비트 양자화 손실을 줄이는 재매개변수화
Softmax Reparameterization for Output-Head Quantization
무엇인가
대규모 어휘를 쓰는 소형 언어모델에서 출력 헤드는 무시하기 어려운 추론 비용이다. 어휘 크기 V와 은닉 폭 d에 대해 출력 투영은 V·d개의 가중치를 가지며, 최근 3~4B 모델에서는 대략 4억~7억 개, 명목 모델 크기의 10~17%를 차지한다. 입력 임베딩이 선택된 행만 읽는 것과 달리 출력 투영은 디코딩 스텝마다 전체 어휘를 계산하고, BF16 기준 스텝당 약 2Vd 바이트의 가중치를 읽는다. 단일 토큰 디코딩은 메모리 바운드이므로 이 가중치 이동 자체가 비용이고, 어휘가 커질수록 헤드 압축의 필요성도 커진다. 그런데도 실무 양자화 레시피는 출력 헤드를 더 높은 정밀도로 남겨두는 경우가 많다. 헤드는 다음 토큰 확률을 직접 결정하며 오차를 흡수할 후속 학습 레이어가 없기 때문이다.
어떻게 동작하나
이 논문이 제안하는 것은 softmax 재매개변수화다. softmax는 로짓의 절대값이 아니라 상대값에만 의존하므로, 모든 어휘 행에서 같은 벡터를 빼도 출력 분포는 정확히 같다. 저자들은 어휘 행 평균 μ = (1/V)W^T·1 방향으로 W_t = W − t·1·μ^T 라는 1차원 스칼라 패밀리를 정의하고, 양자화 후의 검증 KL을 최소화하는 t를 고른다. t=0은 원래 헤드, t=1은 통상적인 평균 중심화(mean-centering)에 해당하므로, 이 탐색은 두 방법을 후보로 포함하면서 그보다 나은 대표를 찾는다. 계수는 RTN, activation-weighted MSE, full-Hessian GPTQ 각각에 대해 따로 선택하며, [-2, 8] 구간의 14점 격자를 쓴다. 학습된 디코더는 건드리지 않고 재학습도 필요 없다. tanh soft-capping처럼 비선형 로짓 경로가 있는 헤드(Gemma 4)에는 rank-one 보정을 적용한다. 선택은 양자화기 적합에 쓰는 상태 집합과 분리된 검증 집합에서 이루어지고, 0이 격자에 포함되므로 검증 KL 기준으로는 항상 원래 헤드보다 나쁘지 않다는 보장이 있다.
무엇과 다른가
실험은 출력 헤드만 양자화하고 디코더는 고정한 채, 일곱 개 헤드(Gemma 3/4, Qwen3.5, Phi, BLOOM-1.7B, BLOOMZ-1.7B, XGLM-1.7B)를 대상으로 한다. W4는 G128 그룹, BF16 디코더, RTN/AW-MSE/full-Hessian GPTQ 조합이고, 적합에는 WikiText 128개 문서에서 8개 상태씩, 검증과 테스트에는 각각 16개 문서를 쓴다. 기준 양자화가 예측을 크게 왜곡하는 헤드에서 이득이 집중된다. XGLM의 RTN 테스트 KL은 2.13에서 0.143으로 93% 줄고, AW-MSE에서 Phi는 0.936에서 0.256(73%), BLOOM은 0.594에서 0.136(77%), BLOOMZ는 0.662에서 0.158(76%)로 떨어진다. 반면 Gemma 3/4와 Qwen3.5는 W4 기준선 KL이 이미 낮아 변화가 작다. 압축 스트레스 테스트로 쓴 W2에서는 거의 모든 모델-양자화기 조합에서 이득이 나타나지만, 여러 설정의 절대 품질은 여전히 나쁘다고 저자들은 밝힌다.
어떻게 쓰나
고정 평균 중심화와의 비교도 분명하다. Phi에서 t=1은 RTN 테스트 KL을 약 1.23에서 0.81로 낮추지만, 검증이 고른 t=4는 약 0.35까지 낮춰 고정 중심화 대비 약 57% 추가 개선을 준다. t=1은 full-precision 가중치 노름을 최소화하지만 양자화 후 KL을 최소화하지는 않는다는 것이 요점이다. 더 강한 GPTQ 캘리브레이션(1,024 → 65,536 상태)에서도 원래 헤드의 W4 테스트 KL이 0.160에서 0.089로 내려간 뒤 재매개변수화가 0.035까지 추가로 낮춘다(61%). 정확한 per-channel 스케일링과 정수 영점 affine 양자화 위에서도 효과가 겹친다(스케일드 AW-MSE 0.240 → 0.172, 스케일드 RTN 0.337 → 0.242, affine RTN 0.750 → 0.259). 이전에 쓰지 않은 WikiText 26개 문서 홀드아웃에서 Phi는 AW-MSE KL 73.6% 감소(0.943 → 0.249), RTN 71.1% 감소, BLOOM-1.7B는 AW-MSE 74.5% 감소(0.558 → 0.142), RTN 51.7% 감소를 재현했다. WikiText에서 고른 계수를 재조정 없이 C4와 OpenWebMath에 적용했을 때, 계수가 1이 아닌 18개 비교 모두에서 고정 평균 중심화를 앞섰고 나머지 6개 XGLM 사례에서는 동률이었다.
전제와 한계
배포 관점의 결과도 있다. Marlin INT4 커널과 vLLM으로 패킹된 W4 실행을 할 때, Phi 헤드에서 t=4 이동은 16,352개 WikiText 토큰 기준 AW-MSE 퍼플렉서티를 31.65에서 15.40으로, BF16 대비 KL을 0.97에서 0.28로 낮춘다. 가장 강한 배포 기준선인 full-Hessian GPTQ에서도 퍼플렉서티가 13.67에서 12.23으로, KL이 0.15에서 0.06으로 내려가 BF16의 11.65 대비 약 5% 초과 수준이 된다. 선형 softmax 헤드에서는 이동이 양자화 전에 가중치에 접혀 들어가므로 패킹 레이아웃이 유지되고 추론 연산이 추가되지 않는다. 실제로 이동 후 지연은 거의 그대로였고, 이동된 AW-MSE 헤드는 BF16 헤드 기준선보다 배치1 지연이 10.8%, 배치16이 9.4% 낮았다. 다만 tied 모델에서는 BF16 입력 임베딩을 유지한 채 별도 패킹 출력 헤드를 추가하므로 상주 가중치가 7.17GiB에서 7.47GiB로 늘어난다.
왜 이런 일이 가능한지에 대한 분석은 재구성 오차와 예측 충실도가 다른 축이라는 점을 보여준다. 저자들은 softmax Fisher 행렬 F(p) = diag(p) − pp^T를 진단 도구로 쓴다. Phi-4-mini에서 t=0에서 t=4로 갈 때 raw 로짓 오차 에너지는 RTN 3.0배, AW-MSE 2.7배로 늘어나지만 Fisher 가중 오차는 72~73% 줄어든다. GPTQ 65,536 캘리브레이션에서도 로짓 오차 에너지 D_z는 0.025%에서 0.067%로 2.66배 커지는데 Fisher/2 비용은 0.0888에서 0.0347로, 테스트 KL은 0.0894에서 0.0349로 떨어진다. 확률 구간별로 보면 p < 10⁻⁶인 항목(확률 질량의 0.20%)이 오차 에너지에서 차지하는 비중이 92.4%에서 98.2%로 커지고, 질량의 84.8%를 가진 p ≥ 0.01 항목의 오차는 절반 이하로 줄어든다. 즉 선택된 대표는 전체 재구성 오차를 키우면서 그 오차를 softmax가 둔감한 방향으로 밀어 넣는다.
개발자 입장에서 이 방법은 기존 파이프라인에 얹기 쉬운 편이다. 이미 쓰는 양자화기(RTN, AW-MSE, GPTQ)를 그대로 두고, 양자화 직전에 스칼라 하나만 탐색하면 된다. 재학습이 없고 디코더를 바꾸지 않으며, 선형 softmax 헤드에서는 추론 시 추가 연산이 없어 W4 패킹 실행과 지연 이득을 그대로 유지한다. 다만 계수는 양자화기별로, 그리고 데이터 분포별로 따로 골라야 한다. RTN에서 고른 계수를 AW-MSE로 옮기면 충실도가 떨어질 수 있고, 더 큰 C4/OpenWebMath 검증 집합은 WikiText와 다른 계수를 고를 수 있다. 적용 전에 확인할 것은 세 가지다. 헤드가 선형 softmax인지(비선형 soft-cap이면 rank-one 보정 필요), tied 임베딩 모델이라면 별도 출력 사본에 따른 메모리 증가를 감당할 수 있는지, 그리고 기준 양자화 KL이 이미 낮은 헤드라면 얻을 게 거의 없다는 점이다.
저자들이 명시한 한계도 분명하다. 탐색은 어휘 행 평균 방향의 1차원 경로에 한정되며, 그 경로가 최적 대표를 포함한다고 주장하지 않는다. t=0이 격자에 있으므로 검증 KL이 나빠지지 않는다는 보장은 선택에 쓴 양자화기와 검증 데이터에만 적용되고, 보지 못한 데이터나 다른 양자화기로 확장되지 않는다. 선호되는 계수는 분포 의존적이다. W2 결과는 이득의 범위가 넓다는 것을 보여주는 압축 스트레스 테스트이며, 여러 설정에서 절대 품질은 여전히 낮다. 또한 초기 평가에 쓴 문서들이 이전 탐색 과정에서 이미 접한 것이어서, 저자들은 별도의 미사용 홀드아웃으로 재검증했다. tied 모델의 저장 비용 증가와 비선형 로짓 경로에 대한 rank-one 보정 필요성도 전제 조건으로 남는다.