기술 요약: 이산 확산 모델에서의 글로벌에서 요소별 전문가 구성으로의 전환

문제 정의

이산 확산 모델(Discrete diffusion models)은 여러 사전 학습된 "전문가(expert)" 모델을 결합하여 개별 학습 분포를 넘어 일반화하는 **구성적 생성(compositional generation)**을 통해 복잡한 추론 작업에서 강력한 프레임워크로 부상했습니다. 최근의 이론적 발전(예: SuperDiff, RNE, Feynman-Kac Correctors)은 시간 의존적 혼합 가중치를 도입하여 구성된 역동성을 타겟 분포와 정렬함으로써 구성적 샘플링을 개선했습니다.

그러나 이러한 최첨단 방법들은 근본적인 한계를 공유합니다. 즉, 이들은 **샘플 단위(per-sample basis)**로 작동하며, 각 타임스텝마다 각 전문가에게 단일한 전역 스칼라 가중치를 할당합니다. 이러한 "단일 구조적(monolithic)" 가정은 모든 전문가가 전체 상태 공간에 대해 동일하게 지식을 갖추고 있다고 암묵적으로 간주합니다. 경계(boundaries)를 전문으로 하는 전문가와 질감(textures)을 다루는 또 다른 전문가가 공존하는 공간 추론과 같은 복잡한 생성 작업에서, 이러한 전역 가중치는 타협을 강요합니다. 이는 특정 영역에서 올바른 전문가의 신호를 희석시키며, 서로 다른 전문가들의 공간적 및 기능적 독립성을 포착하지 못하게 만듭니다.

방법론: FactorDiff

저자들은 생성 과정을 샘플 수준에서 요소(factor) 수준(예: 공간적 픽셀 또는 그리드 셀)으로 분해하는 **요소별 구성 프레임워크(factor-wise composition framework)**인 FactorDiff를 제안합니다. 전체 생성된 상태에 단일 가중치를 적용하는 대신, FactorDiff는 각 요소를 가장 관련 있는 전문가에게 동적으로 라우팅합니다.

핵심 구성 요소

1. 요소별 라우팅(Factor-Wise Routing): 상태 공간은 사용자 정의 요소로 분해됩니다. 샘플링 과정 동안, 라우팅 맵은 각 요소(예: 그리드의 픽셀)를 특정 전문가에게 할당합니다. 이를 통해 하나의 생성된 상태 내에서 이질적인 전문가 기여를 허용합니다.
2. 신뢰도 기반 선택(Confidence-Based Selection): 이 프레임워크는 추론 시 각 위치에서 가장 높은 신뢰도를 가진 전문가를 동적으로 선택함으로써 라우팅을 구현합니다. 변수 $v$에 대한 라우팅 변수 $z_v$는 다음과 같이 결정됩니다:
   $$z_v = \arg\max_{k} C(p_k(x_v | x_t))$$
   여기서 $C$는 신뢰도 지표(상위 두 예측 확률 사이의 마진)입니다. 구성된 역방향 단계(reverse step)는 해당 특정 위치에 대해 선택된 전문가의 예측을 선택합니다.
3. 네이티브 2D 마스크 확산 아키텍처(Native 2D Masked Diffusion Architecture): 이를 지원하기 위해 저자들은 그리드 토폴로지를 보존하는 2D 마스크 확산 아키텍처를 도입합니다. 그리드를 1D 시퀀스로 평탄화하여 토폴로지 불일치를 유발하는 기존 방식과 달리, 이 접근 방식은 그리드를 격자 그래프(lattice graph)로 취급합니다.
   • 전문가(Experts): 세 가지 상보적인 모델 유형이 학습됩니다: 전체 모델(Full model) (모든 색상), 점유 모델(Occupancy model) (배경 대 전경), 그리고 색상 모델(Color model) (전경 색상만).
   • 백본(Backbone): 디노이징 네트워크는 공간적 국부성을 보장하기 위해 학습 가능한 2D 위치 임베딩을 사용하는 반복적 정제(iterative refinement) 아키텍처(컨볼루션 또는 트랜스포머 기반)를 사용합니다.

이론적 정당성

기존 방법들은 타겟 분포를 전문가의 가중 곱($p(x) \propto \prod p_m(x)^{\alpha_m}$)으로 정의하는 경우가 많습니다. 저자들은 상보적인 전문가들의 경우, 전역 지수 $\alpha_m$이 잘못 지정되었다고 주장합니다. 왜냐하면 이는 전문가의 예측(심지어 그들의 도메인 밖의 오류가 있는 예측까지도)이 전체 상태에 영향을 미치도록 강제하기 때문입니다. FactorDiff는 이를 요소별 곱(factor-wise product)으로 대체합니다:
$$p_{comp}(x) \propto \prod_{f \in \mathcal{F}} \prod_{m=1}^M (p_m(x_f))^{\alpha_{m,f}}$$
이 공식화는 요소에 따라 전문성이 달라질 수 있게 합니다. 하드 라우팅(hard-routing)의 경우, $\alpha_{m,f}$는 원-핫(one-hot) 형태가 되어, 추가적인 정규화나 중요도 교정 없이도 유효한 조건부들의 곱이 유효한 결합 전이(joint transition)를 정의하므로, 요소당 단일 전문가를 선택합니다.

주요 기여

1. FactorDiff 프레임워크: 전문가 조합을 요소에 대한 라우팅으로 규정하는 새로운 구성 방법이며, 특히 그리드 구조의 확산 모델을 위한 위치 수준의 전문가 선택으로 구체화되었습니다.
2. 2D 마스크 확산 아키텍처: 그리드 토폴로지를 유지하는 네이티브 아키텍처로, 추론 시 구성될 수 있는 상보적인 전문가(점유, 색상, 전체)를 학습할 수 있게 합니다.
3. ARC-AGI에 대한 실증적 검증: ARC-AGI 벤치마크에 대한 광범위한 평가를 통해, 단순한 요소별 라우팅이 복잡한 전역 스칼라 가중치 방식보다 일관되게 우수한 성능을 보임을 입증했습니다. 특히 논리적 일관성과 공간적 분리(spatial disentanglement)가 필요한 작업에서 그러했습니다.

실험 결과

저자들은 IID(in-distribution) 평가를 위해 RE-ARC 생성기를 사용하여 120개의 ARC-AGI 태스크 서브셋에 대해 FactorDiff를 평가했습니다.

• 특화 영역(Specialization Regime, RQ1): 고도로 특화된 전문가들을 결합할 때(예: 점유 + 전체, 또는 색상 + 점유), FactorDiff는 샘플당 방식의 베이스라인(SuperDiff, RNE, FKC)을 크게 앞질렀습니다.
  • 색상 + 점유(Color + Occupancy): 개별 전문가들은 전체 태스크에서 거의 0에 가까운 점수를 기록했습니다(정확도 각각 3.3% 및 0.2%). FactorDiff는 이들을 결합하여 **90.5%**의 정확도를 달-성했으며, 이는 가장 우수한 샘플당 베이스라인(78.4%)보다 훨씬 높은 수치입니다.
  • 점유 + 전체(Occupancy + Full): FactorDiff는 **94.0%**의 정확도를 달성하여, 단일 최고 전문가(93.0%) 및 모든 샘플당 구성 방식보다 뛰어난 성능을 보였습니다.
• 제너럴리스트 영역(Generalist Regime, RQ2): 두 개의 강력한 제너럴리스트 모델(Full + Full)을 결합할 때, FactorDiff는 성능 저하 없이 가장 우수한 샘형당 베이스라인(95.2%)과 대등한 수준을 유지하며 "안전함"을 보여주었습니다. 이는 전문가들이 광범위하게 동의할 때 공간적으로 변화하는 선택이 불안정성을 유발하지 않음을 나타냅니다.
• 절제 연구(Ablation): 동일한 목적 함수를 가지지만 아키텍처가 다른 두 개의 풀 모델(Trans-IR vs. Conv-IR)을 결로할 때도, FactorDiff는 개별 모델보다 높은 정확도(94.8%)를 달성하여 두 모델의 올바른 영역을 성공적으로 결합해 냈습니다.

의의 및 주장

이 논문은 FactorDiff가 단일 확산 궤적 내에서 임의의 전문가들을 결합하는 단순하고 효과적인 메커니즘을 제공한다고 주장합니다. 이 연구의 주요 의의는 전문가들이 공간적 또는 기능적으로 특화된 도메인을 가진 작업에서는 샘플당 전역 가중치가 불충분하다는 것을 입증했다는 데 있습니다. 요소별 구성을 통해 이 방법은 전역 가중치가 표현할 수 없는 상보적 전문성을 포착하며, 엄격한 논리적 일관성이 요구되는 작업에서 우수한 일반화 성능을 달성합니다.

저자들은 현재 구현이 (신뢰도가 정답과 상관관계가 있다고 가정하는) 신뢰도 기반 라우팅에 의존하고 있지만, 이 프레임워크가 학습된 또는 태스크 조건부 라우팅으로 자연스럽게 확장될 수 있음을 언급했습니다. 이들은 이 연구를 현재의 최첨단 방법들이 가진 "단일 구조적" 가정을 넘어, 이산 확산에서 더 세밀한 구성으로 나아가는 단계로 제시합니다.