Resumo Técnico: Repensando o Viés de Encolhimento (Shrinkage Bias) no Pré-treinamento de LLMs em FP4

Declaração do Problema

À medida que os Modelos de Linguagem de Grande Escala (LLMs) escalam, os custos computacionais e de memória do pré-treinamento exigem uma transição para formatos de menor precisão. Embora o FP8 tenha sido amplamente adotado, o hardware emergente (ex: NVIDIA Blackwell/Rubin, AMD MI350) suporta computação nativa em FP4. As atuais receitas de treinamento para FP4, como a NVFP4, baseiam-se predominantemente no formato E2M1 (2 bits de expoente, 1 bit de mantissa) combinado com a Transformada de Hadamard Aleatória (RHT) para mitigar a sensibilidade a outliers.

No entanto, essas receitas baseadas em E2M1 sofrem de instabilidade de treinamento e uma degradação significativa de perda em relação às bases em BF16. O artigo identifica uma limitação fundamental, anteriormente negligenciada: o Viés de Encolhimento (Shrinkage Bias). Este é um erro de arredondamento negativo sistemático inerente a formatos de quantização não uniformes como o E2M1. A assimetria geométrica dos bins representáveis do E2M1 causa um "encolhimento" sistemático de magnitudes sob o arredondamento Round-to-Nearest-Even (RTNE). Crucialmente, o artigo argumenta que, embora a RHT seja projetada para dispersar outliers e melhorar a utilização do codebook, ela inadvertidamente exacerba este viés no E2M1 ao deslocar a massa dos tensores para as regiões de arredondamento mais assimétricas do formato, levando a um decaimento multiplicativo do sinal através de redes profundas.

Metodologia

1. Análise Teórica do Viés de Encolhimento

Os autores formalizam o Viés de Encolhimento como um erro esperado negativo no espaço de magnitude normalizado causado pela assimetria geométrica dos bins RTNE.

• Origem Geométrica: Em grades não uniformes como o E2M1, o espaçamento entre valores representáveis muda (ex: de 0,5 para 1,0 para 1,5). Para um nível $q_i$, as larguras dos bins à esquerda e à direita ($\ell_i$ e $r_i$) são frequentemente desiguais. Se $r_i > \ell_i$, o erro esperado condicional sob densidade local uniforme é negativo ($\frac{\ell_i - r_i}{2} < 0$).
• Impacto Sistêmico: Diferente do ruído de média zero que se cancela, este viés atua como um fator de atenuação sistemática ($\alpha < 1$) em multiplicações de matrizes (GEMMs). Através de $K$ camadas, esses fatores se acumulam multiplicativamente ($\prod \alpha_k \approx \exp(-\sum \delta_k)$), causando um decaimento sistemático do sinal.
• Interação com RHT: A RHT transforma os tensores de serem "limitados pelo alcance dinâmico" para serem "limitados pela resolução local". Embora isso melhore a utilização dos buckets, força os dados para as regiões de magnitude média onde a assimetria do E2M1 é mais pronunciada, amplificando assim o viés de encolhimento.

2. A Receita UFP4

Com base na descoberta de que grades uniformes contornam este erro geomético, os autores propõem o UFP4 (Uniform FP4), uma receita de treinamento construída sobre uma grade uniforme do tipo E1M2/INT4.

• Grade Uniforme: Ao contrário do E2M1, a grade E1M2 mantém o espaçamento igual entre os níveis ($\ell_i = r_i$), eliminando a fonte geométrica do Viés de Encolhimento.
• Cobertura Total de RHT: Como a grade uniforme é imparcial, o UFP4 aplica a RHT com segurança a todos os três caminhos de GEMM do treinamento: Forward (FPROP), Gradiente de Dados (DGRAD) e Gradiente de Peso (WGRAD). As receitas anteriores de E2M1 restringiam a RHT ao caminho do gradiente de peso para evitar a composição de erros em caminhos que não fossem folhas.
• Arredondamento Estocástico (SR): O SR é aplicado apenas ao gradiente ascendente ($dY$) para preservar as expectativas do gradiente, enquanto o RTNE é usado em outros lugares.

Principais Contribuições

1. Identificação do Viés de Encolhimento: O artigo define e demonstra formalmente que formatos não uniformes (E2M1) sofrem inerentemente de um erro de arredondamento negativo sistemático devido à assimetria geométrica, distinto de erros de quantização induzidos pela distribuição.
2. Explicação Unificada da Instabilidade: Fornece uma explicação teórica e empírica para a instabilidade de treinamento nas atuais receitas baseadas em E2M1, mostrando como o Viés de Encolhimento se acumula multiplicativamente e é exacerbado pela RHT.
3. Proposta do UFP4: Uma nova receita de treinamento de 4 bits utilizando grades uniformes (E1M2/INT4) que permite a aplicação pervasiva de RHT em todos os GEMMs de treinamento, convertendo a melhoria na utilização de buckets em maior qualidade de quantização.
4. Validação em Escala Industrial: A abordagem é validada em modelos Dense 1.5B, MoE 7.9B e MoE 124B, demonstrando que grades uniformes são formatos práticos de primeira classe para treinamento em FP4.

Resultados Experimentais

Diagnósticos de Quantização e GEMM

• Efeito RHT: Em tensores com muitos outliers, a RHT melhora significativamente a razão efetiva de buckets para ambos os formatos. No entanto, a Razão Sinal-Ruído de Quantização (SQNR) resultante diverge: o E1M2 mostra uma SQNR melhorada pós-RHT, enquanto o E2M1 frequentemente degrada ( $\Delta$SQNR negativa) porque os dados deslocados caem em bins assimétricos.
• Saídas de GEMM: Em testes de GEMM único, o E1M2 preserva ou melhora a SQNR de saída através dos caminhos forward e backward com RHT, enquanto o E2M1 sofre degradação, particularmente nos caminhos de backward.

Pré-treinamento de Ponta a Ponta

• Degradação de Perda: Em tarefas de pré-treinamento de longa duração, o UFP4 (E1M2) alcança consistentemente uma degradação de perda relativa ao BF16 menor comparado a bases fortes de E2M1.
  • Dense 1.5B: Erro relativo reduzido de 1,2570% (E2M1) para 0,9673% (UFP4).
  • MoE 7.9B: Reduzido de 2,3596% para 1,8469%.
  • MoE 124B: Reduzido de 1,7308% para 1,3863%.
• Leis de Escalonamento (Scaling Laws): A vantagem do UFP4 persiste através de escalas de modelo (10M a 324M parâmetros). As curvas de perda ajustadas mostram que a lacuna do E1M2 em relação ao BF16 permanece menor que a lacuna do E2M1, e a penalidade de FP4 não cresce com a escala de computação.
• Estudos de Ablação:
  • Escopo RHT: A cobertura total de RHT (todos os três GEMMs) produz os melhores resultados para o UFP4, enquanto restringir a RHT (como feito em receitas de E2M1) limita o desempenho.
  • Restrição de Alcance: Tentativas de emular uma grade uniforme restringindo o alcance dinâmico do E2M1 (ex: limitando o valor máximo) falharam em igualar o desempenho do E1M2 nativo, sacrificando o alcance dinâmico sem eliminar totalmente o viés.

Eficiência

• Fusão de Kernel: Os autores demonstram que a RHT pode ser fundida com a quantização FP4 de forma eficiente. Em hardware testado (SM90/SM100), a latência da fusão RHT+quantização é de apenas ~1,06x–1,07x em relação à quantização isolada, comparado a ~1,4x–1,6x para implementações não fundidas.

Significância e Alegações

O artigo argumenta que o padrão atual da indústria de tratar o E2M1 como o único formato viável de treinamento FP4 é subótimo devido às suas limitações geométricas. A significância deste trabalho reside em:

1. Desafiar o Status Quo: Sugere que futuros aceleradores de ML devem suportar grades uniformes de 4 bits do tipo E1M2/INT4 como primitivas de treinamento de primeira classe ao lado do E2M1, em vez de depender apenas do E2M1.
2. Descompasso de Regime de Formato: Destaca que a eficácia das técnicas de mitigação de outliers, como a RHT, depende da geometria da grade subjacente. A RHT é benéfica para grades uniformes, mas prejudicial para grades não uniformes em regimes de treinamento profundos.
3. Receita Prática: O UFP4 oferece uma receita concreta e escalável que reduz a lacuna entre o treinamento em FP4 e BF16 sem exigir modificações complexas de estimadores, simplesmente alinhando a geometria da grade com a distribuição do tensor pós-RHT.

Os autores concluem que, embora o E2M1 continue adequado para cargas de trabalho limitadas por alcance e inferência, o hardware de treinamento deve evoluir para suportar grades uniformes para realizar plenamente o potencial do pré-treinamento em FP4.