기술 요약: Sparse-LaViDa

문제 정의

마스크된 이산 확산 모델(Masked Discrete Diffusion Models, MDMs)은 이미지 이해, 생성 및 편집을 포함한 멀티모달 작업을 위한 강력한 통합 프레임워크로 부상했습니다. 자기회귀(Autoregressive, AR) 모델과 달리, MDM은 양방향 컨텍스트와 병렬 디코딩을 활용하여 이미지 인페인팅(inpainting) 및 텍스트 인필링(text infilling)과 같은 작업에서 이점을 제공합니다. 그러나 기존 MDM은 다음과 같은 심각한 추론 비효효율성을 겪고 있습니다:

1. 중복 프로세싱: 표준 MDM은 매 샘플링 단계마다 실질적인 정보를 담고 있지 않은 방대한 양의 중복된 마스크 토큰을 포함하여 전체 시퀀스를 처리해야 합니다.
2. KV-캐시 지원 부족: 인과적 어텐션(causal attention)이 아닌 전체 어텐션 메커니즘에 의존하기 때문에, 표준 MDM은 추론 시 KV-캐시 가속을 활용할 수 없습니다.
3. 기존 가속기의 한계: Block Diffusion과 같은 기존의 MDM 가속 시도는 좌-우 블록 인과적 순서(left-to-right block-causal order)를 강요합니다. 이는 KV-캐싱을 가능하게 하지만, 이미지 생성 및 편집에 필수적인 양방향 컨텍스트를 희생시키므로 인페인팅과 같은 작업을 어렵게 하거나 불가능하게 만듭니다.

방법론

저자들은 생성 품질과 양방향 컨텍스트를 보존하면서 불필요한 마스크 토큰을 동적으로 절단(truncating)하도록 설계된 새로운 모델링 프레임워크인 Sparse-LaViDa를 제안합니다. 이 방법은 세 가지 핵심 혁신을 도입합니다:

1. 희소 파라미터화 (Sparse Parameterization)

입력 시퀀스의 모든 마스크 토큰을 실체화하는 대신, Sparse-LaViDa는 다음만을 사용하여 부분적으로 마스크된 시퀀스를 표현합니다:

• 프롬프트 토큰 및 이전에 디코딩된 토큰 (KV 캐시에 저장됨).
• 현재 단계에서 디코딩을 위해 선택된 마스크된 토큰의 서브셋.
• 레지스터 토큰 (Register Tokens): 절단된(제거된) 마스크 토큰의 압축된 표현 역할을 하기 위해 도입된 특수 토큰입니다. 이 레지스터들은 절단으로 인해 손실된 모델링 능력을 복구하는 데 도움을 줍니다. 최종 설계에서는 64개의 레지스터 토큰이 사용되며, 시퀀스의 끝에 배치됩니다.

2. 스텝-인과적 어텐션 마스크 (Step-Causal Attention Mask)

학습과 추론 사이의 일관성을 보장하기 위해, 저자들은 추론 중에 관찰되는 점진적인 토큰 캐싱 동작을 시뮬레이션하는 특화된 어텐션 마스크를 설계했습니다:

• 학습: 시퀀스는 블록으로 분할됩니다. 프롬프트 토큰은 블록 0에 할당됩니다. 클린(clean) 토큰은 블록 $1 \dots M$에 할당되며, 마스크된 토큰은 블록 $M+1 \dots M+N$에 할당됩니다.
• 어텐션 규칙: 특정 블록의 토큰은 자신의 블록 또는 이전 블록의 토큰에 어텐션을 수행할 수 있습니다. 결정적으로, 특정 블록의 마스크된 토큰은 다른 마스크 블록의 마스크된 토큰에는 어텐션을 수행할 수 없지만, 모든 프롬프트 및 클린 토큰에는 어텐션을 수행할 수 있습니다.
• 레지스터 통합: 레지스터 토큰은 각 마스크 블록에 추가되며 동일한 어텐션 규칙을 따릅니다.
  이 설계는 모델이 토큰의 일부만을 관찰하는 추론 과정을 충실히 모사하면서 효율적인 병렬 학습을 가능하게 합니다.

3. 샘플링 전략

프레임워크는 작업에 따라 서로 다른 언마스킹(unmasking) 전략을 지원합니다:

• 시각적 생성 (Text-to-Image, Editing): 높은 생성 품질을 유지하기 위해 신뢰도 점수(confidence scores)에 의존하지 않고 미리 생성된 2D 언마스킹 순서(예: 층화 무작위 샘플링)를 활용합니다.
• 텍스트 생성/이해: 토큰 블록을 왼쪽에서 오른쪽으로 샘플링하되, 블록 내부에서는 동적인 언마스킹을 수행하는 세미-자기회귀(semi-auto-regressive) 전략을 채택합니다. Block Diffusion과 달리, 이 접근 방식은 블록 간의 양방향 컨텍스트를 유지합니다.

주요 기여

1. 효율적인 파라미터화: Sparse-LaViDa는 좌-우 디코딩 순서를 가정하거나 양방향 컨텍스트를 희생하지 않고도 KV-캐싱과 임의의 토큰 절단을 모두 지원하는 최초의 방법입니다.
2. 레지스터 토큰: 압축된 표현으로서의 레지스터 토큰 도입은 공격적인 토큰 제거로 인한 성능 저하를 효과적으로 완화합니다.
3. 학습-추론 일관성: 스텝-인과적 어텐션 마스크는 학습과 추론 사이의 간극을 메워, 모델이 증류(distillation)나 휴리스틱한 수정 없이 희소 파라미터화를 직접 학습할 수 있게 합니다.
4. 통합 프레임워크: 이 접근 방식은 최첨단 통합 MDM인 LaViDa-O를 기반으로 구축되어, 단일 가속화된 프레임워크 내에서 다양한 작업(생성, 편집, 추론)을 처리할 수 있습니다.

실험 결과

저자들은 광범위한 멀티모달 벤치마크에서 Sparse-LaViDa를 베이스라인인 LaViDa-O 및 다른 최첨단 모델들과 비교 평가했습니다.

• Text-to-Image 생성: GenEval 벤치마크에서 Sparse-LaViDa는 LaViDa-O와 대등한 성능(전체 점수 0.78 vs 0.77)을 달성하면서 추론 지연 시간을 21.27초에서 10.86초로 줄여 1.96배의 속도 향상을 기록했습니다. DPG-bench 및 MJHQ-30k에서는 동일한 데이터 서브셋으로 학습했을 때 지각적 지표(예: FID가 8.11에서 7.63으로 개선)를 유지하거나 약간 향상시켰습니다.
• 이미지 편집: ImgEdit 벤치마크에서 Sparse-LaViDa는 전체 점수를 3.71에서 3.79로 높이는 동시에 지연 시간을 63.98초에서 22.55초로 줄여 2.83배의 속도 향상을 달ر성했습니다.
• 시각적 수학 추론: MathVista 벤치마크에서 모델은 베이스라인과 대등한 정확도(56.7 vs 56.9)를 유지하면서 2.80배의 속도 향상(3.72초 vs 10.41초)을 달성했습니다.
• 절제 연구 (Ablation Studies):
  • 레지스터 토큰을 제거하면 미세한 정렬(DPG) 및 이미지 품질 지표(FID, HPS v3)에서 성능 저하가 발생하여, 저수준 시각적 세부 사항을 보존하는 데 있어 이들의 역할을 확인했습니다.
  • 스텝-인과적 어텐션 마스크를 제거하거나 파인튜닝을 건너뛰면 상당한 성능 저하가 발생하여, 제안된 학습 파이프라인의 필요성을 강조했습니다.

의의 및 주장

본 논문은 Sparse-LaVi Da가 사후 증류 단계나 휴리스틱한 무학습 수정 없이도 확장 가능한 학습 및 추론을 가능하게 하는 근본적으로 효율적인 파라미터화를 제공한다고 주장합니다.

저자들은 다음을 강조합니다:

• Sparse-LaViDa는 밀집(dense) 베이스라인과 대등한 생성 품질을 유지하면서 다양한 작업에서 상당한 속도 향상(최대 2.84배)을 달성합니다.
• 기존의 가속 방법(예: Block Diffusion)과 달리, 양방향 컨텍스트를 보존하여 비인과적 의존성에 의존하는 인페인팅, 아웃페인팅, 제약 조건이 있는 캡셔닝과 같은 복잡한 작업을 가능하게 합니다.
• 이 방법은 학습-추론가 일관된 접근 방식을 제공하여, 휴리스틱 기반 가속 기술에서 흔히 나타나는 예측 불가능한 성능 저하를 피합니다.

이 연구는 Sparse-LaViDa를 MDM의 이론적 장점과 배포에 필요한 계산 효율성 사이의 균역을 맞추어, 통합 멀티모달 확산 모델을 실시간 애플리케이션에 실용적으로 만들 수 있는 유망한 경로로 자리매김합니다.