양자화 인식 힐링(Quantization-Aware Healing): 압축된 4비트 모델이 원본 전체 정밀도 모델을 능가하다

저희의 최신 논문, Quantization-Aware Healing: A Practical Recipe for Recovering Compressed, 4-Bit LLMs는 이 분야에서 대부분 해결되지 않은 질문을 던집니다. 즉, 모델이 구조적 압축뿐만 아니라 이미 지나갔다면, 그 복구 단계가 실제로 얼마나 잘 작동하는지, 그리고 올바른 방법은 무엇인지입니다. 저희는 Quantization-Aware Healing (QAH)을 소개하며, 이를 60B 파라미터로 압축되고 MXFP4로 양자화된 GPT-OSS 120B 모델에 적용했을 때, 9개 벤치마크 중 7개에서 자체의 전체 정밀도(bfloat16) 버전을 능가하는 모델을 생성합니다. 이 4비트 모델은 자신이 양자화되었던 체크포인트보다 더 작고, 실행 비용이 저렴하며, 정확도가 높습니다. 이는 일반적인 4비트 모델과 그것이 온 16비트 모델 간의 관계를 역전시키는 것입니다.

대부분의 효율성 파이프라인은 동일한 세 단계를 따릅니다: 아키텍처 압축, 압축된 가중치 양자화, 그리고 손상 복구(healing)입니다. 방법론 간의 차이는 전적으로 마지막 단계에 있습니다.

지배적인 힐링 레시피는 양자화 인식 학습(Quantization-Aware Training, QAT)입니다. 이는 순전파 과정에 가짜 양자화 연산자를 삽입하고 작업 손실로 모델을 계속 미세 조정하여, 가중치가 낮은 정밀도 표현을 견디도록 학습하게 합니다. 실제로는 이미 비용이 많이 드는 다단계 사후 훈련 프로세스(supervised fine-tuning, RLHF, agentic tuning)를 더 노이즈가 많고 낮은 정밀도의 순전파 과정을 통해 재실행한다는 것을 의미합니다. 이는 비용이 많이 들며, 저희 결과에서 보여주듯이 최적점을 지나 너무 오래 학습을 계속하면 불안정해질 수도 있습니다.

대안으로, 양자화 인식 증류(Quantization-Aware Distillation, QAD)는 그러한 이력을 재실행하는 것을 피합니다. 작업 손실(task loss) 대신, 고정된 전체 정밀도 교사 모델(teacher)의 출력을 KL 발산(KL-divergence) 손실을 통해 양자화된 학생 모델(student)에 직접 증류합니다. 이는 변경 사항이 단지 양자화일 경우 잘 작동하는데, 왜냐하면 동일한 모델의 진정한 전체 정밀도 버전이 교사 역할을 할 수 있기 때문입니다. 하지만 모델이 비트 수만 줄인 것이 아니라 구조적 압축을 거쳐 레이어(layer), 헤드(head) 또는 뉴런(neuron) 수가 적어진 경우, 그 가정은 깨집니다. 더 작은 아키텍처의 독립적으로 훈련된 전체 정밀도 버전은 존재하지 않습니다. 유일한 후보 교사는 복구된 bfloat16 체크포인트이며, 이 역시 원래 모델의 증류 근사치에 불과합니다. 여기서 증류를 수행하면 양자화된 학생이 저하된 목표(degraded target)에 고정되고 그 정확도는 해당 복구 체크포인트 자체의 한계에 갇히게 됩니다.

따라서 구조적으로 압축되었고 양자화된 모델을 어떻게 '치유'할 것인가라는 질문은, 지금까지 진정으로 열려 있던 문제였습니다.

QAH는 단 하나의 변화로 그 한계를 제거합니다. 바로 복구된 모델에서가 아니라 원래의 비압축 모델로부터 직접 증류한다는 것입니다. 교사와 학생은 심지어 아키텍처를 공유하지 않습니다. 교사는 전체 크기 및 전체 정밀도이며, 학생은 절반 크기로 MXFP4로 실행됩니다. 교사의 출력 분포는 아키텍처에 구애받지 않기 때문에, 크기나 모양의 불일치와 아무것도 전송을 막지 못합니다. 학생은 절대 하드 라벨(hard labels)을 보지 않고, 오직 교사의 출력 분포만을 받으며 로짓(logits)에 대한 KL 발산을 통해 일치시킵니다.

이것은 양자화 단계가 무엇을 하는지에 대한 관점을 재정립합니다. QAH의 경우, 이는 더 이상 치유가 완료된 후에 적용되는 손실성 후처리 단계가 아닙니다. 이것은 원래 교사 모델에 대한 두 번째 전체 증류 통과이며, bfloat16 체크포인트가 결코 받지 못한 감독(supervision)입니다. 4비트 학생 모델이 양자화로 인해 손실된 정보를 보상하는 것이 아니라, 이전 복구 단계가 시간이나 데이터 측면에서 전송하지 못했던 정보를 포착하고 있는 것입니다.

손실 자체에서 파생되는 안정성 이점도 있습니다. KL 증류(KL distillation)가 학생 모델을 고정된 교사 분포에 연결하기 때문에, 학생 모델이 따라잡은 후에는 더 이상 표류할 압력이 없습니다. 반면에 크로스 엔트로피 태스크 손실(cross-entropy task loss)은 학생 모델을 무기한으로 어려운 레이블 쪽으로 계속 밀어붙입니다. 이 차이는 아래 비교에서 볼 수 있듯이 정확도와 훈련 안정성 모두에 영향을 미칩니다.

QAH를 긴 컨텍스트, 즉 복구 코퍼스(healing corpus)가 최대 32k 토큰의 문서를 포함하는 환경에서 작동시키기 위해, 우리는 효율적인 증류(efficient distillation)에 관한 이전 논문에서 사용한 메모리 효율적인 청크형 KL-다이버전스 손실을 재사용했습니다. 이 손실은 시퀀스의 한 슬라이스씩 KL 값을 계산하며 전체 어휘-시퀀스 그리드(vocabulary-by-sequence grid)를 절대 구현하지 않기 때문에, 32k 토큰 복구 과정이 고정된 GPU 메모리 예산 내에 들어오게 만듭니다. 이 손실의 작동 원리는 이전 게시물에서 다루었습니다.

QAH 개요. 구조적 압축과 양자화(quantization) 후, 성능은 급격히 떨어집니다. QAH는 복구된 체크포인트가 아닌, 오프라인으로 사전 계산된 로짓을 가진 고정 교사 모델로부터 증류합니다. 출처: 논문 Figure 1.

우리는 GPT-OSS 120B 모델에 QAH를 적용하여, 이 모델을 60B 파라미터로 압축하고 bfloat16으로 복구한 다음, QAH 하에서 MXFP4로 재양자화했습니다. 자연스러운 비교 대상은 같은 60B 모델의 bfloat16 체크포인트, 즉 이 아키텍처가 가진 최고의 전체 정밀도 버전입니다. QAH 모델이 9개 벤치마크 중 7개에서 승리합니다.

원래의 120B 교사 모델(teacher)과의 비교 역시 마찬가지로 의미심장합니다. 교사 모델보다 파라미터 수가 절반이고 가중치 메모리도 약 4분의 1만 사용함에도 불구하고, QAH 모델은 LiveCodeBench에서 전체 크기의 교사 모델을 능가하며(66.5 대 66.0), GPQA Diamond에서는 1.6점 차이로 근접합니다(67.4 대 69.0). 교사 모델과 가장 큰 격차가 남아있는 부분은 AA-LCR인데, 이곳은 압축으로 인해 손실된 용량을 회복하기가 본질적으로 가장 어려운 극단적인 장문 컨텍스트 벤치마크이기 때문입니다.

4비트 QAH 모델은 9개 벤치마크 중 7개에서 bfloat16 소스 모델과 일치하거나 능가하며, LiveCodeBench에서는 전체 크기의 교사 모델을 능가합니다. 출처: 논문 Figure 2.

손실 함수(loss function)의 효과를 다른 모든 것들로부터 분리하기 위해, 우리는 QAH를 동일한 조건에서 QAT와 직접 비교했습니다. GPT-OSS 9B 모델을 MXFP4로 양자화하고, 학습이 진행됨에 따라 MMLU-Pro, LiveCodeBench, GPQA Diamond 전반의 평균 성능을 추적했습니다.

두 방법 모두 유사한 최고점(QAH가 54.9, QAT가 54.6)에 도달하므로, 최고의 정확도 측면에서는 사실상 동등합니다. 차이점은 그 지점에 도달하는 방식과 그 이후에 무슨 일이 일어나는지에 있습니다. QAH는 약 100단계 만에 최고점에 도달하며, 이는 QAT의 700단계보다 약 7배 빠르며, 이후 학습 기간 동안 이 최고점에서 약 두 점 범위 내를 유지합니다. 반면 QAT는 최고점을 지나자마자 급격히 무너져(collapse), 1,200단계까지 진행되면서 거의 19점 가까이를 잃습니다.

실질적인 결과는 실제 배포 위험도의 차이로 이어집니다. QAT 체크포인트는 이미 성능 저하가 시작된 모델을 배포하는 것을 막기 위해 보류 신호(held-out signal)를 이용한 세심한 조기 중단(early stopping)이 필요합니다. 반면, 충분히 훈련된 QAH 체크포인트는 단순히 표류(drift)하지 않기 때문에 안전하게 서비스될 수 있습니다. 이는 메커니즘과 일치하는데, 고정된 교사 모델에 대한 KL 증류(KL distillation)는 학생 모델에게 교사 모델과 일치하는 순간 움직일 동기를 주지 않는 반면, 교차 엔트로피 목적 함수(cross-entropy objective)는 어려운 레이블을 계속 밀어붙여 결국 모델이 원래 가지고 있던 능력을 침식시키기 때문입니다.

QAH는 약 100 스텝에서 54.9로 최고점에 도달하여 유지되는 반면, QAT은 700 스텝 근처에서 54.6에 도달한 후 1,200 스텝까지 거의 19점을 잃습니다. 출처: 논문 Figure 3.

이 정확도 이야기는 애초에 압축을 유발했던 효율성 이야기와 함께 나옵니다. 4비트 정밀도로는 QAH 모델이 bfloat16 학생 모델보다 약 4배 적은 가중치 메모리를 사용하며, 120B 교사(teacher)의 파라미터 수의 절반 수준에서 토큰당 연산량을 대폭 줄여, 훨씬 작은 하드웨어에서도 실행할 수 있게 합니다. 만약 모델군이 4비트가 아닌 bfloat16으로 제공된다면, 결합된 파라미터 및 정밀도 감소는 토큰당 약 8배 적은 연산량에 가까워집니다.

핵심은 압축된 4비트 모델이 반드시 전체 정밀도의 대응 모델보다 정확도가 낮을 필요가 없다는 것입니다. 이 '힐링(healing)' 방법을 사용하면 더 작고, 서비스 비용이 저렴하며, 동시에 더 정확할 수 있으며, 이는 QAT 방법론이 필요로 하는 학습 시간의 훨씬 짧은 기간 안에 달성됩니다. 양자화는 효율성을 위해 지불해야 하는 세금(tax)이 아니라, 모델에게 가르칠 수 있는 추가적인 기회가 됩니다.

본 연구는 Multiverse Computing에서 대규모 모델을 사용에 유용한 능력을 포기하지 않으면서 더 작고 저렴하게 실행하는 방법에 대한 지속적인 연구의 일부입니다. 이는 QAH가 의존하는 장문맥(long-context) 학습 메커니즘을 제공하는 효율적 증류(efficient distillation)에 관한 저희 동료 연구와 함께 자리하고 있습니다.

힐링 파이프라인, 장문맥 힐링을 위한 청크 단위 KL 구현, 분산 학습 결과 등 전체 기술 세부 사항을 알고 싶으신가요? 전체 논문을 읽거나, 압축 및 힐링을 귀하의 모델에 적용하는 것에 대해 저희 팀에 문의해 주십시오.