TDLM이 확산 언어모델의 반복 계산을 시간 메모리로 재활용해 생성을 가속한다

무엇인가

마스크 확산 언어모델은 응답의 여러 토큰을 동시에 예측하지만, 아직 확정되지 않은 위치도 매 디노이징 스텝마다 다시 순전파를 통과하며 풍부한 은닉 표현을 계산한다. 문제는 관행적 디노이저가 이미 확정된 토큰만 다음 스텝으로 넘길 뿐, 여전히 마스크인 위치가 자신의 이전 표현을 참조할 경로를 주지 않는다는 점이다. 이 계산을 재사용하지 못하면 확정할 토큰 수를 줄여 순전파를 반복하거나, 불확실한 예측을 밀어붙여 답 전체를 망치거나 둘 중 하나를 골라야 한다. 신뢰도 기반 디코딩의 임계값이 바로 이 긴장을 드러낸다. 논문은 확정되지 않은 위치가 자기 자신의 과거 은닉 표현을 참조하게 만들면, 이전 문맥에서 이미 수행한 계산이 이후 예측에 쓰여 더 이른 시점에 신뢰할 만한 예측에 도달할 수 있는지를 묻는다.

어떻게 동작하나

제안 방법 TDLM은 선택된 Transformer 레이어에서, 현재 마스크인 위치가 같은 위치·같은 레이어의 과거 은닉 표현들에만 어텐션하는 시간 어텐션 모듈을 추가한다. 갱신은 잔차 형태로 이뤄지고 마스크 위치에만 적용되며, 좁은 다중 헤드 투영과 확정 토큰 수를 진행 좌표로 쓰는 temporal rotary embedding, 그리고 학습된 제로값 어텐션 싱크(history를 무시할 수 있는 선택지)로 구성된다. 출력 투영을 0으로 초기화하기 때문에 초기화 시점에는 사전학습된 디노이저의 동작이 그대로 보존되고, 백본 파라미터는 얼어붙은 채 시간 모듈만 학습된다. 현재 표현이 질의가 되고 키·값은 같은 레이어의 시간 보정 이전 표현에서 만들어지는데, 이 성질이 아래의 병렬 학습을 가능하게 한다. 토큰이 확정되면 그 위치의 히스토리는 무효화되지만, 블록이 바뀌어도 다른 위치들의 히스토리는 초기화되지 않는다.

무엇과 다른가

학습은 수작업 정답 없이 이뤄진다. 얼어붙은 백본이 스스로 T 스텝의 궤적을 생성하고, 각 토큰의 최종 값과 그것이 확정된 스텝을 기록해 임의 시점의 입력을 복원한다. 시작점 s0=0과 종료 sK=T를 포함해 K-1개 스텝을 균등 샘플링하고, Si에서 마스크였지만 다음 선택 시퀀스 Si+1에서 확정된 위치만 손실 대상으로 삼아 교차엔트로피를 계산한다. 시간 어텐션의 인과 마스크가 j<i 조건을 걸기 때문에 샘플링된 모든 입력을 레이어 단위로 한꺼번에 평가할 수 있고, 논문은 정확한 산술에서 순차 처리와 동등함을 부록에서 증명한다고 밝힌다.

어떻게 쓰나

실험은 LLaDA-8B-Instruct와 Dream-v0-Instruct-7B 두 백본에 대해 진행됐다. LLaDA에는 16개의 분산된 시간 레이어(4헤드, 폭 64), Dream에는 중간 14개 레이어(4헤드, 폭 84)를 붙였고, 백본당 3,500개의 자기 생성 궤적으로 750회 옵티마이저 업데이트만 수행했다. LLaDA 기준 추가 학습 파라미터는 약 0.84%다. 평가는 MATH-500(500문항), GSM8K(1,319문항), MBPP sanitized(257문항) 전체 테스트셋에서 이뤄졌고, 생성 예산 256토큰·블록 크기 32, 신뢰도 임계값 τ=0.9(LLaDA)·0.95(Dream), d3LLM은 엔트로피 임계값 0.5·0.4를 쓴다고 명시했다.

전제와 한계

결과적으로 TDLM은 원본 DLM 대비 최대 7.15배, Fast-dLLM 대비 최대 1.85배 빠른 생성을 보였다. Fast-dLLM 대비 가장 큰 정확도 개선은 LLaDA MBPP에서 2.72%포인트다. 6개 백본-태스크 조합에서 품질 순위는 1위 1회, 2위 4회, 3위 1회였고 지연 순위는 각각 3회, 2회, 1회로, 품질 변화는 태스크에 의존한다고 저자들은 적었다. 128·384토큰 예산으로 확장한 부록 실험에서는 원본 DLM 대비 가속이 8.54배까지 올라간다.

히스토리의 기여를 분리한 실험이 인상적이다. LLaDA 체크포인트를 고정한 채 시간 모듈이 볼 수 있는 문맥을 현재 평가만, 직전 1회, 3회, 5회, 전체 히스토리로 제한했을 때, 전체 히스토리는 정확도 79.08%, NFE 40.08, 평균 0.585초를 기록했고 직전 1회만 볼 때는 77.62%, 44.06, 0.637초, 현재 히스토리만 볼 때는 76.48%, 42.64였다. 히스토리를 늘릴수록 정확도가 오르면서 평가 횟수와 시간이 함께 줄었다는 뜻으로, 같은 체크포인트에서 확정 행동 자체가 달라졌음을 보여준다.

개발자 관점에서 이 방법의 매력은 적응 비용이 작다는 점이다. 백본을 건드리지 않고 소수의 파라미터만 학습하므로 기존 체크포인트 위에 얹을 수 있고, 별도 정답 데이터 없이 모델 자기 출력만으로 학습이 된다. 다만 공짜는 아니다. 시간 레이어 하나당 히스토리 검색에 O(NHd_t)가 추가되고, 전체 히스토리 저장에 O(L_t·N·H·d_t) 메모리가 필요하며, 제한 없는 어텐션 연산은 디노이징 스텝 수에 대해 제곱으로 늘어난다. 논문 스스로 절약되는 백본 평가가 이 오버헤드를 상쇄할 때 유용하다고 못박으므로, 도입 시 NFE와 실제 wall time을 함께 재고 자신의 스텝 수·응답 길이 분포에서 이득이 나는지 확인해야 한다. 신뢰도 임계값도 백본마다 다르게 튜닝해야 한다(LLaDA 0.9, Dream 0.95).

저자들이 밝힌 한계도 분명하다. 자기 증류는 교사 궤적의 품질 한계를 그대로 물려받으며, 얼어붙은 백본에서 뽑은 토큰으로는 체계적인 교사 오류를 교정할 증거가 부족하다. 백본과 시간 어텐션을 함께 파인튜닝하거나 더 높은 품질의 학습 데이터를 쓰는 것이 향후 과제로 남는다. 또한 희소하게 샘플링된 교사 궤적과 학생의 다중 토큰 확정이 서로 다른 히스토리를 노출할 수 있고, 긴 응답·반복 측정·공간 KV 캐시 및 조기 종료와의 결합은 아직 평가되지 않았다. 현재 은닉 표현만 쓰는 용량 일치 어댑터와의 비교도 후속 분석 과제로 제시된다.