LaTIM: Mamba 모델의 잠재적 토큰 간 상호작용 측정

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

February 21, 2025

저자: Hugo Pitorro, Marcos Treviso

cs.AI

초록

Mamba와 같은 상태 공간 모델(SSMs)은 장문 맥락 시퀀스 모델링에서 트랜스포머의 효율적인 대안으로 부상하고 있습니다. 그러나 이러한 모델의 활용이 증가함에도 불구하고, SSMs는 어텐션 기반 아키텍처를 이해하고 개선하는 데 중요한 역할을 해 온 해석 가능성 도구를 갖추지 못하고 있습니다. 최근 연구들은 Mamba의 내부 메커니즘에 대한 통찰을 제공하지만, 토큰 단위 기여를 명시적으로 분해하지 않아 Mamba가 계층 간에 어떻게 선택적으로 시퀀스를 처리하는지에 대한 이해에 공백이 남아 있습니다. 본 연구에서는 Mamba-1과 Mamba-2 모두에 적용 가능한 새로운 토큰 수준 분해 방법인 LaTIM을 소개하며, 이를 통해 세밀한 해석 가능성을 가능하게 합니다. 우리는 기계 번역, 복사, 검색 기반 생성 등 다양한 작업에 걸쳐 이 방법을 광범위하게 평가하며, 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.