Resumo Técnico: BHyT – Uma Alternativa Estável e Eficiente à Normalização Pré-Camada (Pre-Layer Normalization)

1. Declaração do Problema

Os Grandes Modelos de Linguagem (LLMs) dependem fortemente da Normalização Pré-Camada (Pre-LN), tipicamente implementada via RMSNorm, para estabilizar o treinamento e permitir um aprendizado por transferência eficaz. No entanto, à medida que a profundidade do modelo aumenta, o design padrão de Pre-LN sofre de duas limitações críticas:

1. A Maldição da Profundidade: Análises recentes indicam que a interação entre conexões residuais e Pre-LN impulsiona um crescimento rápido na magnitude e variância dos estados ocultos através das camadas. Essa escalada empurra o Jacobiano do bloco em direção à matriz identidade, fazendo com que as camadas mais profundas se comportem como mapeamentos de identidade dispendiosos e desestabilizando o treinamento.
2. Compromisso entre Eficiência e Estabilidade: As soluções existentes frequentemente abordam apenas um aspecto. Métodos orientados à estabilidade, como o Peri-LN (normalizando antes e depois das subcamadas) ou o LayerNorm Scaling (LNS), suprimem o crescimento da variância, mas incorrem em um overhead computacional significativo devido a cálculos estatísticos repetidos. Por outro lado, métodos orientados à eficiência, como o Dynamic Tanh (DyT), removem o overhead estatístico, mas falham em controlar explicitamente o crescimento residual em profundidade, levando ao potencial de saturação e instabilidade em redes profundas.

O desafio central é projetar um mecanismo que combine o baixo overhead de transformações limitadas com o controle explícito da escala e variância do fluxo residual para evitar o desvio em profundidade sem sacrificar o rendimento (throughput) do treinamento.

2. Metodologia: Bounded Hyperbolic Tanh (BHyT)

Os autores propõem o BHyT, um substituto direto para o Pre-LN que acopla uma não-linearidade de tangente hiperbólica com um limitador de entrada baseado em dados.

Mecanismo Central

Ao contrário do DyT, que depende de um escalar global aprendível para reescalar entradas, o BHyT limita explicitamente a entrada da função $\tanh(\cdot)$ usando estatísticas dependentes dos dados.

• Limitação Probabilística de Entrada: Para evitar a saturação, o BHyT restringe o argumento pré-$\tanh$ para que ele permaneça dentro de um intervalo predefinido $[-\lambda, \lambda]$ com alta probabilidade. Usando a desigualdade de Chebyshev, o método deriva uma regra de escala baseada na média ($\mu_x$) e no desvio padrão ($s_x$) da entrada.
• Forma Ideal ($BHyT^*$): Teoricamente, a transformação reescala a entrada $x$ como:
  $$BHyT^*(x) = \gamma \odot \tanh\left(\frac{\lambda}{\kappa s_x + |\mu_x|} x\right)$$
  onde $\gamma$ é um parâmetro de escala aprendível, $\lambda$ define o intervalo alvo, e $\kappa$ controla o conservadorismo da escala com base em uma probabilidade alvo $p$. Isso garante que a entrada permaneça na região linear não-saturante da $\tanh$ com alta probabilidade.

Implementação Prática: Aproximação de Variância

Calcular a média e a variância exatas em cada local de Pre-LN dentro de um bloco Transformer anularia os ganhos de eficiência. Para resolver isso, o BHyT introduz uma aproximação de variância em nível de bloco:

1. Primeiro Local (Antes da Atenção): Calcula a variância exata por instância da entrada.
2. Segundo Local (Antes do MLP): Em vez de recomputar a variância exata da saída residual ($x + h_{Attn}$), aproxima a variância usando a variância da entrada previamente calculada e uma estimativa do modelo sobre a variância da saída da atenção.
   • A variância da saída da atenção é aproximada como $\tilde{s}^2_{hAttn} \approx \frac{1}{Td} \|W_V W_O\|_F^2 \cdot \frac{\lambda_{Attn}^2}{\kappa^2}$, assumindo pesos de atenção uniformes e entradas gaussianas de média zero.
   • Esta aproximação permite que a segunda normalização reutilize estatísticas ou utilize parâmetros pré-computados do modelo, evitando uma segunda operação de redução no caminho crítico de propagação (forward path).

Estabilidade Teórica

O artigo fornece um limite de propagação de variância de profundidade finita. O Teorema 3.5 afirma que, para uma rede de profundidade $L$, se os hiperparâmetros satisfizerem $\lambda/\kappa < 1/\sqrt{L}$, a variância de saída sob o BHyT é estritamente menor do que sob o LayerNorm Scaling (LNS) para cada camada até a profundidade $L$. Isso oferece uma garantia teórica de que o BHyT suprime o crescimento de variância em profundidade de forma mais eficaz que o LNS.

3. Principais Contribuições

1. Proposta do BHyT: Um substituto direto para o Pre-LN que combina uma transformação $\tanh(\cdot)$ limitada com um limitador de entrada baseado em dados para controlar o crescimento das ativações em profundidade.
2. Eficiência Teórica e Prática:
   • Derivação de um limite de propagação de variância de profundidade finita provando que o BHyT produz um limite de propagação de variância menor que o LNS sob condições específicas.
   • Introdução de uma aproximação de variância em nível de bloco que reduz as computações exatas de variância repetidas, diminuindo significativamente a latência.
3. Validação Empírica: Avaliação abrangente nos modelos Llama-374M, Llama-1B e Llama-3B. O estudo demonstra que o BHyT melhora o compromisso entre estabilidade-eficiência enquanto mantém um forte desempenho em tarefas subsequentes (downstream).

4. Resultados Experimentais

Os experimentos foram realizados em modelos do tipo Llama (374M, 1B e 3B) pré-treinados no corpus C4 com orçamentos de tokens proporcionais ao tamanho do modelo (aprox. 20 tokens por parâmetro).

• Estabilidade: O BHyT suprimiu com sucesso o crescimento em profundidade na magnitude e variância das ativações, superando o RMSNorm e o DyT. As estatísticas de ativação permaneceram controladas através das camadas, comparáveis a baselines orientados à estabilidade como o Peri-LN.
• Desempenho:
  • Pré-treinamento: O BHyT alcançou a menor perda (loss) e perplexidade de pré-treinamento em todas as escalas de modelo. Também obteve a maior precisão média em benchmarks de 3-shot (ex: ARC, MMLU, Winogrande), superando o segundo melhor método em até 0,6 pontos no modelo 1B.
  • Pós-SFT: Após o ajuste fino supervisionado (SFT) no Lima1K, os modelos BHyT mostraram menor perplexidade no C4 e WikiText-2 (indicando menos esquecimento e melhor generalização) e maior precisão em tarefas subsequentes comparado ao RMSNorm, Peri-LN, LNS e DyT.
  • Escala: Em uma comparação de 20 bilhões de tokens no Llama-3B, o BHyT superou o forte baseline orientado à estabilidade Peri-LN em 3,4 pontos percentuais em precisão média de tarefas subsequentes (pré-treinamento) e 2,5 pontos (pós-SFT).
• Eficiência:
  • Rendimento de Treinamento (Throughput): O BHyT alcançou uma média de 1,6% de velocidade de treinamento superior em relação ao RMSNorm.
  • Rendimento de Geração: O BHyT alcançou 1.199,9 tokens/s no Llama-1B, superando o RMSNorm em 1,77% e o Peri-LN em 18,0%. Embora o DyT tenha permanecido o mais rápido (devido à ausência de estatísticas), o BHyT ofereceu um compromisso superior entre velocidade e estabilidade.
  • Ablação: O mecanismo de aproximação de variância proporcionou um ganho de velocidade de 14,9% sobre a forma exata do BHyT* sem degradar materialmente a convergência ou o desempenho final.

5. Significância e Alegações

O artigo afirma que o BHyT oferece um caminho prático para modelos de fundação mais profundos e eficientes. Ao unificar o baixo overhead de transformações limitadas com o controle explícito da variância do fluxo residual, o BHyT aborda a "maldição da profundidade" sem as penalidades computacionais de esquemas de dupla normalização.

Os autores sustentam que o treinamento escalável pode melhorar a eficiência sem sacrificar a estabilidade. O BHyT é apresentado não apenas como uma alternativa teórica, mas como um substituto robusto e direto para o RMSNorm que permite um pré-treinamento estável e um forte desempenho de transferência em várias escalas de modelo (de 374M a 3B parâmetros), reduzindo simultaneamente a latência tanto no treinamento quanto na inferência.