Resumo Técnico: Memória Tensorial: Estado Recorrente de Tamanho Fixo para Transformers de Longo Alcance

1. Declaração do Problema

Transformers processam imagens e vídeos ao achatar as dimensões espaciais e temporais em longas sequências de tokens. Embora os mecanismos de atenção e o cache KV preservem características passadas, eles sofrem com duas limitações primárias à medida que o horizonte da sequência cresce:

1. Escalabilidade de Memória: A inferência prática depende do armazenamento de chaves e valores passados, o que escala linearmente com o número de tokens. Isso torna-se computacionalmente proibitivo e consome muita memória para vídeos longos e documentos.
2. Falta de Estado Persistente: Transformers padrão carecem de um estado explícito e persistente. Eles atendem a características passadas, mas não mantêm uma representação estruturada que possa ser atualizada e consultada conforme a entrada muda. Isso leva a fragilidade sob oclusão, observações ausentes e à necessidade de reconstruir repetidamente o contexto global a partir de um histórico em constante crescimento, em vez de manter um estado compacto e persistente.

Abordagens existentes, como atenção esparsa, compressão de histórico ou contextos aumentados por recuperação, frequentemente estendem o comprimento do contexto, mas armazenam informações como coleções não estruturadas de vetores de tokens ou slots. Elas não preservam explicitamente a estrutura espacial nem mantêm um "estado do mundo" compacto adequado para tarefas baseadas em estado, como compreensão de vídeo.

2. Metodologia: Memória Tensorial

Os autores propõem a Memória Tensorial, um módulo leve que aumenta os blocos padrão de Transformer com um tensor de memória recorrente 3D de tamanho fixo. Este módulo desacopla a capacidade do estado do comprimento da entrada, preservando ao mesmo tempo um viés indutivo espacial.

Arquitetura Central

A memória é definida como um par de volumes (estado oculto $h_t$ e estado da célula $c_t$) com dimensões $C \times D \times H \times W$, onde as dimensões espaciais são constantes, independentemente do comprimento da entrada. O módulo opera em um ciclo por fatia de tokens:

1. Fatiamento (Chunking): Os tokens de entrada são agrupados em fatias. Uma "consulta de leitura" é derivada do primeiro token da fatia, e um "resumo de escrita" é uma projeção aprendida de toda a fatia.
2. Endereçamento Contínuo: Um preditor de coordenadas compartilhado (usando pesos acoplados) gera coordenadas 3D contínuas para leitura ($\mu_{read}$) e escrita ($\mu_{write}$) dentro da grade de voxels.
3. Leitura (Amostragem Contínua): O modelo lê do estado oculto $h_t$ em $\mu_{read}$ usando interpolação trilinear. A característica recuperada é projetada de volta para a dimensão do token e fundida no fluxo de tokens por meio de uma conexão residual com portão (gated). Um portão escalar aprendido $\gamma$ controla a força dessa fusão, permitindo que o modelo suprima o caminho da memória se não for benéfico para uma tarefa específica.
4. Escrita (Escrita Suave Diferenciável): O modelo escreve informações prevendo um vetor de conteúdo e um parâmetro de dispersão $\sigma$. Ele deposita esse conteúdo como um volume ponderado por Gaussiana ao redor da localização 3D prevista $\mu_{write}$. Esta "escrita suave" evita a atribuição rígida de voxels, permitindo que os gradientes moldem onde a informação é colocada.
5. Atualização Recorrente: O volume de escrita é concatenado com o estado oculto atual. Um operador 3D fatorizado e leve (convoluções 3D profundas seguidas por uma mistura pontual) calcula portões no estilo ConvLSTM ($i_t, f_t, o_t, g_t$). Esses portões atualizam os estados da célula e oculto, permitindo que a informação persista, propague-se localmente e permaneça disponível mesmo quando as observações estão ausentes.

Escolhas de Design Chave

• Estado de Tamanho Fixo: O tamanho do tensor de memória é constante, prevenindo o crescimento linear de memória associado ao cache KV.
• Estrutura Espacial: A grade de voxels 3D preserva explicitamente as relações espaciais, ao contrário de slots de tokens não estruturados.
• Escritas Suaves Diferenciáveis: O uso de volumes ponderados por Gaussiana permite atualizações suaves e contínuas, em vez de atribuições discretas e rígidas.
• Fusão com Portão: O portão aprendido $\gamma$ garante que o módulo possa ser anexado a modelos pré-treinados sem desestabilizar o treinamento; se a memória for desnecessária, o portão suprime o caminho.

3. Contribuições Principais

• Mecanismo de Memória Novel: Introdução de um estado de voxel 3D recorrente de tamanho fixo, para o qual os tokens podem escrever e do qual podem ler por meio de coordenadas contínuas, atuando como um estado espacial explícito.
• Integração: O módulo integra-se limpa aos pipelines de treinamento padrão de Transformer e pode ser anexado ou removido de blocos existentes sem alterações arquitetônicas.
• Suite de Diagnóstico: Desenvolvimento de uma suite de diagnóstico controlada "de brinquedo" projetada para isolar cenários onde o estado persistente é benéfico (por exemplo, ligação de coordenadas, oclusão, construção de mapas) versus cenários onde não é.

4. Resultados Experimentais

Os autores avaliaram a Memória Tensorial em benchmarks padrão nas áreas de linguagem, imagem e vídeo, bem como na suite de diagnóstico.

• Modelagem de Linguagem (WikiText-2): A Memória Tensorial alcançou os maiores ganhos absolutos, reduzindo a perplexidade de teste de 130,30 (Transformer Base) para 100,07. Ela superou baselines incluindo Transformer-XL, atenção local e atenção linear.
• Reconstrução de Patch de Imagem (CUB-200-2011): Em uma tarefa de reconstrução de patch mascarado, a variante com Memória Tensorial (Spatial ViT) mostrou melhorias consistentes em PSNR e SSIM em todas as taxas de máscara (25%, 50%, 75%) em comparação com Base ViT, Local ViT e Registers ViT.
• Reconhecimento de Ação em Vídeo (UCF-101): O modelo superou tanto uma baseline de Transformer espaço-temporal conjunta quanto uma forte baseline de "token registrador" em precisão Top-1 (89,21% vs 82,12%) e Top-5. Os autores notam que isso veio ao custo de um tempo de treinamento significativamente maior devido à recorrência de voxel por passo.
• Diagnósticos de Brinquedo:
  • Ligação de Coordenadas: A Memória Tensorial manteve alta precisão (75–91%) à medida que o número de escritas aumentava, enquanto as baselines baseadas em atenção colapsaram para desempenho próximo ao aleatório (7–14%).
  • Controle Sem Dano: Em tarefas que não exigem memória, o portão aprendido $\sigma(\gamma)$ permaneceu próximo de sua inicialização (suprimindo o caminho da memória), e todos os modelos alcançaram 100% de precisão.
  • Construção de Mapas: A Memória Tensorial superou significativamente as baselines em tarefas de navegação em grade de longo alcance onde a informação deve ser retida ao longo do tempo.

5. Significado e Alegações

O artigo afirma que a Memória Tensorial oferece uma maneira prática de lidar com o raciocínio de longo alcance, fornecendo aos Transformers um estado explícito, compacto e estruturado.

• Desacoplamento de Capacidade do Comprimento: Ao usar um tensor de tamanho fixo, o modelo evita a escalabilidade linear de memória dos caches KV, tornando-o teoricamente mais escalável para sequências longas.
• Robustez à Oclusão: O estado persistente permite que o modelo mantenha um "estado de crença" sobre a cena, tornando-o mais robusto a observações ausentes ou oclusão em comparação com modelos que dependem exclusivamente da reconstrução do contexto a partir de um fluxo de tokens.
• Versatilidade: Embora motivada por vídeo de longo alcance, o mecanismo mostrou-se benéfico para linguagem (acumulando contexto global) e imagens (agregação espacial estruturada).
• Fundação para Trabalho Futuro: Os autores veem a implementação atual como uma fundação. Eles sugerem direções futuras incluindo memórias de dimensões mais altas (por exemplo, 4D para memória espacial consciente do tempo), dinâmicas mais ricas (por exemplo, warping de fluxo óptico) e atualizações esparsas para melhorar a eficiência.

Os autores permanecem modestos sobre as limitações atuais, reconhecendo que as atualizações densas de voxel são um gargalo de throughput e que o estado de tamanho fixo pode saturar se uma tarefa exigir o armazenamento de um vasto número de detalhes distintos. No entanto, eles afirmam que adicionar um pequeno estado estruturado é um caminho viável para sistemas de percepção e raciocínio de longo alcance mais capazes.