Resumen Técnico: Enrutamiento Residual Multirresolución para Transformers de Decodificador Único Profundos

Planteamiento del Problema

Los Transformers de decodificador único estándar dependen de conexiones residuales PreNorm donde las actualizaciones de las subcapas se agregan con pesos unitarios fijos. Si bien esta estructura permite un entrenamiento estable, impone una regla de acumulación rígida que puede diluir las contribuciones de las capas individuales a medida que aumenta la profundidad de la red. Avances recientes, como los Residuales de Atención (Attention Residuals) y los Residuales de Atención de Bloque (Block AttnRes), abordan esto reemplazando la acumulación fija por un enrutamiento aprendido y dependiente del contenido sobre las salidas previas. Específicamente, Block AttnRes mejora la eficiencia comprimiendo múltiples capas en resúmenes residuales a nivel de bloque ($C_b$).

Sin embargo, los autores identifican una limitación crítica en Block AttnRes: un único resumen de bloque representa únicamente el componente de baja frecuencia (DC) de la trayectoria residual dentro de un bloque. Al almacenar solo la suma de las actualizaciones ($\sum u_{b,i}$), el mecanismo descarta información estructural direccional valiosa, tal como:

1. Desequilibrio entre subcapas: Si el bloque está dominado por las actualizaciones de Atención o de MLP.
2. Dinámicas temporales: Si las actualizaciones tempranas dentro del bloque apuntan en una dirección diferente a las actualizaciones tardías.

Los autores plantean la hipótesis de que, a medida que la longitud de la trayectoria residual crece, esta información direccional perdida se vuelve cada vez más vital para un enrutamiento efectivo por profundidad.

Metodología: WAV v1

Los autores proponen WAV v1 (Waveform Attention v1), una extensión ligera y multirresolución de Block AttnRes. En lugar de representar un bloque únicamente por su suma total, WAV v1 aumenta cada representación de bloque con dos bases de detalle direccional de suma cero, preservando al mismo tiempo los módulos originales de Atención y MLP.

1. Representación Multirresolución

Para un bloque $b$ que contiene $m$ actualizaciones de subcapa $\{u_{b,i}\}_{i=1}^m$, WAV v1 calcula:

• Base de Estado ($C_b$): La suma de bloque estándar, que representa el estado de baja frecuencia.
  $$C_b = \sum_{i=1}^m u_{b,i}$$
• Base de Fase ($D^{phase}_b$): Una fuente de detalle que contrasta las actualizaciones de Atención frente a las de MLP. Asigna un signo de $+1$ a las actualizaciones de Atención y $-1$ a las de MLP.
  $$D^{phase}_b = \sum_{i=1}^m a_i u_{b,i}, \quad \text{donde } a_i \in \{+1, -1\}$$
• Base de División ($D^{split}_b$): Una fuente de detalle que contrasta la primera mitad frente a la segunda mitad del bloque. Asigna $+1$ a la primera mitad y $-1$ a la segunda mitad.
  $$D^{split}_b = \sum_{i=1}^m r_i u_{b,i}, \quad \text{donde } r_i \in \{+1, -1\}$$

Estas bases se acumulan en línea a medida que se producen las actualizaciones, lo que conlleva una sobrecarga computacional insignificante.

2. Mecanismo de Enrutamiento

El mezclador softmax de profundidad recibe un conjunto de fuentes expandido. Para Block AttnRes, el conjunto es $\{e, C_0, \dots, C_{b-1}, PC\}$. Para WAV v1, este se expande para incluir las bases de detalle de cada bloque completado y el bloque parcial actual:
$$S_{WAV} = \{e, C_0, \tilde{D}^{phase}_0, \tilde{D}^{split}_0, \dots, C_{b-1}, \tilde{D}^{phase}_{b-1}, \tilde{D}^{split}_{b-1}, PC, \tilde{P}^{phase}, \tilde{P}^{split}\}$$

3. Mecanismos de Estabilidad

Para evitar la desestabilización durante las primeras etapas del entrenamiento, WAV v1 emplea dos estrategias conservadoras:

• Sesgo de Detalle Negativo: Las fuentes de detalle se inicializan con un fuerte sesgo negativo ($\beta_D = -2.0$), asegurando que el modelo se comporte como el Block AttnRes estándar en la inicialización y solo utilice los detalles cuando sea beneficioso.
• Emparejamiento de RMS Desacoplado: Los tensores de detalle se escalan en relación con su correspondiente resumen de bloque ($C$) utilizando una relación de RMS desacoplada, limitada por un factor de escala máximo $\rho$. Esto evita que los detalles dominen simplemente debido a magnitudes brutas mayores.
• Lectura Conservadora: Por defecto, el mezclador de predicción final lee solo las fuentes de embedding y $C$, no las fuentes de detalle, lo que garantiza aún más la estabilidad.

Contribuciones Clave

1. Formulación: El artículo enmarca el enrutamiento residual de bloque como un problema de representación multirresolución, distinguiendo entre la información de estado de baja frecuencia ($C_b$) y la información de detalle direccional intra-bloque.
2. Algoritmo: Introducción de WAV v1, una extensión de aplicación directa (drop-in) que aumenta los resúmenes de bloque con bases de detalle de fase y división sin modificar las funciones centrales de Atención o MLP.
3. Evidencia Empírica: Experimentos de escalado preliminares que demuestran que la utilidad del enrutamiento multirresolución depende de la profundidad, siendo significativamente beneficiosa solo en arquitecturas más profundas.

Resultados Experimentales

Los autores evaluaron modelos GPT de decodificador único de 12, 24 y 48 capas en los conjuntos de datos a nivel de carácter TinyStories y Text8.

• Tendencia Dependiente de la Profundidad:
  • 12 Capas: WAV v1 rinde por debajo de Block AttnRes, lo que sugiere que los detalles direccionales son redundantes o ruidosos en redes poco profundas.
  • 24 Capas: WAV v1 se vuelve competitivo, mostrando mejoras leves en TinyStories y un rendimiento comparable en Text8.
  • 48 Capas: WAV v1 logra la mejor pérdida de validación en ambos conjuntos de datos, superando sustancialmente a Block AttnRes, ReZero y LayerScale.
• Ganancias Cuantitativas (48 Capas):
  • TinyStories: Redujo la pérdida de validación de 0.4960 (Block AttnRes) a 0.4738 (WAV v1), una reducción de la perplejidad (PPL) del 2.19%.
  • Text8: Redujo la pérdida de validación de 0.9363 a 0.9305, una reducción de la PPL del 0.57%.
• Dinámica de Entrenamiento: A 48 capas, WAV v1 se separa de Block AttnRes temprano en el entrenamiento y mantiene una ventaja persistente, lo que indica que la mejora no es un artefacto de la etapa final, sino un cambio fundamental en la trayectoria de optimización.

Significancia y Reivindicaciones

El artículo afirma que el enrutamiento residual profundo se beneficia no solo de seleccionar qué bloque leer, sino también de enrutar la estructura direccional interna de cada bloque. Los resultados sugieren que, a medida que la trayectoria residual se alarga, la "forma" de la trayectoria de actualización (por ejemplo, la dominancia de la atención o la dinámica temprana-vs-tardía) se convierte en una señal informativa para el mezclador de profundidad.

Los autores posicionan a WAV v1 como un enfoque complementario a los métodos existentes:

• A diferencia de ReZero o LayerScale, que modulan la magnitud de las actualizaciones, WAV v1 cambia la representación disponible para el enrutador.
• A diferencia de Delta Attention Residuals, que realiza el enrutamiento sobre deltas individuales (alto costo), WAV v1 ofrece una aproximación multirresolución a nivel de bloque que equilibra la eficiencia con la riqueza estructural.

Los autores concluyen que las evaluaciones futuras de los métodos de enrutamiento residual deberían considerar la profundidad, ya que los beneficios de la información multirresolución pueden estar ligados a la longitud y estructura de la trayectoria residual. Reconocen limitaciones actuales, incluyendo la falta de desviación estándar entre semillas, la restricción a modelos pequeños a nivel de carácter y el uso de bases de detalle diseñadas a mano, lo que sugiere que el trabajo futuro debería explorar bases aprendibles y modelos de mayor escala a nivel de token.