HESTIA: 헤시안 기반의 차등 양자화 인식 훈련 프레임워크를 이용한 극저비트 LLM

HESTIA: A Hessian-Guided Differentiable Quantization-Aware Training Framework for Extremely Low-Bit LLMs

대규모 언어 모델(LLM)이 지속적으로 확장됨에 따라, 배포 과정에서 메모리 병목 현상이 심화되고 있으며, 이는 극저비트 양자화로의 전환을 촉진하고 있습니다. 그러나 대부분의 양자화 인식 훈련(QAT) 방법은 훈련 초기에 하드 라운딩과 스트레이트 스루 추정기(STE)를 적용하여 최적화 지형을 조기에 이산화시키고, 잠재 가중치와 양자화된 가중치 간의 지속적인 기울기 불일치를 유발하여 양자화된 모델의 효과적인 최적화를 저해합니다. 이러한 문제를 해결하기 위해, 우리는 극저비트 LLM을 위한 헤시안 기반의 차등 QAT 프레임워크인 Hestia를 제안합니다. Hestia는 경직된 스텝 함수를 온도 제어된 소프트맥스 이완으로 대체하여 훈련 초기 단계에서 기울기 흐름을 유지하면서 점진적으로 양자화를 강화합니다. 또한, Hestia는 텐서 단위의 헤시안 추적 메트릭을 경량화된 곡률 신호로 활용하여 미세 조정된 온도 어닐링을 수행함으로써, 모델 전체에 걸쳐 민감도 기반의 이산화를 가능하게 합니다. Llama-3.2 모델에 대한 평가 결과, Hestia는 기존의 3값 QAT 기준 성능을 꾸준히 능가하며, 1B 모델과 3B 모델에서 각각 평균 5.39%와 4.34%의 향상을 보였습니다. 이러한 결과는 헤시안 기반 이완이 표현 능력을 효과적으로 회복하며, 1.58비트 LLM을 위한 더욱 강력한 훈련 경로를 구축함을 보여줍니다. 코드는 다음 주소에서 확인할 수 있습니다: https://github.com/hestia2026/Hestia.

As large language models (LLMs) continue to scale, deployment is increasingly bottlenecked by the memory wall, motivating a shift toward extremely low-bit quantization. However, most quantization-aware training (QAT) methods apply hard rounding and the straight-through estimator (STE) from the beginning of the training, which prematurely discretizes the optimization landscape and induces persistent gradient mismatch between latent weights and quantized weights, hindering effective optimization of quantized models. To address this, we propose Hestia, a Hessian-guided differentiable QAT framework for extremely low-bit LLMs, which replaces the rigid step function with a temperature-controlled softmax relaxation to maintain gradient flow early in training while progressively hardening quantization. Furthermore, Hestia leverages a tensor-wise Hessian trace metric as a lightweight curvature signal to drive fine-grained temperature annealing, enabling sensitivity-aware discretization across the model. Evaluations on Llama-3.2 show that Hestia consistently outperforms existing ternary QAT baselines, yielding average zero-shot improvements of 5.39% and 4.34% for the 1B and 3B models. These results indicate that Hessian-guided relaxation effectively recovers representational capacity, establishing a more robust training path for 1.58-bit LLMs. The code is available at https://github.com/hestia2026/Hestia.