기술 요약: BlockGen: 하이브리드 샘플러를 이용한 유연한 블록 단위 시퀀스 모델링

문제 정의

자기회귀(Autoregressive, AR) 트랜스포머는 시퀀스 모델링을 지배하고 있으나, 순차적 생성의 한계(처리량/지연 시간)와 양방향 문맥이 필요한 추론 작업에 제약을 주는 인과적 어텐션(causal attention) 문제를 겪고 있다. 이산 확산 모델(Discrete diffusion models)은 노이즈가 섞인 시퀀스를 반복적으로 정제함으로써 병렬 업데이트를 허용하는 대안을 제시한다. 그러나 최근 **마스크 확산 모델(Masked Diffusion Models, MDMs)**과 균일 상태 확산 모델(Uniform-State Diffusion Models, USDMs) 간의 비교 연구는 샘플링 방식에 따라 상충되는 통찰을 보여준다:

1. 조상 샘플링(Ancestral Sampling): USDM은 퍼플렉시티(perplexity) 측면에서는 열세임에도 불구하고 다운스트림 작업에서 MDM보다 우수한 성능을 보이는 경우가 많은데, 이는 USDM이 오류를 스스로 수정할 수 있는 능력(토큰이 모든 상태 사이를 전이할 수 있음) 덕분이다. 반면 MDM은 마스킹이 해제된 토큰을 고정한다.
2. 블록 단위 생성(Blockwise Generation): 기존 연구는 주로 전체 시퀀스 확산을 비교해 왔다. USDM의 이점이 토큰을 블록 단위로 생성하는(준-자기회귀 방식) 설정에서도 유지되는지는 여전히 불분명하며, 이 설정은 KV 캐싱과 빠른 추론을 가능하게 한다.
3. 정보 기반 교정기(Informed Correctors): 기존의 비교는 무작위 위치에 노이즈를 재주입하는 "비정보형(un-informed)" 예측자-교정자(predictor-corrector, PC) 샘플러를 주로 사용한다. 보조 검증 모델 없이 "정보 기반" 교정기(잠재적 오류를 타겟팅하는 방식)가 MDM과 USDM의 순위에 어떤 영향을 미치는지에 대해서는 알려진 바 없다.

방법론: BlockGen

저자들은 다양한 블록 크기의 혼합(mixture) 위에서 학습되며, 블록 내에서 마스크 또는 균일 확산 중 하나로 구현되는 블록 단위 시퀀스 모델인 BlockGen을 소개한다.

1. 혼합 공식화 및 학습

BlockGen은 시퀀스 밀도를 $M$개의 블록 크기 $S = \{s_1, \dots, s_M\}$에 대한 혼합으로 정의한다:
$$p_{\theta}^{\text{BlockGen}}(x) = \sum_{i=1}^M \gamma_i p_{\theta}^{(s_i)}(x)$$
여기서 $p_{\theta}^{(s_i)}$는 블록 크기 $s_i$로 팩터화된 유효한 밀도 함수이다.

• 공유 디노이저(Shared Denoiser): 별도의 모델을 학습시키는 대신, 단일 공유 디노이저 $\hat{x}_\theta$가 블록 크기 $L'$을 입력으로 받아 그에 따라 어텐션 패턴을 조정한다.
• 우도 경계(Likelihood Bounds): 논문은 혼합 밀도에 대한 두 가지 경계를 도출한다. 평가에는 더 타이트한 log-sum-exp 경계를 사용하고, 계산 효율성을 유지하기 위해 학습에는 더 느슨한 혼합 우도 경계(블록 크기에 대한 기대값)를 사용한다.
• 층화 추출(Stratified Sampling): 멀티 GPU 학습 중 분산을 줄이기 위해, 블록 크기는 i.i.d. 샘플링 대신 층화 추출을 통해 선택된다.

2. AR 정보 기반 예측자-교정자 (ARPC)

핵심 기여는 보조 검증 모델 없이 BlockGen의 혼합 구조를 활용하여 정보 기반 교정을 수행하는 ARPC 샘플러이다.

• 메커니즘: 생성 과정 중, 표준 조상 확산 단계 이후 샘플러는 "클린 제안(clean proposal)" 단계를 실행한다. 그 후, 제안된 토큰의 점수를 매기기 위해 동일한 디노이저를 사용하여 AR 모드($L'=1$)로 **두 번째 순전파(forward pass)**를 수행한다.
• 스코어링: AR 로그 우도가 낮은 토큰은 "있을 법하지 않은(unlikely)" 것으로 식별되어 이전 노이즈 레벨로 재노이즈(교정)된다.
• 장점: 이를 통해 모델이 자신의 AR 예측을 사용하여 확산 교정을 가이드할 수 있게 하며, 단일 모델 내에서 AR과 확산 패러다임을 효과적으로 인터리빙(interleaving)한다.

3. 블록 수준 어텐션

BlockGen은 **블록-인과적 어텐션(block-causal attention)**을 사용한다: 블록 내의 토큰들은 양방향으로 어텐션을 수행하지만, 블록 자체는 이전의 모든 블록에 대해 인과적으로 어텐션을 수행한다. 이를 통해 디코딩된 블록 간의 KV 캐싱이 가능해져, 준-자기회귀 생성의 효율성 이점을 보존한다.

주요 결과

1. 수학적 추론 (GSM8K)

TinyGSM으로 학습된 모델을 사용하여 GSM8K에서 평가하였다:

• 조상 샘플링: 블록 단위 생성 설정에서 USDM이 MDM보다 우수하며, 특히 적은 NFE(low NFE) 환경에서 그러하다. 이는 전체 시퀀스 USDM의 이점이 블록 단위 생성으로도 이어진다는 것을 확인해 준다.
• ARPC 샘플링: 고(高) NFE 환경에서는 USDM과 MDM 사이의 성능 격차가 좁혀지거나 역전된다. 16-토큰 블록을 사용한 ARPC 환경에서 MDM이 GSM8K 정확도에서 USDM을 약간 앞선다.
• ARPC vs. 베이스라인: ARPC는 대부분의 NFE 예산에서 조상 샘플링 및 엔트로피 정보 기반 PC(EIPC)보다 일관되게 우수한 성능을 보인다. MDM을 사용한 ARPC는 그리디 디코딩을 사용한 AR 샘플링의 성능에 근접한다.

2. 언어 모델링 (OpenWebText & LM1B)

• 우도(Likelihood): 다양한 블록 크기에 대해 학습하는 것은 고정된 블록 모델에 비해 퍼플렉시티를 크게 개선한다. OpenWebText에서 Masked Diffusion을 사용한 BlockGen은 17.5 Val. PPL을 달성했다(AR의 16.7 및 최적의 고정 블록 BDM의 21.6 대비). 이는 확산 모델과 자기회귀 모델 간의 격차를 좁힌다.
• 샘플 품질:
  • 조상 샘플링 하에서는 USDM이 MDM보다 낮은 생성 퍼플렉시티(Gen. PPL)를 달성한다.
  • ARPC 하에서는 높은 NFE(예: 64 NFE/block)에서 추세가 역전되어, Masked-ARPC가 Uniform-ARPC보다 더 나은 Gen. PPL 프런티어를 달성한다.
  • ARPC는 AR과의 격차를 좁히지만, 실질적인 저-온도(low-temperature) 영역 내에서 AR의 Gen. PPL을 능가하지는 못한다.

의의 및 주장

본 논문은 이산 확산에 관한 두 가지 열린 질문을 해결한다고 주장한다:

1. 블록 단위 유효성: 전체 시퀀스 생성에서 관찰된 균일 상태 확산의 이점은 조상 샘플링 하의 블록 단위 생성에서도 유지된다.
2. 정보 기반 교정: MDM과 USDM의 순위는 샘플링 방식에 따라 달라진다. 단순한 조상 샘플링이나 비정보형 샘플러를 사용할 때는 USDM이 뛰어나지만, **정보 기반 교정기(특히 ARPC)**를 사용하면 이 격차를 좁히고 높은 컴퓨팅 예산(NFE)에서 MDM이 우위를 점하도록 순위를 역전시킬 수 있다.

저자들은 BlockGen이 AR과 순수 확산 사이를 보간하는 유연한 프레임워크를 제공하며, 단일 모델이 블록 단위 샘플러이자 AR 검증자로서 기능할 수 있음을 강조한다. 또한, 실험적 주장이 테스트된 예산 범위 내에서는 견고하지만, 170M 파라미터 백본에 국한된 실험을 수행했으므로 이러한 트레이드오프가 대규모 언어 모델(LLM) 규모에서도 지속될지는 확립하지 않았음을 명시한다. 이 연구는 확산 패러다임(마스크 vs. 균일)의 선택이 보편적으로 우월한 것이 아니라, 샘플링 전략과 컴퓨팅 예산에 따라 달라짐을 시사한다.