DLM-Scope: 희소 오토인코더를 활용한 확산 언어 모델의 메커니즘 해석

DLM-Scope: Mechanistic Interpretability of Diffusion Language Models via Sparse Autoencoders

희소 오토인코더(SAE)는 자기 회귀 대규모 언어 모델(LLM)에서 메커니즘 해석을 위한 표준 도구로 자리 잡았으며, 연구자들이 모델의 동작을 이해하고 제어할 수 있는 희소하고 사람이 이해하기 쉬운 특징을 추출할 수 있도록 합니다. 최근 확산 언어 모델(DLM)이 자기 회귀 LLM의 유망한 대안으로 떠오르면서, 이러한 새로운 모델 클래스에 적합한 메커니즘 해석 도구를 개발하는 것이 중요합니다. 본 연구에서는 DLM-Scope을 제시합니다. DLM-Scope은 DLM을 위한 첫 번째 SAE 기반 해석 프레임워크이며, 학습된 Top-K SAE가 신뢰성 있게 해석 가능한 특징을 추출할 수 있음을 보여줍니다. 주목할 만한 점은 SAE 삽입이 DLM에 미치는 영향이 자기 회귀 LLM과 다르다는 것입니다. 일반적으로 LLM에 SAE를 삽입하면 성능 저하가 발생하지만, DLM의 초기 레이어에 SAE를 적용하면 교차 엔트로피 손실을 줄일 수 있으며, 이는 LLM에서는 나타나지 않거나 훨씬 약한 현상입니다. 또한, DLM의 SAE 특징은 확산 과정에서의 효과적인 개입을 가능하게 하며, 종종 LLM 제어보다 더 나은 성능을 보입니다. 게다가, 우리는 DLM에 대한 새로운 SAE 기반 연구 방향을 개척했습니다. SAE가 DLM의 디코딩 순서에 대한 유용한 정보를 제공할 수 있으며, SAE 특징은 DLM의 사후 학습 단계에서도 안정적임을 보여주었습니다. 본 연구는 DLM의 메커니즘 해석을 위한 기반을 마련하고, SAE를 DLM 관련 작업 및 알고리즘에 적용할 수 있는 엄청난 잠재력을 보여줍니다.

Sparse autoencoders (SAEs) have become a standard tool for mechanistic interpretability in autoregressive large language models (LLMs), enabling researchers to extract sparse, human-interpretable features and intervene on model behavior. Recently, as diffusion language models (DLMs) have become an increasingly promising alternative to the autoregressive LLMs, it is essential to develop tailored mechanistic interpretability tools for this emerging class of models. In this work, we present DLM-Scope, the first SAE-based interpretability framework for DLMs, and demonstrate that trained Top-K SAEs can faithfully extract interpretable features. Notably, we find that inserting SAEs affects DLMs differently than autoregressive LLMs: while SAE insertion in LLMs typically incurs a loss penalty, in DLMs it can reduce cross-entropy loss when applied to early layers, a phenomenon absent or markedly weaker in LLMs. Additionally, SAE features in DLMs enable more effective diffusion-time interventions, often outperforming LLM steering. Moreover, we pioneer certain new SAE-based research directions for DLMs: we show that SAEs can provide useful signals for DLM decoding order; and the SAE features are stable during the post-training phase of DLMs. Our work establishes a foundation for mechanistic interpretability in DLMs and shows a great potential of applying SAEs to DLM-related tasks and algorithms.