LaTIM: Medición de Interacciones Latentes entre Tokens en Modelos Mamba

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

February 21, 2025

Autores: Hugo Pitorro, Marcos Treviso

cs.AI

Resumen

Los modelos de espacio de estados (SSMs, por sus siglas en inglés), como Mamba, han surgido como una alternativa eficiente a los transformadores para el modelado de secuencias de contexto largo. Sin embargo, a pesar de su creciente adopción, los SSMs carecen de herramientas de interpretabilidad que han sido cruciales para comprender y mejorar las arquitecturas basadas en atención. Aunque esfuerzos recientes proporcionan información sobre los mecanismos internos de Mamba, no descomponen explícitamente las contribuciones a nivel de tokens, dejando vacíos en la comprensión de cómo Mamba procesa selectivamente las secuencias a través de las capas. En este trabajo, presentamos LaTIM, un método novedoso de descomposición a nivel de tokens tanto para Mamba-1 como para Mamba-2 que permite una interpretabilidad de grano fino. Evaluamos exhaustivamente nuestro método en diversas tareas, incluyendo traducción automática, copia y generación basada en recuperación, demostrando su eficacia para revelar los patrones de interacción token a token de 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.