LaTIM: Измерение латентных взаимодействий между токенами в моделях Mamba

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

February 21, 2025

Авторы: Hugo Pitorro, Marcos Treviso

cs.AI

Аннотация

Модели пространства состояний (State Space Models, SSMs), такие как Mamba, стали эффективной альтернативой трансформерам для моделирования последовательностей с длинным контекстом. Однако, несмотря на их растущее применение, SSMs не обладают инструментами интерпретируемости, которые были ключевыми для понимания и улучшения архитектур, основанных на механизме внимания. Хотя недавние исследования проливают свет на внутренние механизмы Mamba, они не проводят явного разложения вкладов на уровне токенов, оставляя пробелы в понимании того, как Mamba избирательно обрабатывает последовательности на разных слоях. В данной работе мы представляем LaTIM — новый метод декомпозиции на уровне токенов для Mamba-1 и Mamba-2, который обеспечивает детальную интерпретируемость. Мы тщательно оцениваем наш метод на различных задачах, включая машинный перевод, копирование и генерацию на основе поиска, демонстрируя его эффективность в раскрытии паттернов взаимодействия токенов в 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.