MetaState: 이산 확산 언어 모델을 위한 지속적인 작업 메모리

MetaState: Persistent Working Memory for Discrete Diffusion Language Models

이산 확산 언어 모델(dLLM)은 마스크된 시퀀스를 반복적으로 노이즈 제거하면서 텍스트를 생성합니다. 자기 회귀 모델과 비교할 때, 이러한 방식은 자연스럽게 병렬 디코딩, 양방향 컨텍스트 및 유연한 생성 패턴을 지원합니다. 그러나 일반적인 dLLM은 각 노이즈 제거 단계에서 현재 마스크된 시퀀스에만 의존하며, 샘플링 및 재마스크 후 중간 연속 표현은 버려집니다. 우리는 이러한 병목 현상을 **정보 격리(Information Island)** 문제라고 부릅니다. 이는 단계 간의 불필요한 재계산을 초래하고 단계 간 일관성을 저하시킬 수 있습니다. 우리는 **MetaState**라는 경량 순환 증강 기술을 통해 이 문제를 해결합니다. MetaState는 동결된 dLLM 핵심에 지속적이고 고정 크기의 작업 메모리를 제공하며, 이는 시퀀스 길음에 독립적입니다. **MetaState**는 세 가지 학습 가능한 모듈로 구성됩니다. 첫째, 핵심 활성화를 메모리 슬롯으로 읽어들이는 크로스 어텐션 믹서(cross-attention Mixer)입니다. 둘째, 노이즈 제거 단계 간 정보를 통합하는 GRU 스타일 업데이트(GRU-style Updater)입니다. 셋째, 업데이트된 메모리를 핵심 활성화로 다시 공급하는 크로스 어텐션 인젝터(cross-attention Injector)입니다. 우리는 이러한 모듈을 K-단계 언롤링을 통해 학습시켜 미세 조정 과정에서 다단계 노이즈 제거 동역학에 노출시킵니다. LLaDA-8B 및 Dream-7B 모델에서 **MetaState**는 무시할 만한 학습 가능한 파라미터를 추가하면서 핵심을 동결된 상태로 유지하며, 동결된 기준 모델보다 일관되게 정확도를 향상시킵니다. 이러한 결과는 지속적인 단계 간 메모리가 이산 확산 언어 모델에서 노이즈 제거 단계를 연결하고 생성 품질을 향상시키는 효과적인 메커니즘임을 보여줍니다.

Discrete diffusion language models (dLLMs) generate text by iteratively denoising a masked sequence. Compared with autoregressive models, this paradigm naturally supports parallel decoding, bidirectional context, and flexible generation patterns. However, standard dLLMs condition each denoising step only on the current hard-masked sequence, while intermediate continuous representations are discarded after sampling and remasking. We refer to this bottleneck as the \textbf{Information Island} problem. It leads to redundant recomputation across steps and can degrade cross-step consistency. We address this limitation with \textbf{MetaState}, a lightweight recurrent augmentation that equips a frozen dLLM backbone with a persistent, fixed-size working memory that remains independent of sequence length. \textbf{MetaState} consists of three trainable modules: a cross-attention Mixer that reads backbone activations into memory slots, a GRU-style Updater that integrates information across denoising steps, and a cross-attention Injector that feeds the updated memory back into backbone activations. We train these modules with $K$-step unrolling to expose them to multi-step denoising dynamics during fine-tuning. On LLaDA-8B and Dream-7B, \textbf{MetaState} introduces negligible trainable parameters while keeping the backbone frozen, and it consistently improves accuracy over frozen baselines. These results demonstrate that persistent cross-step memory is an effective mechanism for bridging denoising steps and improving generation quality in discrete diffusion language models.