Analogies between Transformer Layers and Power Method

Resumen Técnico: Analogías entre Capas de Transformadores y el Método de la Potencia

Enunciado del Problema
El artículo aborda la falta de comprensión sobre la dirección vectorial específica hacia la cual convergen las representaciones de tokens en transformadores profundos, un fenómeno conocido como "suavizado excesivo", "colapso de rango" o "uniformidad de tokens". Si bien está establecido que los transformadores profundos tienden hacia un estado de consenso donde todos los tokens se vuelven idénticos, la naturaleza de este vector de consenso y el mecanismo que lo impulsa permanecen poco claros. La literatura existente a menudo modela la autoatención como un sistema de tiempo continuo o se centra en escenarios de una sola cabeza, dejando menos explorado el comportamiento de tiempo discreto y múltiples cabezas de los modelos prácticos (como GPT2, BERT y ALBERT). Específicamente, el artículo busca determinar si el proceso iterativo de pasar tokens a través de las capas del transformador se asemeja a un método numérico conocido para encontrar vectores propios principales.

Metodología
Los autores modelan un transformador como un sistema dinámico de tiempo discreto donde el índice de capa $t$ representa el tiempo. Se centran en la ley de actualización de los vectores de tokens $x_i(t)$ a través de una sola capa, excluyendo explícitamente la Red de Alimentación Directa (FFN) y utilizando normalización $L_2$ (reemplazando la Normalización de Capa estándar) para aislar los efectos geométricos de los mecanismos de atención y proyección.

El núcleo de la metodología implica definir una "matriz de capa" $F_t$ para cada capa $t$, construida como el producto de la matriz de pesos de salida ($O$) y las matrices de pesos de valor concatenadas ($V$) de las cabezas de atención:
$$F_t = O_t V_t$$
La regla de actualización se analiza como un paso en un sistema dinámico:
$$x_i(t+1) = \text{LayerNorm}\left(x_i(t) + \text{SalidaAtención}\right)$$
Los autores establecen una analogía directa entre esta actualización y el Método de la Potencia, un algoritmo iterativo utilizado para calcular el vector propio principal de una matriz. En el Método de la Potencia, un vector se multiplica repetidamente por una matriz y se normaliza, convergiendo al vector propio dominante si existe una brecha espectral.

El análisis se divide en dos escenarios principales:

1. Pesos Compartidos (Invariante en el Tiempo): Analizado utilizando modelos ALBERT donde los pesos son idénticos en todas las capas ($F_t = F$). Los autores proporcionan caracterizaciones analíticas rigurosas para casos donde el valor propio principal es real y estrictamente dominante, complejo o no dominante.
2. Pesos Variables por Capa (Variable en el Tiempo): Analizado utilizando BERT, GPT2 y GPTNeo. Aquí, $F_t$ cambia en cada capa. Los autores investigan si los tokens aún exhiben una tendencia a alinearse con el vector propio principal $v_{1,t}$ de la capa actual, a pesar del objetivo cambiante.

Para validar la teoría, los autores realizan experimentos numéricos en modelos preentrenados, midiendo la correlación entre los vectores de tokens y los vectores propios principales de las matrices de capa. También proponen un método para "dirigir" la salida modificando la capa final.

Contribuciones Clave

1. Analogía con el Método de la Potencia: El artículo establece que la operación de una capa de transformador (desatendiendo la FFN) actúa análogamente a un solo paso del Método de la Potencia. Los tokens son proyectados y normalizados, tendiendo a inclinarse hacia el vector propio principal de la matriz de capa $F_t$.
2. Caracterización Analítica de la Convergencia:
   • Para transformadores de pesos compartidos (por ejemplo, ALBERT), los autores demuestran que si la matriz de capa $F$ tiene un valor propio real estrictamente dominante, el sistema converge a un estado de consenso donde todos los tokens se alinean con el vector propio principal $v_1$.
   • Identifican condiciones de estabilidad, señalando que si el valor propio principal es complejo, los tokens convergen a un ciclo límite en el subespacio 2D generado por las partes real e imaginaria de los vectores propios.
   • Muestran que el "consenso bipartito" (tokens dividiéndose en dos grupos alineados con $\pm v_k$) es teóricamente posible pero empíricamente raro e inestable en configuraciones de múltiples cabezas.
3. Validación Empírica en Modelos Reales:
   • Los experimentos en ALBERT confirman que los tokens convergen al vector propio principal de la matriz de capa compartida. Sin embargo, la velocidad de convergencia a menudo está limitada por brechas espectrales pequeñas (donde $\lambda_1 \approx \lambda_2$) o valores propios complejos.
   • Los experimentos en modelos de pesos variables por capa (BERT, GPT2, GPTNeo) muestran que, aunque los tokens no acumulan alineación hacia un único vector propio global (debido a $F_t$ cambiante), aún exhiben un aumento estadísticamente significativo en la correlación con el vector propio principal $v_{1,t}$ de la capa actual en cada paso. Esto confirma que el comportamiento de "Método de la Potencia" persiste localmente incluso en sistemas variables en el tiempo.
4. Dirigir Tokens mediante Perturbación de Rango 1: El artículo propone un método constructivo para dirigir los tokens hacia una dirección deseada arbitraria. Al añadir una perturbación de rango 1 suficientemente grande a la matriz de salida de la capa final, los autores pueden crear una brecha espectral grande. Esto fuerza a la matriz de la capa final a tener un nuevo vector propio principal definido por el usuario, haciendo que todos los tokens se alineen con esta dirección específica.

Resultados

• ALBERT: En modelos con pesos compartidos, los tokens convergen rápidamente al consenso. Para ALBERT-base y large, el valor propio principal es real pero cercano al segundo, lo que lleva a una convergencia dentro del span de los dos primeros vectores propios. Para xlarge y xxlarge, los valores propios complejos hacen que los tokens roten en un subespacio 2D.
• Modelos de Pesos Variables por Capa: En variantes de BERT y GPT, la alineación con el vector propio principal de la capa actual ($\gamma_t$) es positiva pero generalmente débil. Esta debilidad se atribuye a las pequeñas brechas espectrales ($\lambda_{1,t} - \lambda_{2,t}$) en transformadores del mundo real. Sin embargo, se observa consistentemente que la probabilidad de que la alineación aumente en cada paso ($\rho_t > 0.5$), validando la analogía.
• Dirigir: Cuando la capa final se modifica con una perturbación de rango 1, la brecha espectral se vuelve grande y los tokens logran una alineación casi perfecta con el nuevo vector propio principal inducido, independientemente de la arquitectura del modelo.

Significado y Afirmaciones
El artículo afirma que observar los transformadores a través de la lente del Método de la Potencia proporciona una explicación fundamentada para el suavizado excesivo y el consenso de tokens. Sugiere que la "dirección" del consenso no es arbitraria, sino que está determinada por las propiedades espectrales del producto de las matrices de pesos de valor y salida.

Los autores enfatizan que, aunque la alineación natural en transformadores estándar es débil debido a brechas espectrales pequeñas, el mecanismo es lo suficientemente robusto para ser explotado. El método de dirección propuesto ofrece una alternativa fundamentada a las técnicas existentes (como la ingeniería de prompts o la edición de activaciones) para controlar la salida del modelo, ya que manipula directamente las propiedades de proyección geométrica de la capa final. El artículo concluye que esta analogía se mantiene incluso al desatender la FFN, sugiriendo que el mecanismo de atención y la normalización de capa son los impulsores principales de este comportamiento de alineación con vectores propios.