LLaDA-o: 효과적이고 길이 적응형인 옴니 확산 모델

기존의 멀티모달 통합 모델들은 텍스트(이산형)와 이미지(연속형)의 서로 다른 특성으로 인해 학습이 불안정하거나 성능이 제한적이었다. LLaDA-o는 Mixture of Diffusion(MoD) 구조를 통해 각 매체에 최적화된 확산 방식을 사용하면서도 단일 신경망으로 통합하여, 고품질 이미지 생성과 정교한 시각적 이해를 동시에 가능하게 한다.

Transformer 기반 모델에서 텍스트는 단어 단위의 이산적(Discrete) 토큰으로 처리되지만, 이미지는 픽셀이나 잠재 공간의 연속적(Continuous) 값으로 표현된다. 기존의 통합 확산 모델은 이 두 가지 이질적인 데이터를 하나의 최적화 목표로 학습시키려다 보니 그래디언트 간섭이 발생하고 성능이 저하되는 한계가 있었다.

LLaDA-o는 Mixture of Diffusion(MoD) 개념을 도입하여 이 문제를 해결한다. 텍스트 이해에는 마스크된 토큰을 예측하는 방식을, 이미지 생성에는 노이즈를 제거하는 연속적 확산 방식을 각각 전담하는 '전문가' 구조를 채택했다. 이는 마치 한 명의 요리사가 한 번에 여러 요리를 섞어 만드는 대신, 각 요리에 맞는 전용 도구를 사용하면서 주방(공통 어텐션 백본)은 공유하는 것과 같다.

또한, Intra-Modality Bidirectional Attention을 통해 이미 처리된 조건부 데이터(이미지나 프롬프트)의 계산 결과를 재사용(KV Caching)할 수 있게 했다. 이를 통해 기존의 전체 양방향 어텐션 방식보다 5.9배 빠른 속도로 결과를 생성하며, 가변적인 길이의 응답도 유연하게 처리할 수 있게 되었다.

Mixture of Diffusion (MoD) 프레임워크는 이해 전문가(Understanding Expert)와 생성 전문가(Generation Expert)로 구성된다. 이해 전문가는 SigLIP 비전 인코더와 MLP 프로젝터를 거친 시각 토큰을 텍스트 토큰과 함께 마스크 확산 모델(MDM) 방식으로 처리한다. 생성 전문가는 VAE를 통해 이미지를 잠재 토큰으로 변환한 후 연속 확산 트랜스포머(DiT) 구조를 사용하여 노이즈로부터 이미지를 복원한다.

Intra-Modality Bidirectional Attention 메커니즘은 입력 시퀀스를 모달리티 블록으로 분할한다. 각 블록 내부에서는 모든 토큰이 서로를 참조하는 양방향 어텐션을 수행하지만, 블록 사이에는 이전 블록만 참조할 수 있는 인과적 제약을 둔다. [입력 시퀀스 → 모달리티별 블록 분할 → 블록 내 양방향/블록 간 인과적 어텐션 연산 → KV 캐시 재사용 가능 구조] 순으로 연산이 진행되어 추론 시 중복 계산을 제거한다.

Adaptive Length Augmentation 전략은 학습 데이터의 응답 부분에 확률적으로 [EOS] 토큰을 추가하거나(pext) 임의의 지점에서 텍스트를 자르는(ptrunc) 방식이다. [원본 응답 → 확률적 절단 또는 토큰 추가 → 변형된 응답으로 학습 → 모델이 스스로 종료 시점을 판단] 하도록 유도하여, 아키텍처 변경 없이도 다양한 길이의 문장을 생성할 수 있게 한다.

멀티모달 이해 벤치마크에서 LLaDA-o는 MathVista(66.1), ChartQA(87.9), DocVQA(91.5) 등 10개 지표에서 기존 확산 기반 통합 모델인 LaViDa-O나 Lumina-DiMOO를 능가하는 성능을 보였다. 특히 수학적 추론과 차트/문서 이해 능력에서 큰 폭의 개선을 달성했다.

텍스트-이미지 생성 성능을 측정하는 DPG-Bench에서 87.04점을 기록하여, 통합 모델 중 가장 높은 점수를 얻었으며 전용 생성 모델인 SD3-Medium(84.08)보다도 우수한 결과를 나타냈다. GenEval에서도 두 개 이상의 객체 생성 및 색상 결합 작업에서 강점을 보였다.

효율성 분석 결과, 제안된 어텐션 기법을 통해 기존 LLaDA-V 대비 약 5.9배의 추론 속도 향상을 확인했다. 또한 신뢰도 임계값(Confidence Threshold)을 0.9로 설정했을 때 정확도와 처리량 사이의 최적의 균형을 달성함을 입증했다.

LLaDA-o는 LLaDA-8B-Instruct를 언어 모델 백본으로, SigLIP을 비전 인코더로 사용한다. 이해 전문가는 이산 상태 공간에서의 마스크 예측 손실 함수(Lund)를 최적화하며, 생성 전문가는 Rectified Flow 또는 Flow Matching 기반의 연속 공간 손실 함수(Lgen)를 사용한다.

핵심 차별점은 Mixture of Diffusion(MoD)을 통한 그래디언트 분리다. 이해 작업과 생성 작업의 목적 함수가 다르기 때문에, 이를 별도의 전문가 레이어(MLP, QKV 프로젝션 등)로 분리하여 학습 안정성을 확보했다. 반면, 어텐션 백본은 공유하여 모달리티 간 상호작용을 유지한다.

Intra-Modality Bidirectional Attention은 KV 캐싱을 확산 모델에 도입하기 위한 설계다. 조건부 블록(이미지/프롬프트)을 고정된 프리픽스로 취급하여 한 번만 계산하고, 생성되는 토큰 블록만 반복 계산함으로써 연산 복잡도를 획기적으로 낮췄다.

학습은 3단계로 진행된다. 1단계는 대규모 이미지 이해 및 생성 기초 학습(512 해상도), 2단계는 고해상도(1024) 및 멀티모달 추론 강화, 3단계는 Adaptive Length Augmentation을 적용한 가변 길이 생성 미세 조정 단계다.

LLaDA-o는 강력한 성능을 보이지만, 여전히 텍스트 전용 모델인 BAGEL(Qwen2.5-7B 기반)에 비해 일부 언어 이해 지표에서 다소 낮은 수치를 기록했다. 이는 사용된 언어 모델 백본의 사전 학습 데이터 규모 차이에서 기인한 것으로 분석된다.