Resumen Técnico: Profundidad de Cálculo Adaptativa mediante Enrutamiento de Tokens Aprendido en Transformadores

Declaración del Problema

Las arquitecturas Transformer estándar aplican un número fijo de capas a cada token en una secuencia, independientemente de la dificultad contextual o la densidad de información de dicho token. Este diseño prioriza la simplicidad arquitectónica sobre la adaptabilidad por token, lo que resulta en una ineficiencia computacional significativa. En la práctica, los tokens comunes en contextos predecibles requieren mucho menos procesamiento que los tokens raros en construcciones novedosas, sin embargo, ambos reciben recursos de cálculo idénticos durante cada paso hacia adelante. Esta ineficiencia es particularmente costosa a escala de inferencia, donde el costo dominante es el paso completo hacia adelante a través de todas las capas para todos los tokens. Las soluciones existentes, como el Tiempo de Cálculo Adaptativo (ACT) para RNNs, los Transformadores Universales, la Mezcla de Profundidades (MoD) y los mecanismos de Salida Temprana, intentan abordar esto pero a menudo dependen de selecciones rígidas, presupuestos acumulados o umbrales de confianza post-hoc que pueden no ser completamente diferenciables u óptimos.

Metodología: Atención Selectiva de Tokens (TSA)

El artículo propone la Atención Selectiva de Tokens (TSA), un mecanismo que introduce una puerta aprendida, por token, en las actualizaciones residuales entre bloques de transformadores consecutivos.

Arquitectura y Enrutamiento

TSA inserta un enrutador ligero $r_l$ después de cada bloque $f_l$ (para $l = 0, \dots, L-2$) en un transformador solo-decodificador con pre-normalización. El bloque inicial $f_0$ siempre se ejecuta incondicionalmente.

• Diseño del Enrutador: Cada enrutador es una Perceptrón Multicapa (MLP) de dos capas con una salida sigmoide, produciendo una probabilidad de parada continua $p_l \in (0, 1)$ para cada token. La dimensión oculta es $d/4$ (redondeada hacia abajo a 16), añadiendo aproximadamente un 1.7% de sobrecarga de parámetros al modelo base.
• Actualización Puerta: Las actualizaciones residuales para el bloque subsiguiente $f_{l+1}$ se escalan por $(1 - p_l)$.
  • Si $p_l = 0$, la actualización es estándar.
  • Si $p_l = 1$, el estado permanece sin cambios, saltando efectivamente el bloque.
  • La interpolación es suave, asegurando que el mecanismo sea completamente diferenciable sin requerir estimadores de paso directo, muestreo Gumbel o aprendizaje por refuerzo.

Entrenamiento y Regularización

Durante el entrenamiento, todas las capas se ejecutan completamente; la puerta escala las actualizaciones residuales pero no omite el cálculo. La pérdida total se define como $L = L_{tarea} + L_{profundidad}$.

• Regularización de Profundidad: Se añade un término $L_{profundidad} = \lambda \cdot \frac{1}{L-1} \sum (1 - p_l)$ para fomentar suavemente la parada temprana.
• Comportamiento Emergente: Cabe destacar, el artículo demuestra que el mecanismo de enrutamiento aprende a omitir operaciones incluso cuando $\lambda = 0$ (sin presión explícita de profundidad), impulsado únicamente por el gradiente de la pérdida de la tarea.

Inferencia

En la inferencia, el mecanismo cambia a un modo "disperso". Los tokens con $p_l > 0.5$ se excluyen de la Red Neuronal de Alimentación Directa (FFN) mediante operaciones de recolección/dispersión, mientras que la atención permanece densa para preservar la semántica exacta clave-valor. Esto genera ahorros reales de FLOPs y aceleraciones en tiempo de reloj.

Contribuciones Clave

1. Enrutamiento de Tokens Diferenciable: Una puerta suave continua simple en las actualizaciones residuales que permite el enrutamiento por token y por capa sin cambios arquitectónicos en el Transformador base.
2. Emergencia Impulsada por Gradientes: Evidencia de que el enrutamiento proporcional a la dificultad emerge naturalmente solo del gradiente de la pérdida de la tarea. Incluso sin regularización de profundidad ($\lambda=0$), el enrutador aprende a omitir ~20% de las operaciones de token-capas.
3. Compensación Eficiencia-Calidad: Validación en modelado de lenguaje a nivel de carácter mostrando ahorros del 14–23% en operaciones de token-capas (TLOps) con una pérdida de calidad menor al 0.5%.
4. Superioridad sobre la Salida Temprana: A eficiencia igualada, TSA logra una pérdida de validación menor que las líneas base de Salida Temprana, lo que sugiere que las decisiones de enrutamiento aprendidas de extremo a extremo son superiores a la umbralización de confianza post-hoc.
5. Aceleración en Tiempo Real: Demostración de una aceleración real en tiempo de reloj (hasta un 2.3%) en hardware comercial (Apple M1 Pro) mediante inferencia dispersa.

Resultados Experimentales

Tareas Sintéticas

En tareas de copia y ordenación (secuencias de longitud 10):

• Tarea de Copia: El enrutador aprendió que la tarea era un mapeo identidad, omitiendo el 54.9% de las TLOps ($\alpha = 0.341$) sin pérdida de precisión.
• Tarea de Ordenación: El enrutador asignó más cómputo, omitiendo el 22.5% de las TLOps ($\alpha = 0.730$), reflejando la mayor dificultad.

Modelado de Lenguaje a Nivel de Carácter

• Tiny-Shakespeare: TSA ahorró un 22.8% de las TLOps ($\alpha = 0.726$) con un aumento en la pérdida de validación de solo 0.006 nats (+0.4%).
• enwik8: En un corpus más diverso (Wikipedia), TSA ahorró un 13.9% de las TLOps ($\alpha = 0.833$) sin degradación de calidad (pérdida de validación marginalmente mejor).
• Robustez: El mecanismo es robusto al hiperparámetro de regularización $\lambda$ a través de dos órdenes de magnitud ($[0, 0.1]$). En $\lambda=0.05$, logró ahorros del 50.4% en TLOps con un costo de calidad negligible.

Comparación con Salida Temprana

Con una fracción activa igualada ($\alpha \approx 0.726$), TSA logró una pérdida de validación un 0.71% menor que una línea base de Salida Temprana. Esto indica que aprender decisiones de enrutamiento de extremo a extremo produce mejores decisiones que aplicar un umbral de confianza en el momento de la inferencia.

Rendimiento en Tiempo de Reloj

• Soft-TSA (Entrenamiento): Añadió ~1% de sobrecarga independientemente de $\alpha$.
• Sparse-TSA (Inferencia): Logró una aceleración del 2.3% en Tiny-Shakespeare ($\alpha=0.726$) y un rendimiento de equilibrio en enwik8 ($\alpha=0.833$). Las aceleraciones requirieron un tamaño de lote de $\ge 64$ para superar la sobrecarga de sincronización CPU-GPU.

Significado y Afirmaciones

El artículo afirma que TSA ofrece una solución mínima y arquitectónicamente simple a la ineficiencia de los Transformadores de profundidad fija. Su principal significado radica en demostrar que la profundidad de cálculo adaptativa puede aprenderse directamente del objetivo de la tarea sin supervisión explícita sobre la dificultad o señales complejas de aprendizaje por refuerzo.

Los autores enfatizan que el enrutador actúa como una "puerta de ruido aprendida", atenuando las actualizaciones cuando las contribuciones residuales son redundantes. Esto permite que el modelo asigne dinámicamente el cómputo basado en la dificultad del token, logrando ganancias significativas de eficiencia (hasta un 55% en tareas sintéticas, 14–23% en modelado de lenguaje) con compensaciones de calidad mínimas. El trabajo sugiere que el enrutamiento diferenciable de extremo a extremo es una estrategia más efectiva para la profundidad adaptativa que la selección discreta o la umbralización post-hoc, allanando el camino para una inferencia más eficiente en hardware comercial.

El artículo se mantiene modesto respecto a la escala, señalando que todos los experimentos se realizaron en modelos con aproximadamente 5–6 millones de parámetros, y el comportamiento de escalado hacia modelos más grandes (10M–100M+) sigue siendo una pregunta abierta. Además, las aceleraciones actuales dependen del tamaño del lote y la dispersión de la atención aún no se ha implementado.