Rethinking Shrinkage Bias in LLM FP4 Pretraining: Geometric Origin, Systemic Impact, and UFP4 Recipe

FP4 훈련은 메모리와 계산비를 크게 절감하지만, E2M1 비균일 격자와 RHT의 조합은 Shrinkage Bias를 유발하고 층을 따라 신호가 체인처럼 감소하는 문제를 초래한다. E1M2/INT4 같은 uniform grid를 사용하면 RHT의 코드북 활용을 퀀타이제이션 품질로 더 잘 전환할 수 있으며, Dense 1.5B, MoE 7.9B, MoE 124B에서 BF16-relative 손실 감소를 일관되게 달성한다. 향후 가속기는 FP4 훈련 프린티브로 uniform grid를 첫-class로 지원해야 한다.

출발점: FP4 훈련은 차세대 대형 모델의 메모리 및 계산 비용을 크게 줄이지만 4비트 포맷의 비효율성은 정확도 저하를 가져올 수 있다. Non-uniform한 E2M1 격자는 RTNE 구간의 기하학적 비대칭으로 Shrinkage Bias를 만들어 층을 지나면서 신호가 체인처럼 약화된다. 중간 기점: RHT는 텐서의 분포를 재배치해 양자화 효율을 높이는 도구이나, E2M1에서의 적용은 바이어스 증가를 야기할 수 있다. Uniform-grid(E1M2/INT4)에서는 RHT가 텐서의 에너지를 균등하게 분산시키고, 양자화 구간 간의 간격이 일정해 바이어스가 제거되며, RHT를 세 GEMM 경로에 걸쳐 적용해도 신호 손실이 더 잘 통제된다. 결론: Uniform-grid 기반의 UFP4는 RHT를 모든 GEMM에서 적용하면서 dY에만 SR을 적용하는 설계로 BF16-relative 손실 감소를 달성하고, 대규모 MoE에서도 확장 가능성을 보인다.

수치적 비교: 4-bit 포맷(E2M1, E1M2, INT4) 및 RHT의 조합에 대해 SQNR, bucket-utilization(B_eff), NMSE 등을 측정한다. 각 계산 경로(GEMM)에서의 양자화 품질 변화를 확인하고, Dense 1.5B, MoE 7.9B, MoE 124B에서 BF16-relative LM 손실을 추적한다.

주요 벤치마크: E1M2 기반 UFP4가 BF16 대비 손실 감소 폭이 큰 것으로 확인되며, 3가지 스케일에서의 손실 지표가 각각 0.9673%, 1.8469%, 1.3863%로 나타난다. Q1에서 RHT의 효과는 텐서의 분포와 격자 형식에 따라 다르며, linear_fc2 같은 경우 E1M2가 ROT 후 더 높은 SQNR를 유도한다. Q3에서 스케일을 확장해도 E1M2의 이점은 지속되며, FP4 대 BF16 차이는 모델 규모에 따라 감소하는 경향을 보인다. Q4에서 전체 RHT 범위와 SR의 조합이 가장 큰 손실 감소를 보이며, E2M1의 모방은 충분치 않다. Q5에서 RHT를 양자화와 합치면 레이턴시 페널티가 작아져 실현 가능성이 높다.

아키텍처: FPROP(fwd_y), DGRAD(bwd_dx), WGRAD(bwd_dw) 3개의 training GEMM에 대해 RHT를 적용하고, 양자화는 G={E1M2/INT4, E2M1}의 코드북으로 수행한다. 스케일링은 FP32 단일 레벨 또는 다단계 스케일 계층을 사용할 수 있다. SR은 dY에만 적용한다. 지표로는 BF16-relative LM loss, SQNR, bucket 활용도(B_eff) 및 NMSE를 사용한다. 봉인된 손실 감소를 통해 E1M2 기반 UFP4가 E2M1 기반 레시피보다 BF16 대비 손실 degradation을 줄임을 보였다.

BF16-Relative 손실은 여전히 존재하고, E1M2 기반의 손실 감소가 모든 상황에서 완전한 BF16과 같지는 않다. E2M1은 여전히 range-limited 워크로드에서 유효할 수 있다. 하드웨어의 완전한 지원 여부에 따라 실제 도입 시 이점이 달라질 수 있다.