기술 요약: LLM FP4 사전 학습에서의 Shrinkage Bias 재고

문제 정의

대규모 언어 모델(LLM)이 확장됨에 따라, 사전 학습의 계산 및 메모리 비용을 줄이기 위해 더 낮은 정밀도의 포맷으로 전환하는 것이 필수적입니다. FP8은 널리 채택되었으나, 차세대 하드웨어(예: NVIDIA Blackwell/Rubin, AMD MI350)는 네이티브 FP4 연산을 지원합니다. 현재 NVFP4와 같은 기존 FP4 학습 레시피는 아웃라이어(outlier) 민감도를 완화하기 위해 E2M1 포맷(지수 비트 2개, 가수 비트 1개)과 **Random Hadamard Transform (RHT)**을 결합하여 사용합니다.

그러나 이러한 E2M1 기반 레시피는 학습 불안정성과 BF16 베이스라인 대비 상당한 손실(loss) 저하 문제를 겪습니다. 본 논문은 이 현상의 근본적이고 이전에 간과되었던 한계인 **Shrinkage Bias(수축 편향)**를 식별합니다. 이는 비균등 양자화 포맷(E2M1 등)에서 발생하는 체계적인 음의 반올림 오차입니다. E2M1의 기하학적 비대칭성으로 인해 Round-to-Nearest-Even(RTNE) 반올림 시 일관된 크기 감소(shrinkage)가 발생합니다. 결정적으로, 본 논문은 RHT가 아웃라이어를 분산시키고 코드북 활용도를 높이도록 설계되었지만, 오히려 E2M1의 가장 비대칭적인 반올림 빈(bin) 영역으로 텐서의 질량을 이동시켜 이 편향을 악화시킨다고 주장합니다. 이는 심층 네트워크 전반에 걸쳐 곱셈적인 신호 감쇠를 초래합니다.

방법론

1. Shrinkage Bias의 이론적 분석

저자들은 Shrinkage Bias를 RTNE 빈의 기하학적 비대칭성으로 인해 발생하는 정규화된 크기 공간에서의 음의 기대 오차로 공식화합니다.

• 기하학적 기원: E2M1과 같은 비균등 그리드에서는 표현 가능한 값 사이의 간격이 변합니다(예: 0.5에서 1.0, 1.5로 변화). 특정 레벨 $q_i$에 대해, 좌우 빈 너비($\ell_i$와 $r_i$)는 종종 불일치합니다. 만약 $r_i > \ell_i$라면, 균등한 국소 밀도 하에서의 조건부 기대 오차는 음수($\frac{\ell_i - r_i}{2} < 0$)가 됩니다.
• 체계적 영향: 제로 평균 노이즈(zero-mean noise)와 달리, 이 편향은 행렬 곱셈(GEMM)에서 체계적인 감쇄 인자($\alpha < 1$)로 작용합니다. $K$개의 레이어를 통과하며 이러한 인자들은 곱셈적으로 누적되어($\prod \alpha_k \approx \exp(-\sum \delta_k)$), 체계적인 신호 감쇠를 일으킵니다.
• RHT와의 상호작용: RHT는 텐서를 "동적 범위 제한형(dynamic-range-limited)"에서 "국소 해상도 제한형(local-resolution-limited)"으로 변환합니다. 이는 버킷 활용도를 개선하지만, 데이터가 E2M1의 비대칭성이 가장 두드러지는 중간 크기 영역으로 몰리게 함으로써 Shrinkage Bias를 증폭시킵니다.

2. UFP4 레시피

균등 그리드가 이러한 기하학적 오차를 우회한다는 발견을 바탕으로, 저자들은 E1M2/INT4 스타일의 균등 그리드를 기반으로 한 UFP4(Uniform FP4) 학습 레시피를 제안합니다.

• 균등 그리드: E2M1과 달리, E1M2 그리드는 레벨 간 간격이 동일하여($\ell_i = r_i$) 기하학적 오차를 제거합니다.
• 전방위적 RHT 적용: 균등 그리드는 편향이 없으므로, UFP4는 세 가지 모든 학습 GEMM 경로인 순전파(FPROP), 데이터 그래디언트(DGRAD), 가중치 그래디언트(WGRAD)에 RHT를 안전하게 적용할 수 있습니다. 기존 E2M1 레시피는 비-리프(non-leaf) 경로에서의 오차 누적을 피하기 위해 RHT를 가중치 그래디언트 경로로 제한했습니다.
• 확률적 반올림 (Stochastic Rounding, SR): 그래디언트의 기댓값을 보존하기 위해 SR은 상위 그래디언트($dY$)에만 적용되며, 그 외의 경우에는 RTNE가 사용됩니다.

주요 기여

1. Shrinkage Bias의 식별: 본 논문은 비균등 포맷(E2M1)이 분포 유도 양자화 오차와는 별개로, 기하학적 비대칭성으로 인해 체계적인 음의 반올림 오차를 본질적으로 겪는다는 것을 공식적으로 정의하고 입증했습니다.
2. 불안정성에 대한 통합적 설명: 기존 E2M1 기반 레시피의 학습 불안정성에 대해 이론적, 경험적 설명을 제공하며, Shrinkage Bias가 어떻게 곱셈적으로 누적되고 RHT에 의해 악화되는지 보여줍니다.
3. UFP4 제안: 균등 그리드(E1M2/INT4)를 사용하는 새로운 4비트 학습 레시피를 제안하여, 개선된 버킷 활용도가 더 높은 양자화 품질로 이어지도록 합니다.
4. 산업 규모 검증: Dense 1.5B, MoE 7.9B, MoE 124B 모델을 통해 검증함으로써, 균등 그리드가 실질적인 퍼스트 클래스 FP4 학습 포맷임을 입증했습니다.

실험 결과

양자화 및 GEMM 진단

• RHT 효과: 아웃라이어가 많은 텐서에 대해 RHT는 두 포맷 모두에서 유효 버킷 비율을 크게 개선합니다. 그러나 결과적인 신호 대 양자화 노이즈비(SQNR)는 상이하게 나타납니다. E1M2는 RHT 이후 SQNR이 개선되는 반면, E2M1은 데이터가 비대칭 빈에 위치하게 됨에 따라 SQNR이 저하(음의 $\Delta$SQNR)되는 경-향을 보입니다.
• GEMM 출력: 단일 GEMM 테스트에서, E1M2는 RHT 적용 시 순전파 및 역전파 경로 모두에서 출력 SQNR을 유지하거나 개선하지만, E2M1은 특히 역전파 경로에서 성능 저하를 겪습니다.

엔드 투 엔드(End-to-End) 사전 학습

• 손실 저하 (Loss Degradation): 장기 실행 사전 학습 작업에서, UFP4(E1M2)는 강력한 E2M1 베이스라인에 비해 일관되게 낮은 BF16 대비 손실 저하를 달성했습니다.
  • Dense 1.5B: 상대 오차가 1.2570%(E2M1)에서 0.9673%(UFP4)로 감소했습니다.
  • MoE 7.9B: 2.3596%에서 1.8469%로 감소했습니다.
  • MoE 124B: 1.7308%에서 1.3863%로 감소했습니다.
• 스케일링 법칙 (Scaling Laws): UFP4의 이점은 모델 규모(10M ~ 324M 파라미터) 전반에 걸쳐 지속됩니다. 적합된 손실 곡선을 보면 E1M2와 BF16 사이의 격차가 E2M1보다 작으며, FP4 페널티가 연산 규모에 따라 커지지 않음을 보여줍니다.
• 절제 연구 (Ablation Studies):
  • RHT 범위: 전방위적 RHT 적용(모든 GEMM)이 UFP4에서 최상의 결과를 냈으며, E2M1 레시피처럼 RHT를 제한하면 성능이 제한되었습니다.
  • 범위 제한: E2M1의 동적 범위를 제한(예: 최대값 캡핑)하여 균등 그리드를 모사하려는 시도는, 동적 범위를 희생하면서도 편향을 완전히 제거하지 못해 E1M2의 성능을 따라잡지 못했습니다.

효율성

• 커널 퓨전 (Kernel Fusion): 저자들은 R-H-T를 FP4 양자화와 효율적으로 퓨전할 수 있음을 입증했습니다. 테스트된 하드웨어(SM90/SM100)에서 퓨전된 RHT+양자화 지연 시간은 단독 양자화 대비 약 1.06x1.07x 수준이며, 이는 미퓨전(unfused) 구현의 약 1.4x1.6x와 대조적입니다.

의의 및 주장

본 논문은 E2M1을 유일한 실행 가능한 FP4 학습 포맷로 취급하는 현재의 산업 표준이 기하학적 한계로 인해 최선이 아니라고 주장합니다. 이 연구의 의의는 다음과 같습니다:

1. 현상 유지에 대한 도전: 미래의 ML 가속기가 E2M1에만 의존하기보다, E1M2/INT4 스타일의 균등 4비트 그리드를 핵심 학습 프리미티브로서 지원해야 함을 시사합니다.
2. 포맷-레짐 불일치 (Format-Regime Mismatch): RHT와 같은 아웃라이어 완화 기술의 효과는 밑바탕이 되는 그리드 기하학 구조에 의존한다는 점을 강조합니다. RHT는 균등 그리드에는 유익하지만, 심층 학습 환경의 비균등 그리드에는 해롭습니다.
3. 실용적인 레시피: UFP4는 복잡한 추정기 수정 없이, 단순히 그리드 기하학을 RHT 후의 텐서 분포와 일치시킴으로써 FP4와 BF16 학습 사이의 격차를 줄이는 구체적이고 확장 가능한 레시피를 제공합니다.

저자들은 E2M1이 범위 제한적인 워크로드나 추론에는 여전히 적합할 수 있으나, 학습용 하드웨어는 FP4 사전 학습의 잠재력을 완전히 실현하기 위해 균등 그리드를 지원하는 방향으로 진화해야 한다고 결론짓습니다.