Resumen Técnico: Las Cabezas Laplacianas Mejoran los Transformers Suavizando las Representaciones de Tokens

Enunciado del Problema

La arquitectura estándar de Transformer actualiza las representaciones de tokens mediante atención multi-cabeza y conexiones residuales. En esta formulación, cada cabeza de atención $i$ calcula una matriz de atención softmax $P^{(i)}$ y actualiza la entrada $X$ mediante una suma ponderada de proyecciones de valor: $X \leftarrow X + \sum_i P^{(i)}XW_{Vi}W_{oi}$.

Aunque efectivo, este mecanismo carece de control explícito sobre la varianza intra-secuencia de las representaciones de tokens. La visión predominante en la literatura, particularmente respecto a las Redes Neuronales de Grafos (GNN) y los Transformers profundos, trata el "suvavizado excesivo" (oversmoothing)—donde las representaciones de tokens se vuelven demasiado similares, perdiendo distintividad—como un modo de fallo perjudicial. En consecuencia, la mayoría de los trabajos anteriores se centran en prevenir el suavizado mediante normalización, eliminación de bordes o funciones de pérdida de diversidad. Este artículo desafía la suposición de que el suavizado excesivo de tokens es inherentemente dañino, proponiendo en su lugar que el suavizado controlado puede ser beneficioso para el aprendizaje de representaciones.

Metodología

Los autores proponen una modificación arquitectónica sin parámetros: reemplazar un subconjunto de cabezas de atención estándar por cabezas laplacianas.

El Operador Laplaciano

En la actualización estándar, una cabeza calcula $P^{(i)}X$, que es un promedio ponderado por atención de los tokens. La cabeza laplaciana propuesta reemplaza la matriz de atención $P^{(i)}$ con el operador laplaciano de grafos $I - P^{(i)}$, donde $I$ es la matriz identidad. La ecuación de actualización se convierte en:
$$X \leftarrow X + \sum_{i \in A} P^{(i)}XW_{Vi}W_{oi} + \sum_{i \in L} (I - P^{(i)})XW_{Vi}W_{oi}$$
donde $A$ es el conjunto de cabezas de atención estándar y $L$ es el conjunto de cabezas laplacianas.

Motivaciones

El artículo proporciona dos motivaciones teóricas para esta modificación:

1. Control Directo de la Media y la Varianza: Las cabezas de atención estándar actualizan la media de las representaciones de tokens. En contraste, el término $(I - P^{(i)})X$ calcula la diferencia entre cada token y su media ponderada por atención. Añadir esta diferencia permite al modelo controlar directamente la varianza intra-secuencia de las representaciones de tokens, una capacidad ausente en los Transformers estándar.
2. Difusión en Grafos: Si los tokens se consideran como nodos en un grafo con pesos de borde definidos por $P^{(i)}$, entonces $I - P^{(i)}$ es el laplaciano de grafo normalizado por paseo aleatorio. Aplicar este operador corresponde a un solo paso de difusión de calor en el grafo de tokens, suavizando progresivamente la señal (representaciones de tokens) a través de la secuencia.

Contribuciones Clave

1. Propuesta Arquitectónica: Una modificación simple y sin parámetros a la arquitectura del Transformer que integra cabezas laplacianas junto con cabezas de atención estándar.
2. Validación Empírica: Demostración de que esta modificación mejora el rendimiento en tres paradigmas de aprendizaje distintos: clasificación supervisada, modelado de lenguaje y aprendizaje auto-supervisado.
3. Caracterización Geométrica: Un análisis detallado de cómo las cabezas laplacianas alteran la geometría de las representaciones de tokens aprendidas, mostrando que inducen formas beneficiosas de "colapso" y suavizado en lugar de un suavizado excesivo perjudicial.

Resultados Experimentales

1. Clasificación de Imágenes (Aprendizaje Supervisado)

Utilizando la familia DeiT-3 (ViT-B) en CIFAR-10, CIFAR-100 e ImageNet-1k:

• Rendimiento: Añadir cabezas laplacianas mejoró consistentemente la precisión Top-1. La mejora fue casi monótona con el número de cabezas laplacianas.
  • En ImageNet, el modelo con 12 cabezas laplacianas (ViT-B-12L) alcanzó una precisión del 82.32% en comparación con el 81.45% de la línea base.
  • En CIFAR-100, la ganancia fue de casi 2% (77.04% frente a 75.12%).
• Geometría:
  • ANOVA: Las cabezas laplacianas redujeron significativamente la Varianza Intra-Secuencia mientras aumentaban la Varianza Inter-Clase. Esto indica que los tokens dentro de la misma secuencia colapsan hacia una media común, mientras que las medias de diferentes secuencias se separan más efectivamente.
  • Colapso Neuronal: Los modelos mostraron una convergencia más fuerte hacia la geometría de Colapso Neuronal. Las representaciones de tokens formaron clusters bien separados, similares a un simplex (Simplex ETF) para cada clase, con las medias de las secuencias alineándose con la geometría global de la clase.

2. Modelado de Lenguaje (Autoregresivo)

Utilizando modelos estilo GPT-2 (561 millones de parámetros) pre-entrenados en FineWeb-Edu y ajustados finamente en benchmarks de razonamiento/código:

• Rendimiento: Los modelos con cabezas laplacianas superaron a la línea base en promedio en ARC-Easy, ARC-Challenge, MMLU, GSM8K y HumanEval.
  • El modelo con 5 cabezas laplacianas alcanzó el mejor rendimiento promedio (30.06%), con ganancias significativas en ARC-Easy (+16%) y ARC-Challenge (+7%).
• Geometría:
  • Decaimiento Espectral: Las cabezas laplacianas produjeron representaciones de tokens con espectros de decaimiento más rápidos. El número de componentes principales necesarios para explicar el 99% de la varianza disminuyó un 21% (de 547.7 a 430.4) cuando todas las cabezas eran laplacianas.
  • Colapso Lingüístico: Los modelos mostraron una mayor separabilidad entre las representaciones de tokens que comparten la misma predicción del siguiente token, sugiriendo una forma más fuerte de "colapso lingüístico" donde los tokens con roles semánticos similares se agrupan juntos.

3. Aprendizaje Auto-Supervisado (DINO)

Utilizando modelos ViT-S en ImageNet-1k para pre-entrenamiento DINO:

• Rendimiento: Las cabezas laplacianas mejoraron la precisión de sondeo lineal en ImageNet y el rendimiento de segmentación en ADE20K.
  • ViT-S-5L alcanzó el mejor mIoU en ADE20K (32.27%) en comparación con la línea base (29.81%).
• Geometría:
  • Segmentación: Las visualizaciones de PCA mostraron que las cabezas laplacianas produjeron límites de objetos más nítidos y preservaron detalles más finos (por ejemplo, formas de cascos, números de camisetas) en comparación con la línea base.
  • Decaimiento Espectral: Similar al modelado de lenguaje, estos modelos exhibieron un decaimiento espectral más rápido, concentrando las representaciones de tokens dentro de la imagen en un subespacio de menor dimensión.

Significado y Afirmaciones

El artículo argumenta que la visión predominante—de que el suavizado excesivo de tokens es universalmente dañino—es incompleta. En cambio, los autores afirman que el suavizado es un mecanismo para el aprendizaje de representaciones que puede ser beneficioso cuando está controlado.

• Beneficios Específicos de la Tarea: El "suavizado" inducido por las cabezas laplacianas no es un modo de fallo, sino una característica que se alinea con los objetivos de la tarea:
  • En clasificación supervisada, facilita el Colapso Neuronal al colapsar tokens dentro de las secuencias y alinear las medias de las secuencias con la geometría de la clase.
  • En modelado de lenguaje, mejora el Colapso Lingüístico, mejorando la separabilidad de tokens con objetivos de siguiente token idénticos.
  • En aprendizaje auto-supervisado, concentra las representaciones en subespacios de menor dimensión, ayudando a tareas posteriores como la segmentación.
• Mecanismo: Al incorporar explícitamente la difusión en grafos (a través del Laplaciano) en la arquitectura, el modelo gana la capacidad de controlar directamente la varianza intra-secuencia, lo que lleva a representaciones más robustas y geométricamente estructuradas.

Los autores concluyen que, aunque su implementación de estrategia fija es probablemente subóptima, los resultados sugieren que el suavizado debe investigarse más a fondo como un componente fundamental del aprendizaje de representaciones, en lugar de únicamente como un problema a mitigar.