Resumo Técnico: Composição de Especialistas de Memória para Modelos de Mundo de Difusão

Declaração do Problema

Os modelos de mundo (MWs) são essenciais para o aprendizado por reforço (RL) e o planejamento, pois preveem estados futuros plausíveis consistentes com observações passadas. No entanto, arquiteturas existentes enfrentam um trade-off fundamental entre fidelidade da memória e eficiência computacional. Modelos baseados em transformadores preservam detalhes locais, mas sofrem com custos de atenção quadráticos, limitando o comprimento do contexto. Por outro lado, redes recorrentes e modelos de espaço de estados (SSMs) escalam de forma eficiente, mas comprimem o histórico em estados ocultos, descartando inevitavelmente detalhes de alta granularidade ao longo do tempo. A simples aumento do comprimento do contexto em uma única arquitetura leva a um treinamento instável e a custos de inferência proibitivos. Além disso, a composição ingênua de múltiplos modelos frequentemente resulta em "colapso de modo", onde a distribuição combinada afina excessivamente regiões comuns e perde diversidade.

Metodologia

Os autores propõem a Composição de Especialistas de Memória (CoME), um framework baseado em difusão que desacopla a consistência futuro-passado de uma única arquitetura monolítica. Em vez disso, a CoME aproveita um conjunto de especialistas especializados, cada um lidando com um papel específico de memória, e os funde probabilisticamente no momento da inferência.

1. Produto de Especialistas Contrastivos (PoCE)

Para abordar as limitações do Produto Padrão de Especialistas (PoE), que pode suprimir a diversidade e amplificar modos espúrios, os autores introduzem uma formulação de Produto de Especialistas Contrastivos.

• Mecanismo: Para um conjunto de especialistas heterogêneos $p_i(x)$, a composição é definida como $p(x) \propto \prod \tilde{p}_i(x)$, onde $\tilde{p}_i(x) \propto p_i(x)^{\alpha_i} \bar{p}_i(x)^{1-\alpha_i}$.
• Linha de Base Contrastiva: $\bar{p}_i(x)$ representa a distribuição de linha de base incondicional do especialista. O coeficiente $\alpha_i > 1$ atua como um fator de contraste.
• Justificativa Teórica: Diferentemente do temperamento ingênuo (que reduz a variância do kernel e distorce a geometria local), esta mistura contrastiva repondera os componentes da mistura enquanto preserva a forma dos kernels. Ela suprime efetivamente modos espúrios (onde um especialista atribui alta probabilidade sem evidência forte) enquanto mantém modos secundários válidos, desde que o especialista condicional atribua maior probabilidade a modos válidos do que a linha de base.

2. Especialistas de Memória Especializados

O framework instancia três especialistas complementares:

• Memória de Curto Prazo (MCP): Um modelo de difusão condicionado a uma janela de contexto recente ($10\text{--}100$ quadros). Ele captura dinâmicas locais finas e dependências temporais imediatas usando uma arquitetura com comprimento de contexto estendido (por exemplo, atenção de janela deslizante).
• Memória de Longo Prazo (MLP): Um modelo de difusão separado que armazena o histórico episódico ($100\text{--}1000$ quadros) diretamente em seus pesos. Isso é alcançado via ajuste fino no momento do teste de adaptadores leves LoRA no contexto de longo prazo. Essa abordagem amortiza o custo da recuperação, permitindo a reutilização em tempo constante de experiências passadas sem alimentar todo o histórico através do modelo a cada passo. Evita o esquecimento catastrófico congelando o backbone e atualizando apenas os adaptadores.
• Memória de Longo Prazo Espacial (MLPE): Uma extensão da MLP que incorpora priores espaciais (por exemplo, poses de câmera, mapas de pontos) extraídos do histórico. Este especialista impõe coerência geométrica e espacial, desambiguando locais em ambientes repetitivos ou visualmente semelhantes.

3. Geração Unificada

A distribuição final é um produto do prior pré-treinado e dos três especialistas contrastivos (MCP, MLP, MLPE). Durante a amostragem, as funções de pontuação desses especialistas são combinadas aditivamente (no espaço log) para guiar o processo de remoção de ruído, garantindo que o futuro gerado respeite as dinâmicas recentes, o histórico de longo prazo e a consistência espacial.

Contribuições Principais

1. Arquitetura de Memória Desacoplada: O artigo propõe resolver o trade-off de memória distribuindo-a entre especialistas especializados em vez de um único backbone, formulado como um problema de Produto de Especialistas.
2. Produto de Especialistas Contrastivos (PoCE): Uma estratégia inovadora para fundir especialistas heterogêneos. Os autores fornecem evidências teóricas e empíricas de que este método previne o colapso de modo e a superconfiança ao fatorar modos redundantes, ao contrário do PoE padrão.
3. Memória Externa de Difusão: A introdução de um módulo MLP que usa ajuste fino no momento do teste de adaptadores LoRA para armazenar conhecimento episódico em pesos. Isso é estendido para MLPE para incorporar estrutura espacial, melhorando a precisão em tarefas de navegação.

Resultados Experimentais

O método foi avaliado em benchmarks sintéticos (Memory Maze, DMLab, Minecraft) e do mundo real (RealEstate10K, RECON).

• Qualidade de Reconstrução: No conjunto de dados Memory Maze, a CoME alcançou desempenho de última geração em LPIPS (0,097), SSIM (0,892) e PSNR (23,07), superando significativamente as linhas de base, incluindo transformadores de atenção total, SSMs e modelos de janela deslizante.
• Escalabilidade: Diferentemente de modelos de atenção total que exigem $60\times$ mais computação do que a linha de base MCP, a CoME alcança consistência superior com uma fração do custo de treinamento e inferência.
• Comprimento de Contexto: O modelo não mostrou saturação na consistência de desempenho até 480 quadros de contexto, demonstrando a eficácia do componente MLP.
• Planejamento e Navegação: Em tarefas de navegação condicionadas a objetivos (benchmark RECON), a CoME reduziu o Erro de Trajetória Absoluto (ATE) para 0,96, superando modelos de mundo de navegação especializados (NWM) e outras linhas de base.
• Capacidade de Recuperação: Em um experimento de trajetória reversa no RealEstate10K, a CoME recuperou com sucesso quadros iniciais após rolagens para frente, enquanto as linhas de base sem memória geraram cenas plausíveis, mas inconsistentes.
• Estudos de Ablação: Remover a formulação contrastiva causou o colapso do desempenho, confirmando a necessidade do mecanismo PoCE. Aumentar o rank do LoRA e o comprimento do contexto para MLP melhorou consistentemente os resultados sem sobreajuste.

Significado e Alegações

O artigo afirma estabelecer um paradigma inovador para construir modelos de mundo de difusão aumentados por memória. Ao distribuir a memória entre especialistas especializados e usar uma estratégia de composição contrastiva, a CoME equilibra eficiência com consistência temporal de alta fidelidade. Os autores argumentam que essa abordagem permite que os modelos de mundo escalem para contextos longos sem incorrer em custos quadráticos ou sofrer com a perda de detalhes inerente à compressão recorrente. O método é apresentado como um framework geral aplicável a vários modelos de mundo baseados em difusão, aumentando sua confiabilidade para planejamento e tomada de decisão em configurações de aprendizado por reforço. O trabalho enfatiza que a memória deve ser uma propriedade distribuída do sistema, e não uma restrição de uma única arquitetura.