Tokens con embeddings casi nulos anulan el 20% de los pasos del ajuste fino de modelos gated en telugu

Si ajustas un modelo abierto con gates (Mamba, Gated DeltaNet o xLSTM, es decir, una arquitectura recurrente en la que unas «compuertas» regulan el acceso al estado acumulado), una parte notable de los pasos de optimización puede no modificar los pesos en absoluto. La causa no es un error en los datos, sino los tokens cuyos embeddings de entrada tienen una norma casi nula. En estos modelos gated, un solo token así amplifica los gradientes en decenas de órdenes de magnitud, y el clipping de gradientes estándar convierte esa explosión en una anulación total de la actualización.

En el continual pretraining en telugu se pierde de este modo cerca del 20% de los pasos, y los logs no muestran ni excepciones ni advertencias.

De dónde salen las filas «nulas» en los checkpoints

Los autores revisaron 291 repositorios abiertos de Hugging Face (280 tablas de embeddings distintas). Leyeron los pesos y el tokenizador directamente mediante peticiones HTTP byte-range. Para descartar las entradas especiales e inalcanzables, sometieron cada fila sospechosa a una comprobación: decodificaron el ID y lo volvieron a tokenizar, tanto de forma aislada como dentro de ocho plantillas de frases. Un token solo se consideraba alcanzable si el resultado era el mismo ID.

De los 291 repositorios, 92 contienen filas con una norma muy pequeña y 67 (23,0%) tienen tokens casi nulos alcanzables. Entre los 101 modelos con matrices de entrada y salida atadas no hay ni un solo token de este tipo. Entre los 190 modelos con matrices separadas aparecen en 67 (35,3%). En los Transformer convencionales la proporción es del 30,9%; en los modelos gated e híbridos, del 52,6% (20 de 38).

La causa está en el preentrenamiento. Un token que nunca ha aparecido en el texto no recibe gradiente en su fila de entrada. Pero el AdamW estándar aplica weight decay a todas las matrices bidimensionales, y la fila sin uso se va encogiendo hacia cero paso a paso.

La fila de salida, en cambio, recibe gradientes de softmax distintos de cero en cada paso y se mantiene sana: en el 92% de los 979 tokens comprobados (901), las normas de salida no están por debajo de la mediana del vocabulario. Con pesos atados, los gradientes de salida mantienen por sí solos una escala normal. En una muestra de 364 millones de tokens, las 251 filas casi nulas pertenecen a tokens con cero apariciones.

Los afectados son tokens corrientes: subpalabras del telugu y del euskera, palabras mongolas en cirílico y fragmentos de código con CRLF. Los modelos ajustados y destilados heredan estas filas de los checkpoints base.

Por qué sufren precisamente los modelos gated

En un Transformer, el problema no aparece en posiciones que no sean la primera. La primera capa de atención ya promedia los valores de los tokens anteriores, que tienen una escala normal, y devuelve la señal al rango de trabajo. Lo mismo ocurre si la secuencia empieza con un BOS.

En las recurrencias gated, la lectura del estado histórico se multiplica por una función del token actual: un gate o una query lineal. Con una entrada nula, el multiplicador también tiende a cero y la contribución del historial se comprime. El vector residual se mantiene casi nulo a lo largo de varios bloques seguidos, y el Pre-RMSNorm de cada capa multiplica los gradientes de retropropagación por un coeficiente grande. El efecto se acumula con la profundidad: los gradientes de los embeddings de entrada crecen en más de 30 órdenes de magnitud (hasta 35). Los híbridos topan con un techo; los modelos gated puros siguen creciendo.

Los autores comprobaron la causalidad con una intervención: abrieron el gate de la primera capa en el paso hacia delante conservando todas las derivadas de retropropagación, y el gradiente cayó en seco. Desactivar únicamente la derivada inversa del gate no eliminó el problema.

La anulación silenciosa del batch y dos formas de curarla

A continuación entra en juego el clipping por norma global. PyTorch, HuggingFace Trainer y FSDP2 suman los cuadrados de los gradientes en FP32. Si un solo valor es demasiado grande, al elevarlo al cuadrado se obtiene infinito. Los gradientes se dividen por una norma infinita y se convierten en ceros. AdamW da el paso solo con el momento anterior y el weight decay del resto de la red, mientras que la nueva fila casi nula, con momento cero, permanece congelada y arruina los batches siguientes. En inferencia y en rollouts de RL on-policy este fallo no se produce, pero una actualización de RL de parámetros completos con prompts que contienen chat tags casi nulos sí lo desencadena.

Los autores proponen dos soluciones sin sobrecoste en tiempo de paso ni en memoria pico:

• Clipping separado de los embeddings de entrada. La norma se calcula aparte del resto de los pesos, en aritmética de 64 bits o con reescalado, y requiere tres líneas de código. Las pérdidas de actualizaciones desaparecen y la fila afectada se «cura» en un solo paso. No hace falta escanear el vocabulario ni cambiar la arquitectura.
• Sesgo positivo en el gate para el preentrenamiento desde cero. El gate permanece positivo con una entrada nula, y la primera capa deja pasar un estado distinto de cero. Los gradientes casi nulos caen en más de 28 órdenes de magnitud en modelos de hasta 370M de parámetros, y la perplexity de validación es comparable a la del modelo base.

Desactivar el weight decay para los embeddings en checkpoints ya publicados no ayudará: la explosión ocurre en el primer paso. Y en el entrenamiento desde cero, esa desactivación por sí sola tampoco basta: los bloques siguen amplificando los gradientes cuando la inicialización de las normas es pequeña.

La conclusión práctica es sencilla. Si ajustas un checkpoint gated sin pesos atados en un idioma de pocos recursos o con tus propios chat tags, comprueba las normas de las filas de entrada de los embeddings y aplica su clipping por separado.