Tokens com embeddings quase nulos zeraram 20% dos passos de fine-tuning de modelos gated em telugu

Se você faz fine-tuning de um modelo aberto com gates — Mamba, Gated DeltaNet ou xLSTM, ou seja, uma arquitetura recorrente em que «portas» (gates) controlam o acesso ao estado acumulado —, uma parte considerável dos passos de otimização pode não alterar os pesos de forma alguma. O culpado não é um erro nos dados, mas tokens cujos embeddings de entrada têm norma quase nula. Nesses modelos gated, um único token desse tipo amplifica os gradientes em dezenas de ordens de magnitude, e o clipping padrão de gradientes transforma essa explosão em um zeramento completo da atualização. No continual pretraining em telugu, cerca de 20% dos passos se perdem assim, sem exceções nem avisos nos logs.

De onde vêm as linhas «nulas» nos checkpoints

Os autores verificaram 291 repositórios abertos do Hugging Face (280 tabelas de embeddings distintas). Os pesos e o tokenizador foram lidos diretamente, por requisições HTTP byte-range. Para descartar entradas de serviço e tokens inalcançáveis, cada linha suspeita passou por uma checagem: o ID era decodificado e tokenizado de novo, isoladamente e em oito modelos de frase. Um token só era considerado alcançável se resultasse no mesmo ID.

Dos 291 repositórios, 92 contêm linhas com norma muito pequena, e 67 (23,0%) têm tokens quase nulos alcançáveis. Entre as 101 modelos com matrizes de entrada e saída amarradas, não há nenhum caso. Entre as 190 modelos com matrizes separadas, 67 (35,3%) apresentam o problema. Nos Transformers convencionais, a proporção é de 30,9%; nos modelos gated e híbridos, de 52,6% (20 de 38).

A causa está no pretraining. Um token que nunca apareceu no texto não recebe gradiente na linha de entrada. Já o AdamW padrão aplica weight decay a todas as matrizes bidimensionais, e a linha não utilizada vai encolhendo rumo a zero, passo a passo.

A linha de saída, por sua vez, recebe gradientes de softmax diferentes de zero em todos os passos e permanece saudável: em 92% dos 979 tokens verificados (901), as normas de saída não ficam abaixo da mediana do vocabulário. Com pesos amarrados, os gradientes de saída mantêm sozinhos a escala normal. Em uma amostra de 364 milhões de tokens, todas as 251 linhas quase nulas pertencem a tokens com zero ocorrências.

As linhas afetadas são comuns: subpalavras de telugu e basco, palavras mongóis em cirílico, trechos de código com CRLF. Modelos ajustados e destilados herdam essas linhas dos checkpoints de base.

Por que os modelos gated sofrem mais

Em um Transformer, fora da primeira posição, o problema não existe. Já a primeira camada de atenção faz a média dos valores dos tokens anteriores, que têm escala normal, e devolve o sinal à faixa de operação. O mesmo acontece se houver um BOS no início da sequência.

Nas recorrências gated, a leitura do estado histórico é multiplicada por uma função do token atual — um gate ou uma query linear. Com entrada nula, o multiplicador também vai a zero, e a contribuição do histórico encolhe. O vetor residual permanece quase nulo por vários blocos seguidos, e o Pre-RMSNorm de cada camada multiplica os gradientes retropropagados por um coeficiente grande. O efeito se acumula com a profundidade: os gradientes dos embeddings de entrada crescem em mais de 30 ordens de magnitude (até 35). Os híbridos chegam a um teto, enquanto os modelos puramente gated continuam crescendo.

Os autores verificaram a causalidade por intervenção: abriram o gate da primeira camada no forward pass, mantendo todas as derivadas do backward, e o gradiente despencou. Desativar apenas a derivada de retorno do gate não eliminou o problema.

O zeramento silencioso do lote e duas formas de tratamento

Em seguida entra em ação o clipping pela norma global. PyTorch, HuggingFace Trainer e FSDP2 somam os quadrados dos gradientes em FP32. Se uma única entrada for grande demais, o quadrado resulta em infinito. Os gradientes são divididos pela norma infinita e viram zeros. O AdamW dá o passo apenas com base no momento antigo e no weight decay do restante da rede, enquanto a nova linha quase nula, com momento zero, permanece congelada e compromete os lotes seguintes. Na inferência e nos rollouts de RL on-policy essa falha não ocorre, mas uma atualização de RL com todos os parâmetros, em prompts com chat tags quase nulas, a dispara.

Os autores propõem duas soluções sem custo adicional no tempo de passo nem no pico de memória:

• Clipping separado dos embeddings de entrada. A norma é calculada à parte dos demais pesos, em aritmética de 64 bits ou com reescalonamento, e isso exige três linhas de código. As perdas de atualização desaparecem, e a linha afetada é «curada» em um passo. Não é preciso varrer o vocabulário nem alterar a arquitetura.
• Viés positivo no gate para pretraining do zero. O gate permanece positivo com entrada nula, e a primeira camada deixa passar um estado diferente de zero. Os gradientes quase nulos caem em mais de 28 ordens de magnitude em modelos de até 370M de parâmetros, e a perplexity de validação fica comparável à da base.

Desativar o weight decay para os embeddings em checkpoints já lançados não ajuda: a explosão acontece no primeiro passo. E, no treinamento do zero, essa desativação isolada também não basta — os blocos continuam amplificando os gradientes quando a inicialização das normas é pequena.

A conclusão prática é simples. Se você faz fine-tuning de um checkpoint gated sem pesos amarrados em um idioma de poucos recursos ou com chat tags próprias, verifique as normas das linhas de entrada dos embeddings e aplique o clipping delas separadamente.