Embeddings de entrada quase zerados podem travar o treinamento de modelos com portas, segundo estudo

Um modelo de linguagem aprende a partir de vetores que representam cada token (palavra ou fragmento de texto). Esses vetores de entrada são ajustados a cada passo de treinamento. Segundo um preprint publicado no arXiv, em alguns modelos esses vetores podem chegar a valores quase nulos (normas abaixo de 10⁻⁶) e ficar presos ali, sem que o treinamento avise de nenhum erro.

O que o estudo levantou

Os autores analisaram 291 repositórios de modelos de linguagem abertos, com 280 tabelas de embeddings distintas. Entre os checkpoints untied com portas ou híbridos, 20 de 38 (52,6%) tinham embeddings de entrada acessíveis com norma abaixo de 10⁻⁶. Entre os modelos com pesos compartilhados (weight-tied), a fonte aponta 0%.

Segundo a fonte, esses vetores quase zerados podem surgir na pré-treinamento quando tokens raros ou de template não recebem gradiente de entrada e o decaimento de peso do AdamW encolhe os embeddings. Em configurações de Transformer testadas, posições não iniciais ou um token BOS evitaram esse problema.

Por que modelos com portas são mais afetados

Em modelos com portas citados no estudo (Mamba, Gated DeltaNet e xLSTM), um token quase zerado pode amplificar gradientes de embedding em mais de 30 ordens de grandeza. Quando a entrada é zero, os fatores de leitura do token atual (portas SiLU ou consultas lineares) se anulam, e a leitura do histórico é suprimida. Com isso, os valores RMS residuais podem cair abaixo da raiz quadrada do piso do RMSNorm em camadas consecutivas, e o Pre-RMSNorm passa a amplificar gradientes de retropropagação em até 10³ por camada.

O ponto mais prático é o recorte de gradiente. Segundo o texto, quando uma entrada do gradiente ultrapassa 1,84 × 10¹⁹, o recorte global em FP32 usado em PyTorch, Hugging Face Trainer e FSDP2 pode estourar para infinito e zerar os gradientes do lote atual. Na continuação de pré-treinamento em telugu do Nemotron-3-Nano-4B, isso descartou atualizações em cerca de 20% dos passos.

As duas propostas

• GroupClip (para checkpoints já lançados): separa o recorte dos embeddings de entrada, usando norma em 64 bits ou reescalada. Segundo os autores, a versão reescalada eliminou a perda de atualizações e recuperou as linhas afetadas em um passo, com sobrecarga de ponta a ponta medida como negligenciável.
• GateFloor (para pré-treinamento): usa SiLU(z) + δ, com δ entre 0,01 e 0,1, para manter a porta positiva com entrada zero. Nos experimentos com até 370 milhões de parâmetros, a fonte relata redução de 28 ordens de grandeza nos gradientes quase nulos, com perplexidade comparável à da linha de base.

O que ainda não está claro

O resumo não descreve um procedimento detalhado para verificar os próprios modelos; o critério citado é a norma dos embeddings de entrada abaixo de 10⁻⁶. Os experimentos de GateFloor foram até 370 milhões de parâmetros, e o caso de recuperação usa um único modelo, o Nemotron-3-Nano-4B. Vale ler o artigo completo antes de generalizar.