Can Cross-Layer Transcoders Replace Vision Transformer Activations? An Interpretable Perspective on Vision

Resumen Técnico: Cross-Layer Transcoders (CLTs) para Vision Transformers

1. El Problema

Los Modelos Fundacionales (FMs), especialmente los Transformers de Visión (ViTs), han demostrado un rendimiento excepcional en diversas tareas. Sin embargo, su funcionamiento interno sigue siendo una "caja negra", lo que limita su interpretabilidad, controlabilidad y fiabilidad.

• Limitaciones de los enfoques actuales: Métodos existentes como los Autoencoders Dispersos (SAEs) se centran en extraer características interpretables, pero operan de manera local (dentro de una sola capa). Esto les impide capturar la estructura computacional cruzada entre capas (cross-layer) y la contribución relativa de cada capa en la formación de la representación final.
• La brecha: No existe un método que pueda descomponer la representación final de un ViT en una construcción aditiva y resuelta por capas, permitiendo una atribución fiel de cómo la información se transforma a través de la profundidad de la red.

2. Metodología: Cross-Layer Transcoders (CLTs)

Los autores proponen la adopción de Transcoders de Capa Cruzada (CLTs) como modelos proxy fiables, dispersos y conscientes de la profundidad para los bloques MLP (Perceptrón Multicapa) en ViTs.

• Arquitectura:

  • Los CLTs utilizan un esquema de codificador-decodificador.
  • Codificación: Transforman las activaciones pre-MLP de capas anteriores en codificaciones dispersas (sparse embeddings) $z_i$. Se utilizan funciones de no linealidad como JumpReLU, ReLU-Top-k y Abs-Top-k.
  • Decodificación Cruzada: Reconstruyen la salida post-MLP de una capa objetivo $j$ ($\hat{y}_j$) como una suma lineal de las contribuciones de todas las capas fuente anteriores ($i \leq j$):
    $$\hat{y}_j = \sum_{i=1}^{j} z_i W_{dec}^{i \to j}$$
  • Esta estructura triangular impone una causalidad estricta, evitando fugas de información de capas futuras.

• Entrenamiento y Reemplazo:

  • Se entrenan CLTs sobre ViTs congelados (CLIP ViT-B/32 y ViT-B/16) para minimizar el error de reconstrucción entre la salida real del MLP y la salida reconstruida por el CLT.
  • Modelo de Reemplazo: Una vez entrenados, los CLTs pueden reemplazar funcionalmente los bloques MLP originales durante la inferencia, permitiendo analizar el modelo original sin comprometer su rendimiento en tareas de clasificación zero-shot.

3. Contribuciones Clave

1. CLTs para Visión: Se investiga y valida la efectividad de los CLTs en el dominio de la visión, logrando una reconstrucción dispersa y cruzada de activaciones post-MLP.
2. Reemplazo Funcional Fiel: Los CLTs pueden reemplazar los bloques MLP, especialmente en capas tardías o para el token [CLS], preservando e incluso mejorando en algunos casos la precisión de clasificación zero-shot. Esto los convierte en un proxy fiable para estudiar la contribución de capas anteriores.
3. Atribución Cruzada Interpretativa: Se introduce un sistema de puntuación de contribución basado en proyección que revela una asimetría estructural:
   • Los tokens de parche (patches) muestran una atribución diagonal dominante (cada capa explica principalmente su propia salida).
   • El token [CLS] muestra una atribución ampliamente distribuida a través de la profundidad, agregando información semántica de muchas capas anteriores.
4. Validación de Atribución Fiel: Experimentos de ablación demuestran que la representación final está concentrada en un subconjunto pequeño de capas fuente. Mantener solo las 4 capas superiores recupera la precisión original, mientras que eliminar la capa de mayor puntuación degrada significativamente el rendimiento.

4. Resultados Experimentales

Los experimentos se realizaron en los conjuntos de datos CIFAR-100, COCO e ImageNet-100 con backbones ViT-B/32 y ViT-B/16.

• Fidelidad de Reconstrucción:

  • Los CLTs logran una alta fidelidad (similitud coseno de 0.92–0.97 y $R^2$ de 0.89–0.95) en la reconstrucción de activaciones post-MLP.
  • La reconstrucción es más fácil en ViT-B/16 (parches más pequeños) que en ViT-B/32.
  • La variante ReLU-Top-k mostró el mejor rendimiento general, especialmente en las capas medias y tardías.

• Rendimiento en Clasificación (Reemplazo Cascada):

  • Token [CLS]: El reemplazo de todos los bloques MLP por CLTs (incluso en todas las capas) mantiene la precisión casi idéntica a la del modelo base (ej. 61.74% vs 61.65% en CIFAR-100).
  • Tokens de Parche: El reemplazo completo de todas las capas y todos los tokens causa una degradación notable, especialmente en las capas iniciales. Sin embargo, el reemplazo de las últimas capas (ej. 10→11) es seguro y preserva el rendimiento.
  • Conclusión: Los CLTs son robustos para el token global ([CLS]) y capas tardías, pero las computaciones de parches en capas tempranas son más sensibles a errores de aproximación.

• Análisis de Atribución:

  • Los mapas de calor de contribución confirman que el token [CLS] integra señales semánticas de toda la red, mientras que los parches se refieren principalmente a transformaciones dentro de la misma capa.
  • La eliminación de la capa de mayor puntuación en la reconstrucción final reduce drásticamente la precisión, validando que las puntuaciones de atribución identifican correctamente las capas críticas.

5. Significado e Impacto

Este trabajo establece un nuevo paradigma para la interpretabilidad en modelos de visión:

• Más allá de la capa única: A diferencia de los SAEs tradicionales, los CLTs capturan la dependencia cruzada entre capas, ofreciendo una visión completa de cómo se construyen las representaciones a lo largo de la profundidad de la red.
• Proxy Confiable: Demuestra que es posible reemplazar componentes internos complejos (MLPs) con modelos interpretativos (CLTs) sin sacrificar el rendimiento funcional, permitiendo un análisis de circuitos y atribución en tiempo real.
• Insights Arquitectónicos: Revela una diferencia fundamental en la dinámica de información entre tokens globales ([CLS]) y locales (parches), sugiriendo que la agregación semántica global es un proceso distribuido a través de la profundidad, mientras que el procesamiento local es más jerárquico y por capas.

En resumen, los Cross-Layer Transcoders se presentan como una herramienta fundamental para desbloquear la "caja negra" de los Vision Transformers, proporcionando tanto una sustitución funcional eficiente como una explicación causal detallada de su funcionamiento interno.