[Paper] MoE-DiffuSeq: Sparse Attention 및 Mixture of Experts를 활용한 Long-Document Diffusion 모델 강화

Source: arXiv - 2512.20604v1

개요

이 논문은 MoE‑DiffuSeq라는 새로운 프레임워크를 소개한다. 이 프레임워크는 mixture‑of‑experts (MoE) 라우팅과 맞춤형 sparse‑attention 메커니즘을 결합하여 확산 기반 텍스트 생성이 매우 긴 문서에서도 실현 가능하도록 만든다. 기존 확산 모델(예: DiffuSeq)의 악명 높은 메모리 및 연산 병목 현상을 해결함으로써, 저자들은 과학 논문 초안 작성, 코드베이스 합성, 다중 턴 대화 봇과 같은 실제 사용 사례에 기술을 한층 가깝게 만든다.

주요 기여

• Sparse‑attention diffusion backbone: 시퀀스 길이에 대해 대략 선형적으로 확장되는 맞춤형 어텐션 스킴으로, GPU 메모리 사용량을 크게 줄입니다.
• Mixture‑of‑Experts routing: 토큰당 소수의 전문가 서브‑네트워크만 동적으로 활성화하여 FLOPs를 추가로 감소시키면서 모델 용량을 유지합니다.
• Soft absorbing state: 확산 디노이징 단계에 통합되어 수렴 속도를 높이고 토큰‑레벨 재구성 정확도를 향상시킵니다.
• Comprehensive benchmarking: 과학 초록, 코드 저장소, 대화 로그와 같은 장문 데이터셋에 대한 실험 결과, 훈련/샘플링 속도가 2–3배 빠르고 BLEU, ROUGE, 인간 평가 일관성에서 측정 가능한 향상을 보입니다.
• Open‑source implementation: 저자들은 코드와 사전 학습된 체크포인트를 공개하여 개발자들이 확산‑기반 생성 모델을 실험하기 위한 장벽을 낮춥니다.

Methodology

1. Base diffusion model – DiffuSeq에서 시작하며, 텍스트 생성을 역확산 과정으로 간주합니다: 잡음이 섞인 토큰 시퀀스를 점진적으로 디노이즈하여 읽을 수 있는 텍스트로 복원합니다.
2. Sparse attention layer – 고전적인 전체‑셀프‑어텐션 (O(N²) 비용) 대신, 모델은 슬라이딩 윈도우와 학습된 “전역” 토큰 집합에 대해서만 어텐션을 계산합니다. 이는 레이어당 복잡도를 O(N·k) 로 낮추며 여기서 k ≪ N 입니다.
3. Mixture‑of‑Experts (MoE) routing – 각 트랜스포머 블록은 여러 전문가 피드‑포워드 네트워크를 포함합니다. 가벼운 게이팅 네트워크가 각 토큰에 대해 상위‑k 전문가를 선택하고, 순방향/역방향 연산 시 해당 전문가만 활성화합니다. 이를 통해 높은 용량의 모델을 얻으면서도 계산량이 비례적으로 증가하지 않습니다.
4. Soft absorbing state – 확산 단계 동안 작은 확률 질량이 “흡수”되어 안정된 상태로 들어가게 허용함으로써, 수렴에 필요한 확산 타임스텝 수를 효과적으로 줄입니다.
5. Training & sampling – 모델은 표준 변분 확산 손실로 학습되지만, 추가적인 MoE 정규화(로드 밸런싱 손실)와 희소 어텐션 마스크가 포함됩니다. 샘플링은 기존의 역확산 스케줄을 따르며, 이제 흡수 상태에 의해 가속됩니다.

결과 및 발견

• 효율성: 2k 토큰 시퀀스에서 학습 속도가 약 2.5× 빨라지고 샘플링 지연시간이 50 % 이상 감소했습니다.
• 품질: 자동 평가지표와 인간 평가 모두에서 일관된 향상이 이루어졌으며, 특히 긴 구간에 걸친 전역 일관성 유지가 크게 개선되었습니다.
• 확장성: MoE 희소성 덕분에 파라미터를 1 B에서 4 B로 확장해도 메모리 사용량이 크게 증가하지 않았습니다.

실용적 시사점

• Developer tooling: IDE 플러그인이 방대한 문서나 코드 스니펫을 자동 생성할 때, 이제 확산 모델을 사용해도 지연 시간이 크게 문제되지 않는다.
• Content platforms: 뉴스룸과 과학 출판사는 MoE‑DiffuSeq를 활용해 긴 기사 초안을 작성할 수 있으며, 구조를 유지한 빠른 초안을 얻을 수 있다.
• Conversational AI: 다중 턴, 컨텍스트가 풍부한 대화를 처리하는 고객 지원 봇은 수백 턴에 걸쳐 일관성을 유지하면서 GPU 비용이 급증하지 않는다.
• Edge‑friendly deployment: 토큰당 활성화되는 전문가가 일부만 있기 때문에, 추론을 여러 GPU에 분산하거나 특수 가속기 클러스터에 나눠 실행할 수 있어 대규모 생성이 비용 효율적이다.
• Open‑source ecosystem: 공개된 코드는 Hugging Face Transformers와 통합되어, 개발자가 MoE‑DiffuSeq를 기존 파이프라인에 최소한의 마찰로 연결할 수 있다.

Limitations & Future Work

• Expert imbalance: 로드‑밸런싱 손실을 적용했음에도 불구하고, 특히 매우 동질적인 코퍼스에서는 일부 전문가가 충분히 활용되지 않을 수 있습니다.
• Sparse‑attention hyper‑tuning: 적절한 윈도우 크기와 글로벌 토큰 수를 선택하는 데 여전히 데이터셋별 실험이 필요합니다.
• Diffusion step count: 소프트 흡수 상태가 단계 수를 줄여 주지만, 모델은 여전히 수십 번의 역확산 반복을 필요로 하며, 이는 초저지연 애플리케이션에 장애가 될 수 있습니다.
• Future directions proposed by the authors include:
  1. 토큰당 활성화되는 전문가 수를 동적으로 학습하는 적응형 게이팅.
  2. 사실 정확성을 더욱 향상시키기 위한 검색‑보강 생성 통합.
  3. 추론 단계 수를 더욱 줄이기 위한 하이브리드 자동회귀‑확산 스케줄 탐색.

저자

• Alexandros Christoforos
• Chadbourne Davis

논문 정보

• arXiv ID: 2512.20604v1
• 분류: cs.CL
• 출판일: 2025년 12월 23일
• PDF: PDF 다운로드