LaTIM: Messung latenter Token-zu-Token-Interaktionen in Mamba-Modellen

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

February 21, 2025

Autoren: Hugo Pitorro, Marcos Treviso

cs.AI

Zusammenfassung

Zustandsraummodelle (State Space Models, SSMs), wie Mamba, haben sich als effiziente Alternative zu Transformern für die Modellierung langer Kontextsequenzen etabliert. Trotz ihrer zunehmenden Verbreitung fehlen SSMs jedoch die Interpretierbarkeitstools, die für das Verständnis und die Verbesserung von auf Aufmerksamkeit basierenden Architekturen entscheidend waren. Während jüngste Bemühungen Einblicke in die internen Mechanismen von Mamba bieten, zerlegen sie die tokenweisen Beiträge nicht explizit, wodurch Lücken im Verständnis bestehen, wie Mamba Sequenzen über verschiedene Schichten hinweg selektiv verarbeitet. In dieser Arbeit stellen wir LaTIM vor, eine neuartige Methode zur tokenweisen Zerlegung für sowohl Mamba-1 als auch Mamba-2, die eine feinkörnige Interpretierbarkeit ermöglicht. Wir evaluieren unsere Methode umfassend in verschiedenen Aufgaben, einschließlich maschineller Übersetzung, Kopieren und retrieval-basierter Generierung, und demonstrieren ihre Effektivität bei der Aufdeckung der Token-zu-Token-Interaktionsmuster von 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.