TL;DR

Nemotron-Labs TwoTower는 사전학습된 Nemotron-3-Nano 30B A3B 백본을 기반으로 AR 문맥 타워를 고정하고 trainable diffusion/denoiser 타워를 추가해 블록 단위 mask-diffusion 생성을 수행하는 모델이다. 디노이저는 블록 내 양방향 self-attention과 레이어 정렬 교차어텐션, 문맥 시드된 Mamba-2 상태를 사용해 마스크된 토큰을 반복적으로 복원하며 adaLN-single로 시간 조건화를 적용해 단계별 거동을 조절한다. 기본 운영점(block_size=16, γ=0.8, BF16, 2×H100)에서 집계 품질은 AR 기준의 98.7%를 유지하면서 wall-clock 생성 처리량은 2.42배에 도달해 속도와 품질 사이의 유의미한 트레이드오프를 제공한다. 이 설계는 프롬프트를 한 번만 인코딩해 블록별 반복에서 중복 계산을 줄이며 디노이저만 적응시키는 접근으로 대규모 생성 워크로드에서 처리량 개선을 노리는 응용에 적합하다. 한편 신뢰도 임계값과 블록 크기의 조정은 품질과 속도의 균형을 직접 좌우하므로 사용자는 운영점 튜닝을 통해 요구 성능에 맞춰야 한다.

핵심 역량

• 블록 단위 mask-diffusion으로 프롬프트를 한 번 인코딩한 뒤 여러 토큰을 병렬로 예측해 생성 처리량을 높일 수 있다.
• 다층 레이어 정렬 교차어텐션과 Mamba-2 상태 시딩을 통해 문맥 정보를 디노이저에 레이어별로 제공해 문맥 일관성 있는 블록 복원을 수행한다.
• AR-only 및 mock-AR 모드를 제공해 단일 GPU 환경에서도 표준 autoregressive 생성과 호환되며, 두 타워 전용 모드로 속도-정확도 트레이드오프를 조정할 수 있다.

강점

• 기본 운영점에서 AR 기반 백본의 집계 품질의 98.7%를 유지하면서 전체 생성의 wall-clock 처리량을 2.42배까지 향상한 점이 검증되어 실사용 처리량 이점을 제공한다.
• 문맥 타워를 한 번만 인코딩하고 레이어별 KV와 Mamba 상태를 재사용하는 구조로 블록 단위 반복 생성의 오버헤드를 낮추어 높은 처리량을 달성한다.
• 백본과 동일한 Mamba-2/MoE 하이브리드 구조를 사용하며 디노이저만 학습하는 접근으로 전체 파라미터를 재학습하지 않고도 디퓨전 생성 기능을 추가했다.

훈련 방식

사전학습된 Nemotron-3-Nano 30B A3B 백본을 초기화로 사용해 AR/context 타워는 고정하고 diffusion/denoiser 타워만 약 2.1T 토큰으로 마스크-디퓨전 목적함수로 추가 학습해 블록 복원 능력을 적응시켰다.

알아두면 좋은 것

• 배포 체크포인트는 AR/context 타워와 diffusion/denoiser 타워를 모두 포함해 두 타워 모두 다운로드 가능하며 디퓨전 추론은 두 GPU(예: 2×H100)를 권장한다.
• 디노이저 타워는 약 2.1T 토큰으로 마스크-디퓨전 목적함수로 적응 훈련되었고, 백본은 약 25T 토큰으로 사전학습되었다.
• 기본 운영점은 confidence unmasking(γ=0.8), block_size=16, BF16 정밀도, 2×H100 또는 2×A100 환경에서 측정된 성능(98.7% 품질 유지, 2.42× 처리량)을 제공한다.
• 라이선스는 NVIDIA Nemotron Open Model License로 상업적 사용이 허용되는 범위와 별도 신뢰성 약관이 명시되어 있다.

기술적 특징

• 두 타워 구조는 AR/context 타워를 frozen으로 두고 diffusion/denoiser 타워만 학습해 문맥 인코딩을 재사용하는 방식으로 동작한다. 이로 인해 프롬프트를 한 번만 처리하면 되므로 블록 단위 반복에서 발생하는 중복 계산을 크게 줄이고 추론 처리량을 개선한다.
• 디노이저는 블록 내에서 bidirectional self-attention을 사용해 마스크된 위치들을 병렬로 예측하고, 레이어 정렬된 교차어텐션으로 각 디노이저 층이 문맥 타워의 동일한 층의 KV를 참조해 문맥 일관성을 확보한다. 이러한 레이어 정렬은 디노이저가 문맥의 중간 표현을 직접 활용할 수 있게 해 복원 품질을 유지하도록 돕는다.
• Mamba-2 상태 시딩은 문맥 타워의 conv/SSM 경계 상태를 디노이저 Mamba 레이어의 초기 상태로 주입하는 방식으로 구현되어, 디퓨전 복원 과정에서 문맥 기반 초기화를 제공해 안정성과 효율을 높인다. 이 시드는 인과적(causal) 방식으로 전달되어 미래 정보 누수 없이 문맥을 반영한다.
• 시간 조건화는 adaLN-single 모듈로 처리되어 각 디노이저 층에 스케일·시프트·게이트를 타임스텝별로 주입한다. 이 모듈은 약 1.5M의 추가 파라미터를 도입해 디노이징 진행도에 따른 레이어 반응을 조정하며, confidence unmasking과 결합해 단계별 복원 전략을 가능하게 한다.
• 활성화된 파라미터 절약을 위해 MoE 전문화가 적용되며 128개의 전문가 중 라우팅으로 6개가 활성화되는 설정을 사용해 토큰당 활성 파라미터를 약 3B 수준으로 유지한다. 이 설계는 계산 효율과 용량 사이의 균형을 제공한다.

차별점

• AR 문맥 타워는 고정한 채로 디노이저 타워만 학습시키는 설계로 문맥 재사용을 통해 반복 생성의 계산 비용을 절감한다.
• 블록 단위 mask-diffusion과 신뢰도 기반 언마스킹(confidence unmasking)을 결합해 반복당 다수 토큰을 확정함으로써 AR 대비 실시간 처리량을 크게 높였다.
• 디노이저의 레이어 정렬 교차어텐션과 문맥 기반 Mamba-2 시딩을 통해 레이어별 중간 표현을 직접 활용하는 점이 일반적인 단일타워 디퓨전 접근과 구별된다.

벤치마크

이미지 분석