기술 요약: Nemotron-TwoTower

문제 정의

자기회귀(Autoregressive, AR) 언어 모델은 텍스트 생성 분야를 주도하고 있으나, 순차적인 한 토큰씩 디코딩하는 과정으로 인해 처리량(throughput) 병목 현상을 겪는다. 이산 확산 언어 모델(Discrete diffusion language models)은 토큰을 병렬로 생성하고 이를 반복적으로 정제함으로써 더 높은 처리량과 더 세밀한 제어 가능성을 제공할 수 있는 유망한 대안이다. 그러나 기존의 확산 방식은 일반적으로 하나의 네트워크가 깨끗한 컨텍스트 토큰을 표현하는 역할과 오염된 토큰을 반복적으로 노이즈 제거(denoising)하는 두 가지 서로 다른 역할을 동시에 수행하도록 강제한다. 이러한 결합은 단일 가중치 집합이 상충하는 두 목적을 모두 수행하게 만들어, 각 역할에서 탁월한 성능을 내는 능력을 제한한다. 또한, 이러한 역할들을 완전히 분리하는 것이 Mamba, 어텐션(attention), 그리고 MoE(Mixture-of-Experts) 레이어가 결래된 현대적인 하이브리드 아키텍처 상에서 대규모로 가능한지 여부도 불분명하다.

방법론

저자들은 Nemotron-3-Nano-30B-A3B 하이브리드 아키텍처(Mamba-2, self-attention, MoE 레이어가 교차 배치됨)를 기반으로, 컨텍스트 표현과 노이즈 제거 역할을 두 개의 별도 타워로 분리한 블록 단위 자기회귀 확산 모델인 TwoTower를 제안한다.

아키텍처

이 모델은 상보적인 역할을 가진 두 개의 사전 학습된 네트워크 복사본을 활용한다:

1. 동결된 AR 컨텍스 타워(Frozen AR Context Tower): 이 타워는 동결된 상태로 유지되며 프롬프트와 이미 확정된(committed) 토큰들에 대해 인과적 자기회귀 모델로 동작한다. 이는 깨끗한 토큰을 처리하여 레이어 정렬된 키-값(KV) 캐시와 Mamba 상태를 생성한다. 이를 통해 백본(backbone) 고유의 자기회귀 능력을 보존하고 블록 간 재사용 가능한 상태를 유지한다.
2. 학습 가능한 디퓨전 디노이저 타워(Trainable Diffusion Denoiser Tower): 이 타워는 노이즈가 섞인 블록을 정제하도록 학습된다. 현재 노이즈가 섞인 블록 내에서는 **양방향 블록 어텐션(bidirectional block attention)**을 사용하여 토큰들이 서로를 참조할 수 있게 하며, 동시에 동결된 컨텍스 타워의 해당 레이어들에 대해 인과적 교차 어텐션(causal cross-attention)을 유지한다. 이러한 레이어 정렬된 교차 어텐션은 디노이저에게 백본 표현에 대한 다중 스케일 접근 권한을 제공하며, 이는 최종 은닉 상태(final hidden state)만을 브로드캐스트하던 기존 방식과 차별화된다.

생성 과정

생성은 블록 단위(Algorithm 1)로 진행된다:

• 컨텍스 타워는 프로프트와 확정된 토큰들을 인과적으로 처리한다.
• 새로운 블록은 $S$개의 마스크된 토큰으로 초기화된다.
• 디노이저 타워는 이 블록을 $T$ 단계에 걸쳐 반복적으로 정제한다. 각 레이어에서 디노이저는 과거 블록들의 컨텍스 타워 KV 캐시와 현재 블록의 양방향 KV 캐시에 어텐션을 수행한다.
• 블록이 충분히 정제되면(신뢰도 임계값 기준), 해당 블록은 확정(commit)된다. 그 후 컨텍스 타워는 새로운 깨끗한 토큰들로 상태를 업데이트하며, 다음 블록을 위한 과정을 반복한다.

학습 및 컨디셔닝

• 마스크된 확산(Masked Diffusion): 디노이저는 마스크된 확산 목적 함수를 사용하여 학습된다. 깨끗한 블록은 선형 스케줄에 따라 토큰을 [MASK]로 대체하여 오염시킨다. 모델은 마스크된 위치에서의 원래 토큰에 대한 음의 로그 가능도(negative log-likelihood)를 최소화한다.
• 시간 컨디셔닝(Time Conditioning): 디노이저는 적응형 레이어 정규화(adaptive layer normalization, adaLN)를 사용하여 확산 타임스텝 $t$에 조건화된다. 이는 이미지 확산에서 흔히 사용되지만 마스크된 언어 모델에서는 덜 활용되는 기술이다. 이는 약 1.5M의 미미한 파라미터를 추가하지만 샘플링 품질을 향상시킨다.
• 커리큘럼(Curriculum): 학습은 2단계 커리큘럼을 따른다. 1단계는 광범한 범위를 커버하기 위해 다양한 데이터 혼합물을 사용하며, 2단계는 고품질 및 STEM 중심 소스의 가중치를 높인다. 학습 블록 크기는 병렬성과 컨디셔닝 품질 사이의 균형을 맞추기 위해 점진적으로 감소시킨다 ($S=32$에서 $S=16$으로).
• 샘플링: 저자들은 신뢰도 기반 언마스킹(confidence-unmasking) 전략(Algorithm 2)을 채택한다. 각 단계에서 디노이저는 모든 마스크된 토큰을 병렬로 예측한다. 신뢰도가 임계값 $\gamma$보다 높은 토큰은 확정되는 반면, 불확실한 토큰은 추가 정제를 위해 마스크 상태로 남는다. 이를 통해 확정되는 토큰의 수가 적응적으로 결정된다.

주요 결과

Nemotron-3-Nano-30B-A3B 백본(30B 파라미터, 하이브리드 MoE)에서 평가한 결과, Nemotron-TwoTower는 다음과 같은 성과를 보였다:

• 품질-처리량 트레이드오프: 모델은 자기회귀 베이스라인의 종합 벤치마크 품질의 **98.7%**를 유지하면서, 실제 시간(wall-clock) 생성 처리량을 2.42배 향상시켰다.
• 벤치마크 성능:
  • 일반 지식: AR 베이스라인과 약 1점 차이 이내의 성능을 유지한다.
  • 코드 및 수학: 베이스라인 대비 약간의 성능 저하를 보인다.
  • 상식 및 다국어: 성능이 회복되거나 약간 향상된다.
• 설계 절제 실험(Ablations):
  • 양방향 Mamba: 양방향 Mamba 변형을 테스트한 결과 미미한 품질 향상이 있었으나 SSM 연산량이 두 배로 증가했다. 최종 설계에서는 Mamba를 인과적(causal)으로 유지한다.
  • 시간 컨디셔닝: 명시적인 시간 컨디셔닝이 생성, 코드 및 수학 점수를 유의미하게 개선했다.
  • 타워 분리: 컨텍스 타워를 동결하고 별도의 디노이서를 학습시키는 것이 최선의 결과를 냈다. 타워 간 가중치를 공유하거나 컨텍스 타워에 AR 학습을 계속하는 것은 심각한 정확도 저하를 초래했다.
• 블록 크기 민감도: 더 작은 학습 블록 크기($S=16$)가 더 큰 크기($S=128$)보다 품질을 개선했으나 처리량은 감소시켰다. 저자들은 품질 유지와 속도 사이의 균형을 위해 $S=16$을 기본값으로 선택했다. 학습 블록 크기보다 큰 블록으로 샘플링할 경우 심각한 성능 저하가 발생했다.

의의 및 주장

본 논문은 Nemotron-TwoTower가 전체 모델을 재학습하지 않고도 사전 학습된 자기회귀 표현이 학습 가능한 확산 디노이저를 지원할 수 있음을 성공적으로 입증했다고 주장한다. 컨텍스와 디노이징 역할을 분리함으로써, 이 접근 방식은 30B 규모의 하이브리드 MoE 모델이 가진 광범위한 능력을 보존하는 동시에 확산 모델의 병렬 생성 이점을 활용한다.

저자들은 이러한 품질 회복이 원래 사전 학습에 사용된 데이터의 일부(25T 토큰 대비 약 2.1T 토큰)에 대한 적응을 통해 달성되었음을 강조한다. 이 연구는 대규모 자기회고 모델의 근본적인 품질을 희생하지 않으면서 더 높은 처리량 생성을 위한 실행 가능한 경로를 제시함으로써, 언어 모델의 새로운 파레토 최적선(Pareto frontier)을 구축한다. 연구를 촉진하기 위해 코드와 모델 가중치가 공개되었다.