NLD-Uni는 교차 패러다임 온폴리시 증류로 확산 LLM의 이해와 생성을 통합한다
Nemotron-Labs-Diffusion-Uni: Elastic Self-Refining Unified Diffusion LLM with Cross-Paradigm On-Policy Distillation
무엇인가
통합 멀티모달 LLM은 크게 두 갈래로 발전해 왔다. 텍스트와 이미지를 모두 이산 토큰으로 바꿔 자기회귀(AR)로 모델링하는 방식과, 텍스트는 AR로 이미지는 연속 확산으로 생성하는 AR+Diff 방식이다. 이 논문이 다루는 것은 세 번째 갈래인 통합 이산 확산 언어 모델(unified DLLM)이다. 마스크 확산(MDM)을 쓰는 이 계열은 여러 토큰을 병렬로 생성해 AR보다 빠르고, 텍스트와 이미지에 하나의 목적함수를 적용해 AR+Diff처럼 상충하는 목표를 조율할 필요가 없다는 장점이 있다. 문제는 성능이다. 기존 통합 DLLM은 도메인 특화 VLM과 연속 흐름 모델에 비해 여전히 큰 격차를 보인다. NLD-Uni는 이 격차를 새 학습 패러다임과 새 아키텍처 두 축으로 좁히겠다고 주장한다.
어떻게 동작하나
아키텍처는 LaViDa-O를 따른 두 개의 탑으로 된 Elastic Mixture-of-Transformer다. 이해 전문가 탑과 생성 전문가 탑이 조인트 어텐션으로 상호작용하고, 이해 탑은 블록 인과 SBD(블록 단위 준자기회귀 확산)를 쓰는 NLD 확산 언어 모델로 초기화하며 생성 탑은 그 복사본에서 출발한다. 이미지는 의미 인코더와 VQ 인코더 양쪽으로 인코딩되며, 토큰화 이후에는 표준 SBD 목적함수를 그대로 적용할 수 있다. SBD는 블록 크기가 1이 되면 AR 모델로 퇴화한다는 성질을 가지는데, 논문은 이 성질을 증류 손실 설계의 발판으로 쓴다.
무엇과 다른가
첫 번째 기여는 자기 정제 확산 과정이다. 표준 MDM은 토큰이 마스크 상태와 깨끗한 상태 사이를 한 번에 오가지만, NLD-Uni는 토큰 나이 a∈{0,…,K}라는 토큰 로컬 이산 인덱스를 도입해 상태 공간을 어휘 집합과 나이의 곱으로 확장한다. 전방 과정에서 토큰은 한 단계씩만 나이를 먹으며, 각 전이는 VQ 임베딩 공간의 볼 트리(ball tree)에서 정의된다. 깊이 d의 공통 조상을 공유하는 토큰 집합을 이용해, 나이가 들수록 더 큰 반경의 이웃으로 균등하게 대체되다가 결국 균등 분포인 완전 마스크 상태에 도달한다. 볼 트리를 쓴 이유는 효율이다. 13만1천 개의 VQ 코드에 대한 쌍별 거리 행렬은 32GB VRAM을 요구하지만, 18깊이 이진 트리는 1D 배열 524KB로 표현된다. 또 DNDM 방식을 확장해 토큰별 점프 시각 구간 (τ, τ′)을 샘플링하고, 하이퍼파라미터 ε이 중간 상태의 지속 시간을 조절한다. ε이 0에 가까우면 DNDM으로 환원되고, 1이면 모든 토큰이 같은 비율로 손상된다. 추론 시 토큰은 [τ′,1]에서는 마스크로 남고, [τ, τ′]에서 점진적으로 정제되며, [0, τ]에서는 고정된다. 이 덕분에 표준 MDM이 할 수 없었던 중간 커밋과 반복 수정이 가능해진다.
어떻게 쓰나
두 번째 기여는 교차 패러다임 온폴리시 증류(OPD)다. OPD는 학생 모델이 스스로 생성한 시퀀스에 교사 감독을 주는 방식인데, 통합 MDM에는 쓸 만한 MDM 교사가 없고 AR VLM이나 연속 흐름 모델은 패러다임이 근본적으로 달라 그대로는 증류할 수 없다. 이해 능력 쪽에서는 학생의 AR 뷰(블록 크기 1)를 AR 교사로 직접 감독하되, 교사 점수를 확산 뷰에 그대로 복사하면 문맥 불일치로 잘못된 신호가 생긴다는 점을 지적한다. 논문이 드는 예는 질의응답에서 학생 답이 "Jason visited San Francisco"일 때, 교사가 앞 두 단어를 보고 "New"에 높은 확신을 준 경우다. 블록 크기 2에서 부분 상태 "[m] Francisco"에 이 점수를 그대로 옮기면 "New Francisco"라는 비문법적 출력을 유도하게 된다. 해결책은 이행적 증류다. 학생의 AR 뷰는 교사 점수로 감독하는 AR OPD 손실을 쓰고, 확산 뷰는 학생 자신의 AR 예측과의 KL 발산으로 감독하는 일관성 손실을 추가한다. 일관성 손실에서는 AR 쪽에 stop-gradient를 걸고 마스크된 위치만 평가한다.
전제와 한계
생성 능력 증류는 더 까다롭다. 연속 흐름 모델은 연속 잠재 공간에서 동작해 정확한 KL 방식의 분포 정합이 불가능하기 때문이다. 논문은 이미지 편집 과제에서 Qwen-Image를 교사로 쓰는 방식을 택한다. 학생 출력을 교사의 연속 VAE로 다시 인코딩하고, 교사 잠재 공간에서 노이즈를 더한 뒤 원본 이미지와 지시를 조건으로 다시 디노이즈해 교사 정제 표적을 만든다. 이렇게 만든 이미지는 학생 생성물과 의미·공간적으로 정렬되면서 아티팩트가 줄고 원본 보존이 좋아, 독립적으로 생성한 교사 출력보다 표적으로 적합하다는 관찰을 제시한다. 저자들도 이를 전통적 의미의 온폴리시는 아니고, 현재 학생 롤아웃에서 교사 정제 표적을 구성하는 방식이라고 인정한다.
실험은 3단계로 진행된다. 1단계는 텍스트-이미지 쌍으로 기본 정렬과 생성 능력을, 2단계는 지시 따르기·이미지 편집·객체 그라운딩·인터리브 생성 데이터를 추가한다. 3단계 통합 OPD에서는 멀티모달 추론에 Nemotron-3 Nano Omni를, 이미지 편집에 Qwen-Image를 교사로 쓰고 T2I SFT 배치를 섞는다. T2I에는 OPD 교사를 두지 않았는데, GenEval과 HPSv3에서 기본 모델이 Qwen-Image나 Flux 같은 교사 후보보다 강했기 때문이다. 학습은 H100 64대로 26일, 그중 최종 OPD 단계가 5일이다. 결과적으로 MMMU, ChartQA, InfoVQA, AI2D와 그라운딩 벤치마크 RefCOCO, ScreenSpot-Pro에서 LLaDa-V, NLD-VLM, LaViDa-O 같은 기존 통합 확산 모델을 큰 폭으로 앞선다. 추론 과제에서는 We-Math, Geometry3k, MathVerse 개선이 두드러지며, OPD-LM 스타일로 교사 점수를 그대로 복사하는 순진한 OPD 베이스라인은 오히려 기저 체크포인트 성능을 떨어뜨렸다. 이미지 편집 ImgEdit에서는 LaViDa-O와 RL 후처리 모델 LaViDa-R1을 앞서고 GPT-4o와 Qwen-Image 교사에만 뒤지며, 기저 대비 +0.36, 독립 교사 샘플을 쓰는 오프폴리시 KD 대비 +0.27을 기록했다. T2I의 HPSv3에서는 Flux.1-dev와 Qwen-Image를 능가한다고 주장한다. 다만 본문에 실린 표의 개별 절대 점수는 발췌문에 수치가 제시되지 않아 인용할 수 없다.
절제 실험은 64 확산 스텝, 2만 스텝 학습으로 통일해 비교했다. 자기 정제 확산은 균등 확산, 마스크 확산, 토큰 편집을 쓰는 마스크 확산보다 평균적으로 우수했고, 객체 검출기로 상위 의미를 보는 GenEval에서는 균등 확산과 비슷했지만 세밀한 DPG와 사람 선호 기반 HPSv3에서는 일관되게 앞섰다. 특히 HPSv3 개선이 가장 커서, 자기 정제가 시각적 충실도 향상에 가장 효과적이라고 해석한다.
개발자 입장에서 이 논문의 실용적 신호는 세 가지다. 첫째, 확산 LLM을 서비스에 올릴 때 병렬 디코딩과 단일 목적함수라는 이점을 유지하면서 특화 모델과의 품질 격차를 줄이는 구체적 경로를 보여준다. 둘째, 이미 학습된 확산 학생에 더 강한 AR 교사를 붙이는 상황을 전제하므로, 교사 점수를 확산 뷰에 직접 복사하지 말고 AR 뷰와 확산 뷰 사이의 일관성 손실로 우회해야 한다는 설계 교훈이 있다. 셋째, 볼 트리 기반 코드북 기하와 토큰 나이 스케줄은 중간 커밋과 반복 수정을 가능하게 하는 추론 절차이므로, 지연 시간과 품질을 맞바꾸는 스텝 스케줄을 튜닝할 여지를 준다.
한계도 명시적이다. 저자들은 더 나은 데이터 혼합과 아키텍처 최적화로 모델을 더 개선할 수 있다고 인정한다. T2I에는 교사를 두지 못했고, 생성 쪽 증류는 KL 정합이 아닌 SFT 손실에 의존하며 업데이트된 학생 분포에 대한 엄밀한 온폴리시도 아니다. 토큰 나이 과정은 τ를 주변화할 수 없어 비마르코프가 되고, 추론 시 (τ, τ′)를 미리 샘플링해야 한다는 제약이 따른다. MDM-prime처럼 서브토큰으로 부분 마스킹을 하는 선행 연구는 32x32 저해상도 이미지에 그쳐 이 규모로 확장되지 않았다는 점도 함께 밝힌다.