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

Título: Alinhamento Representacional entre Camadas de Modelos e Regiões Cerebrais com Transporte Ótimo Multi-Nível (MOT)

1. O Problema

A comparação de representações neurais (tanto em redes neurais artificiais quanto no cérebro biológico) enfrenta desafios fundamentais quando se utiliza métodos padrão de similaridade representacional (como RSA, CKA ou correspondência de camadas greedy). Os principais problemas identificados são:

• Correspondências Rígidas e Assimétricas: Métodos tradicionais emparelham cada camada de origem com a "melhor" camada de destino individualmente. Isso força uma correspondência um-para-um, o que falha quando redes têm profundidades diferentes (número de camadas distinto) ou quando as características de uma camada estão distribuídas entre várias camadas de outro modelo.
• Falta de Consistência Global: Ao otimizar cada par de camadas independentemente, ignora-se a estrutura global de ativação da rede. Isso pode levar a sobreajuste ao ruído em camadas específicas e a mapeamentos assimétricos (o resultado depende de qual rede é considerada "origem").
• Incapacidade de Lidar com Desalinhamento de Profundidade: Não há uma pontuação unificada para o alinhamento global da rede, e métodos existentes não conseguem lidar naturalmente com redes de profundidades diferentes sem descartar camadas ou forçar correspondências inadequadas.

2. Metodologia: Multi-Level Optimal Transport (MOT)

Os autores propõem o MOT, um framework unificado baseado em Transporte Ótimo (Optimal Transport - OT) que realiza o alinhamento em dois níveis simultaneamente, garantindo consistência global.

• Nível Interno (Neurônio a Neurônio):

  • Para cada par de camadas candidatas, calcula-se um custo de alinhamento entre os neurônios (unidades) usando uma formulação de soft-matching OT.
  • Isso permite que um neurônio de origem distribua sua "massa" representacional para múltiplos neurônios de destino, lidando com diferenças no número de unidades (largura das camadas).
  • O resultado é uma matriz de acoplamento suave ($Q_{\ell m}$) que define como os neurônios se correspondem.

• Nível Externo (Camada a Camada):

  • Os custos internos formam uma matriz de custo entre camadas. O MOT resolve um segundo problema de Transporte Ótimo para encontrar o acoplamento global entre as camadas das duas redes.
  • Restrições de Conservação de Massa: O framework impõe que cada camada de origem distribua 100% de sua massa representacional entre as camadas de destino, e que a massa total recebida por cada camada de destino seja balanceada.
  • Resultado: Em vez de um emparelhamento rígido, obtém-se um plano de transporte suave onde uma camada de origem pode se alinhar a uma combinação de várias camadas de destino. Isso resolve naturalmente discrepâncias de profundidade.

• Extensão Invariante a Rotação (MOT + R):

  • Reconhecendo que representações em redes (especialmente Transformers) podem ser invariantes a rotações no espaço de características, os autores introduzem matrizes de rotação ortogonais ($R_{\ell m}$) no processo de otimização.
  • Isso minimiza o erro de reconstrução permitindo que as bases de coordenadas sejam alinhadas, capturando equivalências geométricas que métodos sensíveis à rotação perderiam.

• Avaliação:

  • A qualidade do alinhamento é medida pela capacidade de reconstruir as ativações originais de uma rede usando as ativações da outra rede e os planos de transporte aprendidos. A métrica é a correlação média entre as ativações reais e as reconstruídas em dados de validação.

3. Contribuições Principais

1. Framework de Alinhamento Global (Teórico): O MOT otimiza conjuntamente todas as correspondências de camadas, produzindo atribuições simétricas e globalmente consistentes com uma única pontuação de alinhamento de rede.
2. Tratamento Natural de Desalinhamento de Profundidade (Teórico): Ao permitir mapeamentos "muitos-para-muitos" suaves, o MOT alinha redes de profundidades diferentes sem forçar correspondências um-para-um inadequadas.
3. Extensão Invariante a Rotação (Teórico): A incorporação de transformações ortogonais permite recuperar correspondências mesmo quando as características compartilhadas estão em subespaços rotacionados.
4. Descoberta de Estrutura Hierárquica Emergente (Empírico): Sem impor restrições de ordem, o MOT recupera naturalmente a organização hierárquica (camadas iniciais mapeiam para iniciais, profundas para profundas) em dados de cérebro e modelos, algo que métodos greedy falham em revelar.
5. Distribuição de Características (Empírico): O MOT revela como redes mais profundas espalham computações que redes mais rasas comprimem em menos estágios (uma camada rasa distribui massa para várias camadas vizinhas em uma rede profunda).

4. Resultados Experimentais

O MOT foi avaliado em três domínios distintos:

• Modelos de Linguagem (LLMs):

  • Comparação entre famílias (LLaMA vs. Qwen) e escalas (0.5B a 3B).
  • Resultado: O MOT superou consistentemente os métodos baseline (como correspondência par-a-par greedy) na precisão de reconstrução. Os planos de transporte mostraram uma estrutura diagonal clara, indicando correspondências hierárquicas robustas.

• Córtex Visual Humano (fMRI):

  • Alinhamento entre diferentes sujeitos humanos usando o dataset Natural Scenes Dataset (V1-V4).
  • Resultado: O MOT recuperou correspondências regionais esperadas (ex: V1 de um sujeito alinhado com V1 de outro) com alta consistência, enquanto o OT par-a-par produziu mapeamentos ruidosos e desestruturados.

• Modelos de Visão (Vision Transformers):

  • Comparação entre DINOv2 e ViT-MAE de várias escalas.
  • Resultado: O MOT padrão (sensível à rotação) teve desempenho misto. No entanto, a variante MOT + R (invariante a rotação) superou significativamente todos os baselines, alcançando pontuações de reconstrução muito mais altas e revelando alinhamentos hierárquicos limpos. Isso destaca a importância da invariância a rotação em Transformers.

• Alinhamento de Trajetórias de Treinamento (Nível Triplo):

  • Uma extensão do MOT para alinhar checkpoints de treinamento de dois modelos ao longo do tempo.
  • Resultado: O MOT revelou correspondências coerentes entre os checkpoints de duas trajetórias de treinamento independentes, algo que o emparelhamento par-a-par não conseguiu fazer.

5. Significado e Conclusão

O trabalho demonstra que o alinhamento representacional não deve ser tratado como uma série de problemas locais independentes, mas como um problema global de otimização.

• Interpretabilidade: O MOT fornece planos de transporte que são não apenas métricas de similaridade, mas mapas interpretáveis de como a informação flui e se transforma entre sistemas com arquiteturas diferentes.
• Universalidade: A capacidade de lidar com profundidades variáveis e invariâncias geométricas torna o MOT uma ferramenta poderosa para testar hipóteses sobre a universalidade das representações neurais (biológicas e artificiais).
• Limitações: O método é computacionalmente intensivo ($O(L^2 n^3 \log n)$), o que pode limitar sua aplicação em modelos extremamente grandes sem otimizações algorítmicas futuras.

Em suma, o MOT oferece uma nova base rigorosa para comparar sistemas inteligentes, revelando padrões estruturais ocultos que métodos tradicionais de correspondência de camadas ignoram.