MemDLM: 메모리 강화형 확산 언어 모델 학습

기존 Diffusion Language Model은 학습과 추론 방식의 차이로 인해 긴 문맥에서 성능이 저하되는 문제가 있었다. MemDLM은 학습 과정에 추론 시뮬레이션을 통합하여 모델 가중치에 정보를 직접 저장하는 '파라미터 메모리'를 도입해 이 문제를 해결한다. 이를 통해 별도의 외부 장치 없이도 수만 토큰의 문서에서 정보를 정확히 추출하는 능력이 비약적으로 향상되었다.

Diffusion Language Model(DLM)은 Transformer의 Attention 메커니즘을 활용해 문장의 모든 토큰을 동시에 처리할 수 있는 장점이 있지만, 학습 시에는 정답 데이터에 마스크를 씌워 한 번에 맞히는 방식을 사용한다. 반면 실제 생성 시에는 노이즈가 섞인 자신의 이전 출력물을 바탕으로 여러 단계에 걸쳐 정답을 찾아가는 과정을 거친다. 이 과정에서 발생하는 오차가 누적되면서 모델이 학습 때 보지 못한 불안정한 상태에 빠지는 '노출 편향' 문제가 발생한다.

MemDLM은 이 문제를 해결하기 위해 학습 그래프 내부에 작은 최적화 루프를 하나 더 추가한다. 이를 통해 학습 중에도 실제 추론과 유사한 '단계적 노이즈 제거' 과정을 시뮬레이션한다. 이때 '고속 가중치(Fast Weights)'라고 불리는 별도의 파라미터 집합이 각 데이터 샘플의 특징적인 정보를 실시간으로 학습하여 일시적인 기억 장치인 '파라미터 메모리' 역할을 수행한다.

결과적으로 베이스 모델은 모든 정보를 토큰 표현에만 의존해 기억할 필요가 없어진다. 복잡하고 긴 문맥 정보의 일부를 가중치 공간으로 넘겨줌으로써 모델은 더 안정적으로 긴 시퀀스를 처리할 수 있게 된다. 이는 마치 사람이 긴 글을 읽을 때 머릿속에 핵심 내용을 요약해두고 세부 사항을 찾아보는 것과 유사한 원리로 동작한다.

MemDLM은 전체 모델 파라미터를 베이스 가중치(θ)와 파라미터 효율적인 고속 가중치(ϕ, LoRA 어댑터)로 분리하여 Bi-level Optimization 문제를 구성한다. 내부 루프는 특정 배치 데이터에 대해 K-단계의 노이즈 제거 궤적을 시뮬레이션하며 ϕ를 업데이트한다. [현재 가중치 ϕ_k-1과 내부 손실 함수의 기울기를 입력으로] → [학습률 η를 곱해 가중치를 갱신하는 경사 하강법 연산을 수행하여] → [새로운 가중치 ϕ_k를 얻고] → [이 값은 해당 데이터 샘플의 국소 궤적 정보를 학습한 일시적인 메모리 상태를 의미한다.]

내부 루프는 '앵커 일관성 궤적' 설계를 따른다. 첫 단계에서는 앵커 상태 x_t보다 더 노이즈가 심한 x_tpre에서 x_t를 예측하도록 ϕ를 학습시키고, 두 번째 단계에서는 x_t에서 최종 정답 x_0를 예측하도록 한다. [앵커 상태 x_t와 노이즈가 더 심한 x_tpre를 입력으로] → [x_tpre에서 x_t를 복원하는 방향으로 고속 가중치 ϕ를 업데이트하여] → [국소 궤적 정보가 담긴 파라미터 메모리를 생성하고] → [이 값은 모델이 추론 시 겪게 될 노이즈 제거 과정을 미리 학습한 결과가 된다.]

외부 루프에서는 내부 루프에서 최적화된 파라미터 메모리 ϕ_2를 조건으로 하여 베이스 모델 θ를 업데이트한다. 계산 복잡도를 줄이기 위해 First-Order Approximation을 적용하여 2차 헤시안 행렬 계산을 피하고, θ에 대한 그래디언트만 계산하여 가중치를 갱신한다. [베이스 가중치 θ와 파라미터 메모리 ϕ_2가 적용된 손실 함수의 기울기를 입력으로] → [학습률 β를 곱해 가중치를 갱신하는 연산을 수행하여] → [최종 베이스 모델 θ를 얻고] → [이 값은 특정 샘플의 세부 정보를 메모리에 맡기고 핵심적인 언어 구조 학습에 집중한 최적의 상태를 의미한다.]

LLaDA-MoE-7B 모델을 기반으로 한 실험에서 MemDLM은 RULER Variable Tracking 벤치마크 성능을 기존 78.8%에서 95.8%로 크게 향상시켰다. 또한 BABILong 벤치마크에서도 LLaDA2.1-mini 모델 기준 47.4%에서 57.0%로 성능이 개선됨을 확인했다.

학습 곡선 분석 결과, MemDLM은 표준 MDLM 방식보다 더 빠르게 수렴하며 최종 학습 손실과 평가 손실 모두 일관되게 낮은 수준을 유지했다. 이는 파라미터 메모리가 모델의 최적화 부담을 효과적으로 덜어주고 있음을 보여준다.

추론 시 내부 루프를 활성화하는 적응형 방식을 적용했을 때, 16K 및 32K와 같은 초장거리 문맥 외삽 상황에서도 성능 하락폭이 기본 모델 대비 훨씬 완만하게 나타났다. 이는 파라미터 메모리가 단순히 학습 데이터를 외우는 것이 아니라 유용한 문맥 표현을 보존하는 데 기여함을 입증한다.

MemDLM 아키텍처는 LLaDA-MoE 및 LLaDA2.1을 백본으로 사용하며, 내부 루프 최적화를 위해 별도의 LoRA 어댑터를 할당한다. 오버헤드를 최소화하기 위해 내부 루프는 트랜스포머 마지막 10% 레이어의 FFN 모듈만을 타겟으로 SGD 최적화를 수행한다.

Bi-level Optimization의 핵심은 내부 루프의 손실 함수 설계에 있다. 실험 결과 단순 Cross-Entropy 손실이 KL Divergence나 Hidden-state Distillation보다 더 효과적인 것으로 나타났다. 이는 토큰 레벨의 직접적인 감독 신호가 궤적 정보를 인코딩하는 데 가장 효율적임을 시사한다.

그래디언트 정규화 측면에서, 전역 방식보다 파라미터별 국소 정규화와 그래디언트 클리핑을 결합했을 때 가장 높은 성능을 보였다. 이는 급격한 태스크 로컬 적응이 필요한 내부 루프의 특성이 반영된 결과이다.

추론 단계에서는 프롬프트를 모델 가중치에 내재화하는 'In-weight Retrieval' 메커니즘이 발현된다. 이는 토큰 간 Attention이 희석되는 병목 현상을 파라미터 공간에서의 정보 검색으로 보완하여 Needle-in-a-Haystack 태스크의 정확도를 높인다.

추론 시 내부 루프를 활성화할 경우 생성 전 프롬프트 적응 단계에서 추가적인 연산 시간이 소요된다. 또한 현재의 추론 절차는 생성 과정 중이 아닌 생성 전 프롬프트에 대해서만 내부 루프를 적용하고 있어, 생성 중 동적인 적응에 대한 연구가 향후 과제로 남아 있다.