극단 양자화와 저랭크 보정의 융합: HGF로 1.58비트 LLM 안정화

초록

본 논문은 1.58비트(ternary) 양자화된 대형 언어 모델에 저랭크 FP16 보정 경로와 적응형 게이트를 결합한 Hybrid Gated Flow(HGF) 구조를 제안한다. TinyStories 데이터셋에서 2500·3500 스텝 학습 시 검증 손실 0.9306을 달성해 BitNet(b1.58)의 1.0294 대비 약 55% 품질 격차를 회복했으며, 메모리 오버헤드는 12~15% 수준에 머문다. 또한 양자화를 구조적 정규화로 해석하고, 대규모 1.2B·3B 모델에서도 안정적인 학습 특성을 확인한다.

상세 분석

본 연구는 “Memory Wall”이라는 하드웨어 병목을 극복하기 위해 1.58비트 극단 양자화와 저랭크 보정의 시너지를 탐구한다. 먼저, 논문은 ternary 가중치 {−1,0,1} 을 absmax 스케일링과 STE(직통 추정)를 통해 FP16과 동일한 학습 파이프라인에 삽입한다. 이때 양자화 함수는 거의 모든 위치에서 미분이 0이므로, STE가 필수적인데, 저자는 스케일 γ_W 가 작을수록 STE 편향이 제한적임을 정량화한다.

핵심 기여는 두 번째 스트림인 저랭크 FP16 보정이다. 저자는 양자화 오차 ε_q = X(W−f_W)^T 가 저랭크 서브스페이스에 집중된다고 가정하고, LoRA와 유사한 A∈ℝ^{d_in×r}, B∈ℝ^{r×d_out} ( r≪d )를 도입한다. 여기서 비선형 활성화 SiLU를 삽입해 단순 선형 보정보다 복잡한 오차 형태를 학습하도록 설계했다.

게이트 메커니즘은 스칼라 α 에 tanh를 적용해 g= tanh(α) 로 정의하고, 최종 출력 Y_HGF = Y_tern + g·Y_corr 으로 결합한다. 이때 ∂L/∂α = (∂L/∂Y_HGF)·Y_corr·sech²(α) 임을 증명해, |α|→∞ 시 기울기가 급격히 사라지는 “게이트 포화” 현상을 정규화 효과로 활용한다. 또한 게이트 초기값을 α₀=0.1 로 설정해 초기 g≈0.1 을 확보하고, “live initialization”을 통해 B를 작은 가우시안(σ=10⁻³)으로 초기화함으로써 보정 경로가 초기에 죽는 현상을 방지한다.

학습 단계에서는 메인 파라미터와 게이트 파라미터에 서로 다른 학습률(η_main=2.5e‑3, η_gate=2.5e‑4)을 적용하고, 500~900 스텝 구간에 평균 게이트 크기에 비례한 정규화 손실을 부과한 뒤 900 스텝 이후에는 게이트를 고정한다. 이는 게이트가 과도하게 성장하거나 급격히 사라지는 것을 방지하고, 최종 모델이 안정적인 보정 비율을 유지하도록 만든다.

이론적으로는 ternary 가중치가 ‖Q‖·‖K‖ 을 제한해 attention logits의 분산을 감소시키고, 게이트가 작은 g 값을 유지함에 따라 gradient variance가 FP16 대비 (1+O(g²)) 정도 감소한다는 정리를 제시한다. 실험적으로는 TinyStories에서 2500·3500 스텝 학습 시 BitNet(b1.58) 대비 0.0988 낮은 검증 손실을 기록했으며, 메모리 사용량은 ternary 백본에 비해 12~15% 정도만 추가된다. 또한 “Diff_Only”라는 완전 FP16 differential attention 베이스라인은 손실이 1.68을 초과하며 불안정했지만, HGF는 전 과정에서 안정적인 수렴을 보였다.

대규모 실험에서는 1.2B와 3B 파라미터 모델을 SlimPajama·FineWeb‑Edu 데이터에 적용했으며, 품질 회복 비율이 소규모 실험과 비슷하게 선형적으로 확장됨을 관찰했다. 다만 아직 최종 체크포인트와 완전한 로그가 공개되지 않아 정량적 비교는 제한적이다.

전체적으로 HGF는 극단 양자화의 메모리 효율성을 유지하면서, 저랭크 고정밀 보정과 적응형 게이트를 통해 학습 안정성과 품질을 크게 향상시킨 설계이다. 특히 양자화를 “구조적 정규화”로 해석하고, 게이트 포화 메커니즘을 통해 자동적인 정규화를 구현한 점이 학술적·공학적 의미가 크다.