확산 언어모델의 표현공간 MMD 후학습으로 샘플 품질과 디코딩 병렬성을 함께 올린다

무엇인가

확산 언어모델(DLM)은 노이즈가 섞인 입력을 이산 토큰 공간이나 연속 잠재 공간에서 반복 정제해 텍스트를 만든다. 이산 모델은 토큰별 교차엔트로피로 조건부 토큰 주변분포를 맞추고, 연속 모델은 깨끗한 잠재값의 조건부 평균을 제곱오차로 회귀한다. 문제는 이런 조건부 추정이 반복 샘플링을 안내할 뿐, 적은 스텝에서 고품질 샘플을 보장하지 않는다는 점이다. 기존 확산 증류 기법들은 교사 모델의 목표값이나 생성기 출력으로 함께 학습되는 보조 모델을 요구한다. 이 논문은 그 대신 생성 샘플과 참조 샘플만으로 분포 매칭 학습 신호를 직접 계산할 수 있는지 묻는다.

어떻게 동작하나

제안 방법의 핵심은 비교 공간을 고정된 사전학습 DLM의 특징 공간으로 잡는 것이다. 시퀀스를 하나의 임베딩으로 풀링하지 않고 토큰 위치별 문맥 특징을 그대로 유지해, 추출기 한 번의 순전파로 시퀀스당 여러 개의 관측치를 얻는다. 무조건 생성에서는 모든 위치, 프롬프트 조건 생성에서는 응답 위치, 마스크 확산에서는 마스크된 위치만 점수화한다. 두 시퀀스의 커널 유사도는 각 시퀀스의 점수화 위치 쌍에 대한 평균으로 정의되며(수식 4), 이는 경험적 커널 평균 임베딩의 내적이므로 시퀀스에 대한 양의 준정부호 커널이 된다. MMD 추정에서는 같은 시퀀스 안의 토큰 쌍을 모두 제외해 불편 추정량을 유지한다(수식 5). 조건부 학습에서는 조건당 참조 1개와 서로 독립적으로 샘플링한 생성 시퀀스 B개(≥2)만으로 손실을 계산한다(수식 6). 참조-참조 항은 θ와 무관해 최적화에서 생략된다. 생성-생성 항은 참조 특징에 대한 끌림과 생성 특징끼리의 유사성 억제를 균형 맞추는 역할을 하는데, B=1로 두면 끌림만 남아 이산 모델에서는 성능이 떨어지고 연속 모델에서는 급격히 붕괴한다고 저자들은 밝힌다.

무엇과 다른가

이산 모델 실현에서는 디노이저 한 번의 순전파가 만든 토큰 분포에서 B개 시퀀스를 독립 샘플링한다. 범주형 샘플링은 미분이 불가능하므로 REINFORCE로 최적화하며, 같은 참조와 손상 상태에 대해 G개의 독립적인 MMD 배치를 뽑아 leave-one-out 베이스라인을 만든다(수식 8). 정책 경사 대리 손실은 배치 보상에 배치 내 모든 시퀀스의 로그 확률 합을 곱하는 형태다(수식 9). 마스크 확산용 MDLM-MMD와, 마스크 예측과 자체 토큰 예측 정제를 함께 학습하는 하이브리드 마스크-균일 모델 DMax용 DMax-MMD 두 가지로 구현된다. 연속 모델 실현인 ELF-MMD는 사전학습된 ELF에서 초기화한 1스텝 잠재 생성기를 학습한다. MMD 특징은 ELF를 t=1, 자기조건화 입력을 자기 자신으로 두고 통과시켜 얻으며, 그래디언트가 추출기 입력과 생성 잠재값을 통해 직접 전파된다. 추론 시에는 자기조건화로 예측을 반복 정제해 다단계 생성으로 확장하고(수식 12), 학습 중에는 예비 정제 스텝 수를 무작위로 뽑아 마지막 스텝에서만 역전파한다. 선택적으로 IRD를 적용해 다단계 궤적을 1스텝으로 증류한다.

어떻게 쓰나

무조건 생성 실험은 OpenWebText를 길이 1024로 패킹해 학습하고, 샘플 1000개를 GPT-2 Large 기반 생성 퍼플렉시티와 평균 유니그램 엔트로피로 평가했다. MDLM-MMD는 8·16·32 스텝 모두에서 동일 엔트로피 기준으로 IDLM 계열 베이스라인보다 낮은 생성 퍼플렉시티를 보였고, 참조 데이터 엔트로피 지점에서 곡선을 보간하면 IDLM 대비 약 17~21% 낮다고 보고한다. 연속 모델은 T5-small과 GPT-2 Large의 최종 레이어 은닉 표현을 잠재 공간으로 쓰는 ELF-B로 실험했는데, ELF-MMD는 대부분의 샘플링 예산에서 ELF를 앞섰고 8스텝에서는 ELF-PD보다 낮은 퍼플렉시티와 참조에 가까운 엔트로피를 기록했다. IRD를 붙이면 4스텝에서 T5와 GPT-2 인코더 각각 생성 퍼플렉시티가 약 30포인트, 10포인트 줄어든다.

전제와 한계

조건부 생성은 TinyGSM으로 학습해 GSM8K 테스트 정확도로 평가했고, 최대 길이는 512다. MDLM-MMD는 신뢰도 임계값을 스윕한 결과 약 49스텝에서 약 54% 정확도에 도달해 MDLM, IDLM, IDLM-REINFORCE, DiDi-Instruct를 앞섰다. 연속 모델은 GPT-2 Small 최종 은닉 표현을 잠재 공간으로 쓰는 ELF-B를 주 비교로 삼았고, 타임스텝 시프트를 각각 32와 128로 적용했다. 4스텝에서 정확도가 14.2%에서 20.8%로, 8스텝에서 27.5%에서 32.5%로 오르며, 같은 조건의 ELF-PD는 각각 15.9%와 23.5%였다. 64스텝에서는 ELF-MMD+IRD가 36.3%로 ELF-MMD의 35.2%, ELF의 31.6%를 웃돈다. pass@k도 작은 샘플 예산에서 개선되지만 k가 커지면 격차가 줄어든다.

16B 규모에서는 LLaDa2.0-Mini 기반 하이브리드 마스크-균일 모델인 DMax-Math와 DMax-Coder 체크포인트에서 출발해 DMax-Math-MMD와 DMax-Coder-MMD를 만들었다. 8장의 NVIDIA H100에서 400 스텝을 돌렸고 Math는 약 13분, Coder는 약 19분이 걸려 각각 약 1.7, 2.5 GPU-시간에 해당한다(초기화와 평가 제외). 임계값은 0.85와 0.9를 선택했다. 수학 벤치마크 4종(GSM8K, MATH500, Minerva-Algebra, ASDIV)에서 정확도를 유지하거나 높이면서 포워드당 생성 토큰 수(TPF)를 10.3~16.5% 늘렸고, 코드에서는 HumanEval-Instruct 2.4%포인트, MBPP-Instruct 3.8%포인트 정확도를 올리면서 TPF도 함께 증가시켰다. 손실 설계 절제에서는 선형 커널 MMD(특징 평균 MSE), 평균 풀링한 시퀀스 수준 RBF, 끌림만 남긴 변형, 전체 특징 텐서 회귀를 비교했는데, 토큰 수준 RBF가 가장 좋은 정확도-계산 트레이드오프를 보였고 이산 모델에서 이득이 더 컸다. 끌림 전용과 특징 회귀는 특히 연속 모델에서 정확도가 낮았다.

실무 관점에서 이 방법의 매력은 학습 비용이다. 전체 샘플링 궤적을 저장하거나 보조 판별 모델을 함께 학습할 필요 없이, 이미 갖고 있는 사전학습 DLM을 특징 추출기로 재사용해 사후학습 몇백 스텝만으로 디코딩 병렬성과 정확도를 동시에 조정할 수 있다는 점은 자체 서빙 환경에서 스텝 수를 줄이려는 팀에 실질적이다. 다만 코드 저장소(yandex-research/dlm-mmd)를 확인할 때는 RBF 대역폭, 어떤 레이어의 특징을 쓸지, 참조 샘플을 어떻게 확보할지가 성능을 좌우한다는 점을 먼저 따져야 한다. 조건부 태스크에서는 조건당 참조 1개로도 학습이 되지만 생성 배치 B는 2 이상이어야 하고, 배치 그룹 G와 샘플링 온도 스윕이 결과 곡선을 크게 바꾼다.

저자들이 명시한 한계는 분명하다. 성능이 RBF 대역폭과 분포 매칭에 쓰는 표현 공간 선택에 의존하므로 표현 설계 자체가 추가 탐구 방향이라고 밝힌다. 실험에서는 깨끗한 입력에 대해 단일 레이어의 특징만 사용했고, 중간 확산 타임스텝의 손상된 입력에서는 성능이 좋아지지 않았다. 이론적으로도 주입적 특징 맵에 characteristic 커널을 적용하면 모집단 MMD가 0일 때 원래 분포가 일치한다는 결과는 있지만, 이 논문이 쓰는 토큰 특징 분포를 맞추는 것이 시퀀스 분포를 유일하게 식별한다는 보장은 없다고 스스로 인정한다. 향후 과제로 여러 레이어나 노이즈 수준의 특징 결합, 추출기 앙상블, 다른 학습 목적함수와의 결합을 제안한다.