Representational Alignment Across Model Layers and Brain Regions with Multi-Level Optimal Transport

1. El Problema: Limitaciones de los Métodos de Alineación Actuales

La comparación de representaciones neuronales (entre modelos de IA, entre cerebros humanos o entre ambos) suele realizarse mediante emparejamiento de capas por pares (pairwise layer-wise matching). En este enfoque estándar, cada capa de una red fuente se empareja independientemente con la capa "mejor coincidente" de una red objetivo bajo una métrica de similitud (como RSA, CKA o predictividad lineal).

El paper identifica tres limitaciones fundamentales en este enfoque:

• Correspondencias rígidas y asimétricas: Se fuerza un mapeo uno-a-uno. Si las redes tienen profundidades diferentes ($L \neq M$) o si las características de una capa fuente se distribuyen a través de múltiples capas objetivo, este método falla. Además, los resultados son asimétricos dependiendo de la dirección de la comparación.
• Falta de una puntuación global: No existe una métrica unificada para la alineación de toda la red; solo se obtienen puntuaciones por capa.
• Ignorancia de la estructura global: Al optimizar cada emparejamiento de forma independiente, se ignora la estructura de activación global, lo que puede llevar a sobreajuste al ruido y a ignorar correspondencias jerárquicas significativas.

2. Metodología: Transporte Óptimo Multinivel (MOT)

Los autores proponen MOT (Multi-Level Optimal Transport), un marco unificado que infiere acoplamientos suaves (soft) y globalmente consistentes tanto a nivel de neuronas como de capas.

Estructura del Método

MOT opera en dos niveles jerárquicos:

1. Nivel Interno (Neurona a Neurona):

   • Para cada par de capas candidatas $(\ell, m)$, se calcula un costo de alineación neuronal utilizando Transporte Óptimo (OT) suave.
   • Se construye una matriz de disimilitud basada en las funciones de sintonización (tuning functions) de las neuronas.
   • Se resuelve un problema de OT para obtener un acoplamiento suave $Q_{\ell m}$, que especifica cómo cada neurona de la capa fuente se distribuye entre las neuronas de la capa objetivo. Esto permite que neuronas de una capa se mapeen a mezclas de neuronas en otra, manejando capas de anchos diferentes.

2. Nivel Externo (Capa a Capa):

   • Los costos internos $C_{\ell m}$ forman una matriz de costos entre capas.
   • Se resuelve un segundo problema de OT global para encontrar el acoplamiento de capas $P$.
   • Restricciones de Marginalidad: A diferencia de los métodos anteriores, MOT impone leyes de conservación de masa:
     • Cada capa fuente debe distribuir el 100% de su "masa" representacional a través de las capas objetivo.
     • La masa total recibida por cada capa objetivo debe estar equilibrada.
   • Esto permite que una capa fuente distribuya su representación en múltiples capas objetivo (muchos-a-muchos), resolviendo naturalmente los desajustes de profundidad.

Reconstrucción y Evaluación

La calidad de la alineación se mide reconstruyendo las activaciones de la capa fuente $\hat{X}_\ell$ utilizando las activaciones de la capa objetivo $Y_m$, los acoplamientos de neuronas $Q$ y los acoplamientos de capas $P$:
$$\hat{X}_\ell = \frac{L}{M} \sum_{m=1}^{M} P_{\ell m} Y_m Q_{\ell m}^\top$$
La puntuación global es la correlación media entre las activaciones originales y las reconstruidas en un conjunto de datos de validación.

Extensión Invariante a Rotaciones (MOT + R)

Dado que muchas representaciones en redes neuronales (especialmente en Transformers) son invariantes a rotaciones de la base, los autores proponen una extensión que incorpora matrices de rotación ortogonales $R_{\ell m}$. Se minimiza el error de reconstrucción optimizando alternativamente los acoplamientos de transporte y las rotaciones.

3. Contribuciones Clave

1. Marco de Alineación Global (Teórico): MOT optimiza conjuntamente todas las correspondencias de capas, produciendo asignaciones simétricas y globalmente consistentes con una sola puntuación de alineación. Evita que ciertas capas dominen el mapeo mientras otras son ignoradas.
2. Manejo Natural de Desajustes de Profundidad (Teórico): Al permitir mapeos suaves muchos-a-muchos, MOT alinea redes de diferentes profundidades sin forzar correspondencias uno-a-uno inapropiadas.
3. Extensión Invariante a Rotaciones (Teórico): La variante MOT+R permite recuperar correspondencias incluso cuando las características compartidas están en subespacios rotados, crucial para modelos como Vision Transformers.
4. Mejora en Puntuaciones de Alineación (Empírico): En todos los dominios probados, MOT iguala o supera a los métodos de emparejamiento por pares estándar.
5. Estructura Jerárquica Emergente (Empírico): Sin imponer restricciones de orden, MOT recupera automáticamente la organización jerárquica conocida (capas tempranas a tempranas, profundas a profundas) y revela cómo las redes más profundas distribuyen el procesamiento de una sola capa de redes más superficiales en múltiples etapas.

4. Resultados Experimentales

El método se evaluó en tres dominios diversos:

• Modelos de Lenguaje Grande (LLMs): Comparación entre familias (LLaMA-3.2, Qwen-2.5) y escalas.

  • MOT superó consistentemente a los baselines en la precisión de reconstrucción.
  • Los planes de transporte mostraron una estructura diagonal clara (alineación jerárquica), mientras que los métodos por pares produjeron mapas ruidosos.
  • Se observó que las capas de modelos más pequeños distribuyen su masa en múltiples capas consecutivas de modelos más grandes.

• Corteza Visual Humana (fMRI): Alineación de respuestas de 4 sujetos humanos (regiones V1-V4).

  • MOT recuperó correspondencias región-a-región consistentes entre sujetos, algo que los métodos por pares no lograron de manera estructurada.
  • Mantuvo un rendimiento de reconstrucción comparable a los baselines, pero con una interpretabilidad biológica superior.

• Modelos de Visión (Vision Transformers): Comparación entre DINOv2 y ViT-MAE.

  • El MOT estándar (sensible a rotaciones) no siempre superó a los baselines en precisión.
  • MOT + R (con rotación) logró puntuaciones de reconstrucción significativamente más altas y recuperó alineaciones jerárquicas limpias, demostrando la importancia de la invarianza rotacional en este dominio.

• Alineación de Trayectorias de Entrenamiento (Nivel Tres):

  • Se extendió MOT a un marco de tres niveles para alinear dos redes a lo largo de sus checkpoints de entrenamiento.
  • MOT reveló una correspondencia coherente y diagonal entre los checkpoints de dos entrenamientos (Pythia-14M), mientras que el emparejamiento por pares falló en capturar la trayectoria global.

5. Significado y Conclusión

El trabajo de Shah y Khosla representa un avance significativo en la neurociencia computacional y la IA. MOT proporciona una herramienta rigurosa para comparar sistemas con arquitecturas y profundidades arbitrarias, superando las limitaciones de los métodos "codiciosos" (greedy) actuales.

• Interpretabilidad: Permite entender cómo se distribuyen las representaciones a través de la profundidad de las redes, revelando principios organizativos que antes permanecían ocultos.
• Universalidad: Al demostrar que las correspondencias jerárquicas emergen naturalmente entre modelos artificiales y cerebros biológicos, MOT apoya la hipótesis de que existen dimensiones representacionales universales.
• Escalabilidad Futura: Aunque el costo computacional es alto ($O(L^2 n^3 \log n)$), el marco establece las bases para futuras mejoras algorítmicas y para la comparación de sistemas de inteligencia cada vez más complejos.

En resumen, MOT transforma la alineación representacional de un problema de emparejamiento local y rígido a uno de optimización global y flexible, ofreciendo una visión más rica y matizada de cómo se organizan la información en el cerebro y en las máquinas.