LaTIM: Misurare le Interazioni Latenti Token-to-Token nei Modelli Mamba

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

February 21, 2025

Autori: Hugo Pitorro, Marcos Treviso

cs.AI

Abstract

I modelli dello spazio degli stati (SSM), come Mamba, sono emersi come un'efficiente alternativa ai trasformatori per la modellizzazione di sequenze a lungo contesto. Tuttavia, nonostante la loro crescente adozione, gli SSM mancano degli strumenti di interpretabilità che sono stati cruciali per la comprensione e il miglioramento delle architetture basate sull'attenzione. Mentre gli sforzi recenti forniscono approfondimenti sui meccanismi interni di Mamba, non decompongono esplicitamente le contribuzioni a livello di token, lasciando delle lacune nella comprensione di come Mamba elabori selettivamente le sequenze attraverso i livelli. In questo lavoro, presentiamo LaTIM, un nuovo metodo di decomposizione a livello di token per sia Mamba-1 che Mamba-2 che consente un'interpretabilità dettagliata. Valutiamo ampiamente il nostro metodo su diverse attività, tra cui la traduzione automatica, la copiatura e la generazione basata sul recupero, dimostrandone l'efficacia nel rivelare i modelli di interazione da token a token di 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.