Per-Token Fixed-Point Convergence in Depth-Recurrent Transformers

Resumen Técnico: Convergencia de Punto Fijo por Token en Transformers de Profundidad Recurrente

Planteamiento del Problema
Los transformers estándar de profundidad fija vinculan la calidad del modelo a un único punto operativo de cómputo. El servicio de múltiples niveles de calidad/costo requiere entrenar y almacenar modelos separados, mientras que adaptar el cómputo a la dificultad de la entrada durante la inferencia sacrifica la caché de prompts. Los transformers de profundidad recurrente (que aplican un núcleo de pesos compartidos un número variable de veces) ofrecen una solución teórica, permitiendo que un único conjunto de parámetros opere a través de un rango de profundidades. Las aproximaciones existentes para la asignación de profundidad por token sufren limitaciones significativas: los métodos de salida temprana (early-exit) en pilas no compartidas fallan debido a inconsistencias en la caché de clave-valor; los métodos de mezcla de recursiones (mixture-of-recursions) fijan el presupuesto de cómputo en los pesos durante el preentrenamiento, lo que impide el reajuste post-hoc; y las políticas de parada entrenadas conjuntamente suelen colapsar hacia profundidades degeneradas o requieren un ajuste delicado. Crucialmente, ningún trabajo previo ha medido directamente el cómputo por token de un modelo elástico entrenado para determinar si una política simple y libre de entrenamiento podría superar a las aprendidas.

Metodología
Los autores proponen un enfoque de dos etapas que separa la robustez de la profundidad de la asignación de la misma:

1. Robustez de la Profundidad (Preentrenamiento): Se entrena un backbone de profundidad aleatorizada utilizando una receta mínima. En cada paso del optimizador, se muestrea un recuento de recursión $r$ de una distribución log-normal y se aplica a todo el paso. El modelo utiliza una pérdida de entropía cruzada estándar sin unidades de parada, pérdida de ponderación (ponder loss) o penalizaciones de FLOPs. La arquitectura consiste en un preludio único, un núcleo de pesos compartidos (con reinyección de entrada y normalización tipo sándwich para evitar el colapso del estado) y una coda única.
2. Asignación de Profundidad (Política Post-hoc): Una vez congelado el backbone, se evalúan tres políticas sobre el mismo modelo para determinar la profundidad por token:
   • A0 (Profundidad Uniforme): Aplica un recuento de bucle $r$ fijo a todos los tokens (línea base).
   • A1 (Salida por Convergencia): Una regla libre de entrenamiento que detiene un token una vez que la divergencia KL entre distribuciones de salida sucesivas ($KL(p_i || p_{i-1})$) cae por debajo de un umbral $\epsilon$.
   • A2 (Router Etiquetado por Oráculo): Un probe lineal entrenado en los propios estados ocultos del modelo para predecir si la predicción argmax de un token cambiará en bucles futuros.

Contribuciones Clave

1. Medición Directa por Token: El estudio proporciona la primera caracterización de la convergencia por token en un modelo de lenguaje de profundidad recurrente. Demuestra que el estado recurrente converge a un punto fijo por token. Mientras que el token mediano converge en el sexto bucle, una cola de aproximadamente el 10% de los tokens continúa actualizándose en la profundidad media de entrenamiento de ocho. La profundidad de convergencia está ordenada por tipo de token (los espacios en blanco convergen más superficialmente, las palabras de contenido más profundamente).
2. Lectura vs. Aprendizaje de la Convergencia: El artículo demuestra que leer directamente la dinámica del estado del modelo supera a aprender a predecirlas. La salida de convergencia libre de entrenamiento (A1) logra la calidad de la profundidad-8 uniforme en un promedio de 4.94 bucles (una reducción del 38% en la profundidad promedio). En contraste, el router lineal (A2), entrenado con etiquetas de convergencia extraídas del mismo modelo, requiere casi la profundidad completa (7.99 bucles) y no ofrece reducción alguna.
3. Hipótesis de Saturación de la Media de Entrenamiento: Respaldados por datos internos y verificaciones de consistencia externa, los autores hipotetizan que la profundidad útil de un modelo recursivo se satura en su distribución de bucles media de tiempo de entrenamiento. Los modelos entrenados con una media de 4 se saturan en la profundidad 4; aquellos con una media de 8 se saturan en la profundidad 8.

Resultados

• Elasticidad: El backbone de profundidad aleatorizada exhibe una mejora de calidad monótona desde la profundidad 1 hasta la 8, tras la cual la pérdida de validación se estabiliza (3.20 nats) incluso hasta los 32 bucles. Esto confirma que el modelo es robusto en profundidad sin necesidad de un ajuste específico para cada profundidad.
• Dinámicas de Convergencia: La divergencia KL media de las salidas sucesivas cae de $3.9 \times 10^{-1}$ en el bucle 2 a $8.5 \times 10^{-6}$ en el bucle 16. El cambio de estado decae por paso, confirmando que el sistema alcanza un punto fijo en lugar de girar inútilmente.
• Eficiencia de Asignación: La política A1, libre de entrenamiento, iguala la calidad de un modelo de profundidad-8 uniforme en 4.94 bucles promedio. Supera a la línea base uniforme en profundidades bajas y la iguala en el rango de profundidad promedio. El router aprendido (A2) falla al no proporcionar ahorros, requiriendo casi la profundidad completa para igualar la línea base.
• Comportamiento a Nivel de Token: Los tokens estructurales (espacios en blanco) salen más temprano, mientras que los tokens de contenido computan durante más tiempo. El orden relativo de números y palabras se invierte entre conjuntos de datos (TinyStories vs. FineWeb-Edu), atribuido a cambios en la distribución más que a un fallo arquitectónico.
• Límites de Eficiencia: Aunque el estudio reporta una reducción del 38% en la profundidad promedio (un proxy de FLOPs), señala explícitamente que la aceleración de tiempo de ejecución (wall-clock) realizada depende del procesamiento por lotes continuo por profundidad (depth-wise continuous batching), el cual no fue implementado. La mejor aceleración de tiempo de ejecución estimada es de 1.48–1.51x, mientras que el entorno actual muestra un 0% de aceleración realizada debido al enmascaramiento en lugar de saltarse el trabajo.

Significancia y Reivindicaciones
El artículo afirma que la "elasticidad" de los transformers de profundidad recurrente es una propiedad estable que puede aprovecharse mediante políticas simples y libres de entrenamiento. El hallazgo central es que la variación por token en la profundidad de convergencia es directamente legible y explotable, lo que hace que las complejas políticas de parada aprendidas sean innecesarias y menos efectivas. Los autores posicionan este trabajo como un estudio de medición y asignación más que como una contribución al backbone elástico subyacente (al cual atribuyen a trabajos previos como Huginn y LoopFormer).

El estudio es modesto en su alcance:

• Los resultados se establecen en una escala única (parámetros de la clase 135M) y una única semilla (para el modelo principal).
• No se reclama paridad de FLOPs igualada contra transformers estándar; las comparaciones se hacen contra una referencia de techo (SmolLM2-135M) que tiene una ventaja de 150x en el conteo de tokens.
• La aceleración de tiempo de ejecución realizada no se reporta como una métrica final, sino como un intervalo entre la reducción de FLOPs y el potencial de mejor caso de procesamiento por lotes.
• La hipótesis de que la profundidad de saturación es igual a la profundidad media de entrenamiento se presenta como una hipótesis abierta apoyada por datos limitados, no como una ley probada.

El estudio completo fue reproducible en una sola RTX 4090 en aproximadamente 100 horas de GPU.