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

균일 상태 확산 (uniform-state diffusion) 프레임워크가 이산 확산 (discrete diffusion)을 위한 더 강력한 패러다임일까요? 최근 연구들은 그럴 가능성이 있다고 나타냅니다. 예측기-교정기 (predictor-corrector) 샘플러와 결합했을 때, 균일 상태 확산 모델 (USDMs)은 마스크 확산 모델 (MDMs)보다 더 높은 품질의 샘플을 생성하며, USDMs는 더 높은 퍼플렉시티 (perplexity)를 보임에도 불구하고 다운스트림 태스크 (downstream tasks)에서 MDMs와 대등하거나 이를 능가합니다. 두 가지 문제가 여전히 해결되지 않은 채 남아 있습니다. 첫째, 기존 연구들은 틀릴 가능성이 가장 높은 토큰을 타겟팅하는 대신, 무작위 위치에 노이즈를 재주입하는 정보가 없는 교정기 (un-informed correctors)를 사용하여 균일 확산과 마스크 확산을 비교합니다. 둘째, 이전 연구들은 전체 시퀀스 확산 모델 (full-sequence diffusion models)을 비교했기 때문에, 토큰이 블록 단위로 생성될 때도 동일한 결론이 유지되는지 알 수 없습니다. 이러한 문제를 해결하기 위해, 우리는 마스크 확산과 균일 확산을 모두 사용하여 구현한 블록 단위 시퀀스 모델인 BlockGen을 소개합니다. BlockGen은 혼합된 블록 크기로 학습하며, 그 가능도 (likelihood)는 고정된 블록 크기를 가진 모델보다 자기회귀 (AR) 모델과 순수 확산 (pure diffusion) 모델 사이를 더 미세하게 보간합니다. BlockGen은 보조 검증기 (auxiliary verifier) 없이도 AR과 확산 예측을 결합하여 가능성이 낮은 토큰을 재생성하는 AR 정보 기반 예측기-교정기 (AR-informed predictor-corrector, ARPC) 샘플링을 가능하게 합니다. 조상 샘플링 (ancestral sampling) 하에서, 블록 단위 설정 시 균일 확산은 특히 적은 단계 (few-step) 영역에서 마스크 확산보다 우수한 성능을 보입니다. ARPC 하에서는 그 격차가 줄어들며 높은 NFE (Number of Function Evaluations)에서 역전됩니다. GSM8K 데이터셋에서 블록 크기 $16$을 사용할 때, MDMs는 USDMs보다 약간 더 높은 정확도에 도달하며, OpenWebText의 생성 퍼플렉시티 (Generative Perplexity)에서도 유사한 경향이 관찰됩니다. 우리의 코드는 https://github.com/jdeschena/blockgen 에서 확인할 수 있습니다.