Resumo Técnico: Remoção de Normalização Portada e Ancoragem de Escala em Transformers Pré-Norm

Declaração do Problema

Camadas de normalização (por exemplo, LayerNorm, RMSNorm) são componentes padrão em arquiteturas de transformadores, estabilizando ativações e permitindo o treinamento de redes profundas. No entanto, essas camadas introduzem computações não lineares dependentes de amostra que acarretam sobrecarga significativa durante a inferência. Permanece incerto se essas estatísticas por token são estritamente necessárias durante toda a fase de inferência, particularmente em transformers pré-norm. Embora trabalhos anteriores tenham tentado remover a normalização por meio de truques de inicialização ou ajuste fino sequencial, há uma falta de compreensão sobre os papéis funcionais distintos das camadas de normalização internas versus finais, especificamente no que diz respeito a como elas controlam a escala da representação pré-logit.

Metodologia

Os autores propõem o TaperNorm, uma abordagem de remoção de normalização portada projetada para transformers pré-norm. Este método permite que o modelo transite da normalização padrão para mapas lineares ou afins aprendidos e independentes de amostra durante o treinamento, permitindo a remoção das operações de normalização na inferência.

Mecanismo Central: TaperNorm

O TaperNorm interpola dinamicamente entre um ramo normalizado e um ramo de escala fixa usando um portão global $g \in [0, 1]$. Para um vetor de token $h$:
$$\text{TaperNorm}(h; g) = g \cdot \frac{h}{r(h)} D\gamma + (1 - g) \cdot c \cdot h D\tilde{\gamma}$$
Onde:

• $r(h)$ é o fator de escala da RMSNorm.
• $D\gamma$ e $D\tilde{\gamma}$ são matrizes de escala diagonais aprendíveis.
• $c$ é um escalar por camada.
• $g$ é um portão global que começa em 1 (normalização padrão) e é decaído por cosseno até 0 durante o treinamento.

Procedimento de Treinamento:

1. Aquecimento (Warmup): O portão $g$ é mantido em 1. O modelo treina como um transformer pré-norm padrão.
2. Calibração: No final do aquecimento, o escalar $c$ é computado para minimizar a diferença de mínimos quadrados entre o ramo normalizado e o ramo de escala, garantindo uma transição suave sem deslocamento de distribuição.
3. Afinação (Tapering): O portão $g$ é decaído até 0. Os parâmetros do ramo de escala ($\tilde{\gamma}$) permanecem treináveis.
4. Dobramento na Inferência (Inference Folding): Uma vez que $g=0$, a camada TaperNorm torna-se um simples mapa linear/afim ($c \cdot h D\tilde{\gamma}$). Esta operação pode ser matematicamente dobrada em projeções lineares adjacentes (por exemplo, projeções QKV ou pesos MLP), removendo a operação explícita de normalização do grafo de inferência.

Análise de Ancoragem de Escala

O artigo identifica uma distinção crítica entre camadas de normalização internas e finais:

• Normalização Final como Âncora: A camada de normalização final atua como uma âncora de escala implícita. Matematicamente, ela é 0-homogênea, o que significa que remove o componente de gradiente radial na saída. Sem essa âncora, a minimização da entropia cruzada pode levar o modelo a aumentar indefinidamente as magnitudes dos logits ("perseguição de logits") para reduzir a perda, levando à instabilidade.
• Alternativa Explícita: Quando a normalização final é removida (em configurações "All-Taper"), os autores introduzem uma perda de escala de alvo fixo ($L_{aux}$). Esta perda auxiliar penaliza desvios das escalas de token de um alvo fixo $s_{tgt}$, fornecendo uma força restauradora radial explícita que contrabalança a tendência de perseguição de logits.

Duas configurações experimentais são testadas:

1. Taper Interno: Normalizadores de blocos internos são afinados para mapas lineares, mas a normalização final permanece ativa (âncora implícita).
2. Taper Total (All-Taper): Todos os normalizadores, incluindo o final, são afinados. Isso requer a perda de escala de alvo fixo para manter a estabilidade (âncora explícita).

Principais Contribuições

1. Remoção de Normalização Portada: Um método para afinar módulos RMSNorm/LayerNorm para reescalagens lineares/afins fixas, permitindo a eliminação de estatísticas por token na inferência através do dobramento de pesos.
2. Teoria de Ancoragem de Escala de Saída: Uma análise demonstrando que o papel distinto primário da camada de normalização final é remover gradientes radiais. O artigo mostra que, sem essa âncora, a entropia cruzada induz crescimento ilimitado de logits, e que uma perda de escala de alvo fixo pode efetivamente substituir esse papel estabilizador.
3. Validação Empírica: Experimentos no pré-treinamento de TinyStories e no ajuste fino de GPT-2 mostram que a normalização interna pode ser removida com aumentos mínimos na perda de validação. A abordagem alcança até 1,18× maior vazão na decodificação autoregressiva com cache KV após o dobramento de escalas internas.

Resultados

• Pré-treinamento (TinyStories):
  • Taper Interno (+aux): Alcançou perdas de validação dentro de 0,7–1,5% da linha de base em todos os tamanhos de modelo (1M a 30M parâmetros).
  • Taper Total (+aux): Alcançou perdas dentro de 0,8–1,8% da linha de base, demonstrando que um modelo totalmente livre de normalização é viável com ancoragem de escala explícita.
  • Dinâmicas: Modelos sem uma âncora de escala exibiram crescimento rápido nas magnitudes dos logits ("perseguição de logits"). A perda de escala de alvo fixo suprimiu com sucesso esse comportamento, alinhando as magnitudes dos gradientes com a linha de base ancorada.
• Vazão de Inferência:
  • Em uma NVIDIA H100, a configuração Taper Interno (fundido) alcançou 1,15–1,18× maior vazão de tokens por segundo em comparação com a linha de base RMSNorm padrão. A configuração Taper Interno (não fundido) (operações de escala explícitas) alcançou aceleração de 1,10–1,14×.
• Ajuste Fino (GPT-2):
  • Os métodos propostos Internal-TaperLN e All-TaperLN alcançaram perdas de entropia cruzada competitivas com a abordagem "FakeLN" (Baroni et al., 2025) e modelos Vanilla padrão nos conjuntos de dados OpenWebText e The Pile.
  • Diferentemente de métodos de remoção sequencial, o TaperNorm usa um único portão global para remover todas as camadas selecionadas simultaneamente.

Significado e Alegações

O artigo alega que as camadas de normalização interna e final desempenham papéis distintos em transformers pré-norm. Enquanto normalizadores internos podem ser substituídos por mapas lineares aprendidos independentes de amostra com degradação de desempenho negligenciável, a camada de normalização final serve a uma função única como uma âncora de escala de saída.

O significado deste trabalho reside em:

1. Eficiência: Fornece um caminho prático para remover operações de normalização computacionalmente caras do grafo de inferência, gerando ganhos mensuráveis de vazão (até 1,18×) sem retreinamento do zero.
2. Insight Teórico: Esclarece que a estabilidade fornecida pela camada de normalização final deve-se largamente à sua remoção de gradientes radiais. Esse insight permite o projeto de modelos "totalmente livres de normalização" substituindo a âncora implícita por uma perda de escala de alvo fixo explícita.
3. Robustez: A abordagem mantém dinâmicas de pré-norm padrão no início do treinamento, transitando para mapas lineares apenas mais tarde, o que parece mais estável do que métodos que tentam treinar sem normalização desde o início.

Os autores permanecem modestos quanto ao escopo, observando que os experimentos são limitados a modelos de até 30M parâmetros e pilhas de hardware específicas. Eles reconhecem que a robustez a outliers de ativação em regimes de bilhões de parâmetros e a aplicabilidade a outros objetivos além da entropia cruzada permanecem questões em aberto para trabalhos futuros.