LaTIM: Medindo Interações Latentes entre Tokens em Modelos Mamba

LaTIM: Measuring Latent Token-to-Token Interactions in Mamba Models

February 21, 2025

Autores: Hugo Pitorro, Marcos Treviso

cs.AI

Resumo

Modelos de espaço de estados (SSMs), como o Mamba, surgiram como uma alternativa eficiente aos transformers para a modelagem de sequências de contexto longo. No entanto, apesar de sua crescente adoção, os SSMs carecem das ferramentas de interpretabilidade que têm sido cruciais para a compreensão e melhoria das arquiteturas baseadas em atenção. Embora esforços recentes forneçam insights sobre os mecanismos internos do Mamba, eles não decompõem explicitamente as contribuições em nível de token, deixando lacunas na compreensão de como o Mamba processa seletivamente as sequências através das camadas. Neste trabalho, apresentamos o LaTIM, um novo método de decomposição em nível de token para o Mamba-1 e Mamba-2 que permite uma interpretabilidade refinada. Avaliamos extensivamente nosso método em diversas tarefas, incluindo tradução automática, cópia e geração baseada em recuperação, demonstrando sua eficácia em revelar os padrões de interação token-a-token do Mamba.

English

State space models (SSMs), such as Mamba, have emerged as an efficient alternative to transformers for long-context sequence modeling. However, despite their growing adoption, SSMs lack the interpretability tools that have been crucial for understanding and improving attention-based architectures. While recent efforts provide insights into Mamba's internal mechanisms, they do not explicitly decompose token-wise contributions, leaving gaps in understanding how Mamba selectively processes sequences across layers. In this work, we introduce LaTIM, a novel token-level decomposition method for both Mamba-1 and Mamba-2 that enables fine-grained interpretability. We extensively evaluate our method across diverse tasks, including machine translation, copying, and retrieval-based generation, demonstrating its effectiveness in revealing Mamba's token-to-token interaction patterns.