Focusing on What Matters: Saliency-Harnessing Accurate Routing for Diffusion MoE의 한국어 번역

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

대형 확산 기반 이미지 생성에서 연산 자원 분배는 이미지의 살리언시를 반영할 때 생성 품질이 크게 향상된다는 점이 본문에서 확인됐다. 기존 diffusion MoE 라우터는 노이즈가 있는 latent에 의존해 중요한 토큰을 식별하지 못했고 이로 인해 계산이 비효율적으로 배분됐다. SharpMoE는 이전 타임스텝의 예측된 clean latent를 라우팅 입력으로 사용해 노이즈 영향을 회피함으로써 적은 후학습으로도 pretrained MoE 모델의 생성 품질을 개선했다.

왜 중요한가

대형 확산 기반 이미지 생성에서 연산 자원 분배는 이미지의 살리언시를 반영할 때 생성 품질이 크게 향상된다는 점이 본문에서 확인됐다. 기존 diffusion MoE 라우터는 노이즈가 있는 latent에 의존해 중요한 토큰을 식별하지 못했고 이로 인해 계산이 비효율적으로 배분됐다. SharpMoE는 이전 타임스텝의 예측된 clean latent를 라우팅 입력으로 사용해 노이즈 영향을 회피함으로써 적은 후학습으로도 pretrained MoE 모델의 생성 품질을 개선했다.

핵심 기여

노이즈 의존적 라우팅 문제 규명

기존 diffusion MoE 프레임워크들이 라우터 입력으로 소음이 남아 있는 latent를 사용함으로써 살리언시를 반영한 전문가 할당에 실패한다는 점을 정량적으로 규명했다. 라플라시안 연산으로 추출한 토큰별 살리언시와 할당된 전문가 수의 분포를 분석해 DiffMoE 등의 방법이 살리언시-민감성이 낮음을 실험적으로 확인했다. 이 진단은 이후 clean latent 기반 라우팅 설계의 필요 근거가 됐다.

Saliency-Harnessing Accurate Routing 설계

각 타임스텝의 라우팅 입력으로 직전 타임스텝에서 모델이 예측한 clean latent x^0를 도입하는 Saliency-Harnessing Router를 제안했다. 이 라우터는 기존의 pretrained 라우터와 병렬로 동작하며 최종 라우팅 점수는 두 라우터 출력의 합으로 계산된다. 새 라우터의 가중치는 0으로 초기화되어 사후 학습 동안 안정적으로 살리언시 정보를 통합했다.

Recursive Full-Trajectory Training 기법

Saliency-Harnessing Router가 필요로 하는 이전 단계의 clean prediction을 얻기 위해 연속된 T-step 롤아웃을 수행하는 Recursive Full-Trajectory Training을 도입했다. 각 롤아웃에서 Flow Matching 손실을 계산하고 다음 단계의 clean prediction을 생성해 다음 타임스텝의 라우팅 입력으로 전달한다. 실험에서는 T=10 설정을 기본으로 사용했고 T를 5에서 20까지 바꿔도 성능이 안정적임이 확인됐다.

Trajectory Routing Loss 도입

전체 롤아웃에서 토큰별 전문가 활성화 누적치 A를 계산하고 라플라시안 기반 살리언시 맵 M과의 정규화 분포 간 KL-divergence를 최소화하는 Trajectory Routing Loss를 제안했다. 손실식은 L_routing = D_KL(softmax(A) || softmax(M)) 로 정의되며 하이퍼파라미터 lambda_routing은 논문 실험에서 0.001로 설정됐다. 이 손실은 전체 생성 궤적에 걸쳐 계산 자원을 살리언시에 비례해 재분배하도록 유도했다.

플러그인형 사후학습으로 광범위한 개선 입증

SharpMoE는 사전학습된 여러 diffusion MoE 백본에 플러그인 방식으로 적용되어 100K 후학습만으로 일관된 품질 향상을 보였고 DiffMoE-L 백본에서 cfg=1.5일 때 FID=3.10, IS=228.88의 결과를 달성했다. 다양한 모델 규모(S, B, L)와 CFG 스케일에서 우수한 성능 개선이 보고됐다. 이 결과는 노이즈 없는 라우팅 신호가 생성 충실도를 높이는 핵심 요인임을 실증했다.

핵심 아이디어 이해하기

기존 확산 기반 MoE 접근은 라우터 입력으로 현재 시점의 노이즈가 포함된 latent x_t를 사용해 전문가를 할당했다. 연속적 denoising 과정의 초중반부에서는 잔여 노이즈가 구조적·텍스처 정보를 가려 라우터가 살리언시를 잘 구분하지 못했고 이로 인해 중요한 토큰에 충분한 계산 자원이 할당되지 않았다. 따라서 라우팅 신호 자체를 노이즈에서 분리해 주는 것이 핵심 문제로 규정됐다. SharpMoE의 해결 원리는 단순하다. 모델이 이전 타임스텝에서 예측한 clean latent xhat_0^{t_{k-1}}는 노이즈가 제거된 추정된 이미지 표현이므로 구조적 경계와 텍스처 밀도를 더 분명히 포착한다. 이 clean prediction을 라우터의 입력으로 사용하면 노이즈에 의해 가려진 살리언시 신호를 회복할 수 있고 그 결과 더 많은 전문가가 높은 살리언시 토큰에 할당되는 경향이 형성된다. 추가로 전체 궤적 관점의 규제가 필요했다. 단일 타임스텝에서의 라우팅 점수는 롤아웃 전반의 누적 계산량을 반영하지 못하므로 Trajectory Routing Loss를 통해 여러 타임스텝에 걸친 전문가 활성화 누적치와 라플라시안 기반 살리언시 맵을 정렬했다. 이로 인해 단기적 잡음에 흔들리지 않는 장기적 자원 배분이 가능해졌다.

방법론

전체 아키텍처는 DiT 백본의 FFN을 SharpMoE 블록으로 교체하는 방식으로 구성됐다. 각 SharpMoE 블록은 pretrained router R_pre와 saliency-harnessing router R_sal을 병렬로 갖고, 최종 라우팅 점수 S는 두 라우터 출력의 합 S = R_pre(x_tk) + R_sal(xhat_0^{t_{k-1}})으로 계산된다. R_sal은 두 계층 MLP와 SiLU 활성화를 사용하며 학습 초기에는 가중치를 0으로 초기화해 사후학습 동안 점진적으로 관여시켰다. 학습은 Recursive Full-Trajectory Training으로 수행됐다. 매 반복에서 연속된 T개의 타임스텝을 샘플링하고 각 타임스텝에서 v_k = v_theta(x_tk, t_k, stop-gradient(xhat_0^{t_{k-1}}))를 예측해 Flow Matching 손실을 계산한다. 각 단계에서 xhat_0^{t_k} = x_tk - t_k * v_k를 도출해 다음 단계의 saliency 라우터 입력으로 전달했다. 첫 스텝은 x_1을 대리 입력으로 사용하며 실험에서는 t1=0.999를 사용했다. 전체 손실은 평균 Flow Matching 손실과 lambda_routing * Trajectory Routing Loss의 가중합으로 정의됐다. Trajectory Routing Loss는 각 토큰 i에 대해 모든 타임스텝과 모든 MoE 레이어의 라우팅 점수를 집계해 얻은 누적 할당치 A_i와 라플라시안 기반 살리언시 맵 M_i의 softmax 간 KL-divergence로 계산됐다. 실험에서는 lambda_routing=0.001, 기본 롤아웃 길이 T=10을 사용했고 T에 대한 민감도는 낮게 관찰됐다.

주요 결과

주요 정량 결과는 ImageNet 256x256 설정에서 500K 단계로 사전학습된 모델에 대해 100K 단계의 사후학습을 적용한 후 평가한 것이다. DiffMoE-L 백본에 SharpMoE를 적용했을 때 cfg=1.5에서 FID=3.10 및 IS=228.88을 기록했고, 다양한 백본(TC-DiT, EC-DiT, DiffMoE)과 규모(S/B/L)에서 일관된 FID 및 IS 개선이 보고됐다. 모든 비교는 Flow Matching Euler 250 스텝과 50,000 샘플로 측정된 FID50K와 IS를 기준으로 했다. 구성 요소별 분석에서 Saliency-Harnessing Routing만 추가해도 cfg=1.5 환경에서 FID가 8.03에서 6.95로 개선됐고 Trajectory Routing Loss를 더하면 6.66으로 추가 개선이 발생했다. 사전학습 단계에 무관하게 SharpMoE는 100K 사후학습만으로도 성능 향상을 재현했고 롤아웃 길이 T를 5에서 20으로 변경해도 성능 편차가 작아 방법의 실용성이 확인됐다.

기술 상세

아키텍처 관점에서 SharpMoE 블록은 기존 DiT의 FFN 자리에 삽입되며 각 블록은 N_E개의 전문가를 보유하는 표준 MoE 구조를 유지한다. 라우팅은 두 개의 라우터 출력을 합쳐 얻으며 Top-K 선택은 기존 MoE와 동일한 방식으로 수행된다. Saliency-Harnessing Router는 입력으로 xhat_0^{t_{k-1}}를 받아 라우팅 점수를 계산하므로 각 학습 반복에서 이전 단계의 clean prediction을 반드시 생성해야 한다. 알고리즘적 기반은 Flow Matching 기반의 Rectified Flow 프레임워크를 사용해 v_theta를 직접 회귀하는 방식에 있다. 각 타임스텝에서 v_k를 예측하면 xhat_0^{t_k} = x_tk - t_k * v_k로 clean prediction을 도출하고 이 값을 다음 타임스텝의 라우팅 입력으로 사용한다. Trajectory Routing Loss는 식 9과 11에 따라 토큰별 누적 할당치 A_i를 합산하고 라플라시안으로 산출한 살리언시 맵 M의 softmax와 KL-divergence를 계산해 추가적인 그레이디언트를 제공한다. 학습 하이퍼파라미터는 논문 실험 표준을 따랐고 optimizer로 AdamW, lr=1e-4, batch size=256, EMA decay=0.9999를 사용했다. Trajectory 구성은 T=10을 기본으로 사용했고 초기 타임스텝 t1은 0.999로 설정해 pure Gaussian noise 초기화로 인한 학습 불안정을 회피했다. R_sal의 가중치는 0으로 초기화해 사후학습 중에만 점진적으로 살리언시 신호가 반영되도록 제어했다.

실무 활용

SharpMoE는 사전학습된 diffusion MoE 모델에 플러그인 방식으로 적용하는 사후학습 프레임워크로 설계돼 실무 도입 장벽이 낮다. 적은 수의 사후학습 단계(논문에서 100K)를 통해 품질 향상을 얻을 수 있어 기존 대형 모델을 다시 학습시키는 비용을 크게 줄인다. 다만 공개된 코드 저장소는 제공되지 않아 구현은 논문의 알고리즘과 의사코드를 참조해 재현해야 한다.

사전학습된 diffusion MoE 기반 생성 모델의 품질을 추가 사후학습으로 개선해야 하는 연구 및 엔지니어링 작업
리소스가 제한된 환경에서 이미지의 중요한 영역에 계산을 집중해 생성 충실도를 높이고자 하는 응용
이미지 생성 모델의 라우터 동작을 진단하고 살리언시 기반 라우팅 분배를 도입하려는 모델 개발 파이프라인

코드 공개 여부: 미확인

키워드

Mixture-of-Experts(Mixture-of-Experts)Diffusion Models(확산 모델)DiT(Diffusion Transformer)Trajectory Routing Loss(궤적 라우팅 손실)Saliency-Harnessing Router(살리언시 라우터)

추가 이미지 분석

Chart
이 Figure는 타임스텝별 분석에서 DiffMoE가 살리언시-민감성이 낮아 전문가 할당이 균일하게 분포되는 반면 SharpMoE는 살리언시 수준이 높을수록 더 많은 전문가를 할당하는 뚜렷한 단조 증가 경향을 보인다는 점을 시각화했다. 특히 고노이즈 초기 단계에서 SharpMoE의 개선 폭이 더 크게 나타나며 이는 clean latent 입력의 시간적 안정성이 고노이즈 구간에서 결정적임을 시사한다. 전체 궤적 관점에서도 SharpMoE는 누적 할당이 살리언시 맵과 더 높은 상관을 보였다.
각 타임스텝 및 전체 궤적에 대해 살리언시 수준과 할당된 전문가 수의 분포를 비교한 막대그래프 모음이다.

용어 해설

Mixture of Experts: — Mixture-of-Experts는 입력 토큰별로 소수의 전문가(expert)만 활성화하여 전체 모델 용량을 크게 확장하면서 연산 효율을 유지하는 구조이다. 라우터가 각 토큰과 전문가 간 유사도 점수를 예측하고 Top-K 기반으로 전문가를 선택하며 선택된 전문가들의 출력을 가중합해 최종 출력을 만든다. 이 논문에서는 MoE 라우팅이 노이즈에 민감해 시각적 중요도를 반영하지 못하는 문제를 다루고 있다.
Full-Trajectory Training: — Full-Trajectory Training은 확산 생성의 여러 연속된 타임스텝을 롤아웃하여 각 단계의 예측(clean latent)을 다음 단계의 라우팅 입력으로 전달하는 학습 방식이다. 이 방식은 이전 단계에서 생성한 노이즈 없는 추정치를 라우터에 제공해 시간적 재귀 의존성을 학습하게 한다. SharpMoE는 이 방식을 통해 노이즈가 많은 초기 단계에서도 안정적인 라우팅 신호를 확보한다.
Trajectory Routing Loss: — Trajectory Routing Loss는 전체 롤아웃에서 각 토큰에 할당된 전문가 활성화의 누적 합을 이미지의 라플라시안 기반 살리언시 맵과 정규화된 분포로 맞추는 손실 함수이다. 구체적으로 토큰별 누적 할당 점수 A를 softmax로 정규화하고 라플라시안으로 계산한 살리언시 맵 M의 softmax와 KL-divergence를 최소화한다. 이 손실은 중요한 영역에 더 많은 계산 자원을 집중시키도록 라우터를 유도한다.
Saliency-Harnessing Router: — Saliency-Harnessing Router는 현재 노이즈화된 latent와 직전 타임스텝에서 예측된 clean latent를 모두 입력으로 받아 라우팅 점수를 생성하는 보조 라우터이다. 이 라우터의 가중치는 초기화 시 0으로 설정되어 사후 학습(post-training) 동안 점진적으로 살리언시 신호를 도입한다. 결과적으로 노이즈에 의해 가려진 구조적 정보를 복원해 살리언시 기반의 정확한 전문가 할당을 가능하게 한다.

코드 예제

text

Algorithm 1 Recursive Full-Trajectory Training for SharpMoE
Input: SharpMoE model v_theta with saliency-harnessing router R_sal, Pre-trained model weights W_pre, Rollout steps T, Loss hyperparameter lambda_routing.
Initialize: theta_R_sal <- 0; theta_others <- W_pre.
while not converged do
  Sample x0 ~ X0 and x1 ~ X1.
  Sample {t_k}_{k=1}^T subset U(0,1) and sort such that t1 > t2 > ... > tT.
  Set t1 = 0.999.
  L_total = 0.
  for k = 1 to T do
    x_tk = t_k * x1 + (1 - t_k) * x0.
    if k == 1 then xhat0^{t0} = x_{t1}.
    v_k = v_theta(x_tk, t_k, sg(xhat0^{t_{k-1}})).
    L_fm = ||(x1 - x0) - v_k||^2.
    L_total += L_fm.
    xhat0^{t_k} = x_tk - t_k * v_k.
    Collect routing scores S_k.
  end for
  Calculate L_routing with {S_k}.
  L_total = (1/T) * L_total + lambda_routing * L_routing.
  Update theta by minimizing L_total.
end while
return v_theta

이 코드는 SharpMoE의 재귀적 full-trajectory 학습 절차를 의사코드 형태로 나타낸 것이다. 각 학습 반복에서 연속된 T개의 타임스텝을 롤아웃하며 이전 단계의 clean prediction을 다음 단계의 라우터 입력으로 사용하고 Flow Matching 손실과 Trajectory Routing Loss를 결합해 최종 손실을 계산한다.