MIPIC: Matryoshka Representation Learning via Self-Distilled Intra-Relational and Progressive Information Chaining

1. Declaração do Problema

Matryoshka Representation Learning (MRL) visa criar embeddings aninhados onde um único vetor de alta dimensão contém prefixos de baixa dimensão significativos. Isso permite inferência adaptativa, onde um modelo pode truncar embeddings para se ajustar a diferentes orçamentos computacionais sem retreinamento.

No entanto, os métodos existentes de MRL enfrentam dois desafios críticos:

1. Incoerência Estrutural: Abordagens atuais frequentemente dependem de alinhamento ao nível de sentença ou supervisão independente em diferentes níveis de truncamento. Elas falham em organizar explicitamente como a informação semântica está arranjada internamente através das dimensões de embedding e das camadas da rede. Consequentemente, prefixos de baixa dimensão frequentemente perdem relações estruturais críticas encontradas no espaço de dimensão completa.
2. Gargalo de Compressão Semântica: Comprimir conhecimento rico e específico de tarefas das camadas profundas em prefixos estreitos é difícil. Métodos padrão frequentemente aplicam supervisão apenas na camada final, forçando o modelo a realizar um "salto semântico" que pode levar a instabilidade e perda de informação em representações iniciais de baixa capacidade.

2. Metodologia: Framework MIPIC

Os autores propõem MIPIC, um framework de treinamento unificado que sincroniza a aprendizagem de representações aninhadas com o fluxo de informação interno do modelo. Ele emprega dois mecanismos sinérgicos: Alinhamento Intra-Relacional Auto-Distilado (SIA) e Encadeamento Progressivo de Informação (PIC).

A. Alinhamento Intra-Relacional Auto-Distilado (SIA)

O SIA aborda a consistência estrutural cruzada de dimensões. Em vez de apenas alinhar vetores de sentença, ele reorganiza a hierarquia de características interna para priorizar características estruturais centrais dentro de subespaços de baixa dimensão.

• Correspondência de Distribuição de Atenção: A representação de dimensão completa atua como um "professor". Ela alinha a ordem de importância relativa dos tokens (derivada de pontuações de atenção) entre a dimensão completa e os prefixos truncados. Isso garante que os tokens mais informativos permaneçam enfatizados mesmo após a compressão.
• Alinhamento de Estados Ocultos Top-k via CKA: Para evitar saturação de informação em gargalos pequenos, o SIA usa uma estratégia de Seleção Rígida Top-k. Ele seleciona os $k$ tokens mais importantes (baseados em pontuações de atenção) e alinha a estrutura geométrica desses subconjuntos entre as representações do aluno (truncado) e do professor (completo) usando Alinhamento de Kernel Centralizado (CKA).
  • O CKA é escolhido por sua invariância a transformações ortogonais e escalonamento isotrópico, tornando-o robusto para comparar representações de diferentes dimensões.
  • Isso cria uma "estrutura dupla Matryoshka": núcleos de baixa dimensão capturam um esqueleto semântico grosseiro, enquanto dimensões mais altas adicionam detalhes mais finos.

B. Encadeamento Progressivo de Informação (PIC)

O PIC aborda a consolidação semântica ao longo da profundidade. Ele guia o fluxo de informação específica de tarefas das camadas mais profundas (mais expressivas) para as camadas anteriores.

• Alinhamento com Suporte (Scaffolded): Diferentemente de métodos que supervisionam apenas a saída final, o PIC constrói um pipeline onde cada camada atua como uma ponte semântica. Ele propaga dicas relevantes para a tarefa das camadas superiores para as camadas inferiores.
• Objetivo InfoNCE: O método maximiza a informação mútua entre checkpoints adjacentes (camadas) na rede. Ele usa um projetor leve para mapear representações a montante de baixa dimensão para o espaço de destino a jusante, minimizando uma perda InfoNCE.
• Supervisão Seletiva: A supervisão é restrita ao subespaço compacto reorganizado pelo SIA. Isso permite que as dimensões restantes continuem modelando características linguísticas de baixo nível (sintaxe, semântica geral) sem super-regularização, garantindo uma base flexível para refinamento específico de tarefas.

C. Objetivo Geral de Treinamento

A função de perda total combina a perda de tarefa padrão de MRL com as perdas de alinhamento propostas:
$$L_{MIPIC} = \alpha L_{MRL} + (1 - \alpha)[L_{SIA} + L_{PIC}]$$
Onde $\alpha$ equilibra o desempenho da tarefa e o alinhamento estrutural/semântico.

3. Contribuições Principais

1. Framework Unificado: O MIPIC é o primeiro framework a coordenar explicitamente tanto o alinhamento cruzado de dimensões (via SIA) quanto a consolidação de informação ao longo da profundidade (via PIC) para resolver incoerência estrutural e semântica no MRL.
2. Mecanismos Novos:
   • SIA: Introduz auto-distilação CKA Top-k para impor consistência topológica, forçando estruturas semânticas críticas em prefixos de baixa dimensão.
   • PIC: Introduz uma estratégia de aprendizado com suporte que permite que camadas iniciais adquiram informação relevante para a tarefa progressivamente, atuando como uma "ponte semântica".
3. Validação Empírica Extensa: O método foi testado em backbones diversos (TinyBERT, BERT-base, BGEM3, Qwen3-0.6B) e benchmarks (STS, NLI, Classificação).

4. Resultados Experimentais

• Desempenho: O MIPIC alcança desempenho competitivo ou superior comparado a baselines de última geração (MRL, Sentence Embeddings Espresso, SimCSE) em todas as capacidades.
• Vantagem de Baixa Dimensão: Os ganhos mais significativos são observados em regimes extremos de baixa dimensão (16 e 32 dimensões). Por exemplo, no conjunto de dados Banking77 com TinyBERT, o MIPIC alcançou 48,93% de precisão em 16 dimensões, superando significativamente o MRL (40,64%) e o ESE (43,30%).
• Generalização: O método demonstra robustez em configurações dentro do domínio e fora do domínio (OOD), mantendo alta densidade semântica mesmo sob truncamento severo.
• Estudos de Ablação:
  • Remover o SIA causa quedas acentuadas de desempenho em tamanhos pequenos, confirmando que a ordem geométrica é fundamental para a densidade de informação.
  • Remover o PIC leva à perda semântica nas etapas iniciais, destacando a necessidade de orientação progressiva.
  • Escalonamento progressivo de dimensões (usando diferentes dimensões em diferentes camadas) supera manter todas as camadas em tamanho completo.

5. Significado e Trade-offs

• Eficiência de Inferência: O MIPIC introduz zero latência adicional de inferência ou parâmetros. Todos os projetores auxiliares e módulos de alinhamento são descartados após o treinamento. O modelo se comporta exatamente como um backbone padrão durante a inferência, mas com representações truncadas superiores.
• Custo de Treinamento: O método incorre em uma sobrecarga de treinamento maior (aproximadamente 45% de throughput mais lento que o ESE) devido ao cálculo de perdas auxiliares CKA e InfoNCE através de múltiplas camadas. No entanto, os autores argumentam que isso é um investimento único justificável, dados os ganhos massivos na qualidade da representação e eficiência de compressão, especialmente para modelos implantados em ambientes com recursos limitados.
• Impacto: O MIPIC fornece uma abordagem principista para organizar informação semântica, tornando embeddings Matryoshka viáveis para cenários de compressão extrema onde métodos anteriores falharam em manter a expressividade semântica.