PixelMoE가 픽셀 공간 확산 모델에 LLM식 MoE 희소 확장을 이식한다

무엇인가

픽셀 공간 확산 모델은 VAE 토크나이저 없이 원본 이미지를 직접 모델링한다는 점에서 매력적이지만, 지금까지는 대부분 dense하게 스케일해 왔다. 모델 용량을 키우면 활성 연산량도 같이 늘어나는 구조였다. 반면 LLM에서는 MoE가 용량과 활성 연산을 분리하는 표준 스케일링 패러다임으로 자리 잡았는데, 이를 확산 모델에 그대로 옮기면 dense 대비 뚜렷한 이득이 일관되게 나오지 않았다. 그래서 기존 연구들은 Expert-Choice 라우팅이나 확산 전용 라우팅 가이드·보조 목표를 도입했다. 이 논문은 여기서 질문을 던진다. 표준 LLM식 MoE를 픽셀 공간 확산에 직접 이식할 수 있는가. 저자들의 답은 '그렇다'이며, 그 결과가 PixelMoE다.

어떻게 동작하나

구조부터 보면, PixelMoE는 초기 트랜스포머 블록은 dense FFN으로 유지하고 나머지 블록의 FFN을 희소 MoE 레이어로 교체한다. 각 MoE 레이어는 셀프 어텐션 뒤에 붙고, 이미지 토큰마다 표준 Token-Choice 라우팅으로 상위 A개의 세분화된 라우티드 전문가를 고른다. 라우터는 bias 없는 선형 계층에 시그모이드를 적용해 점수를 만들고, 전문가 선택 시에만 밸런싱 bias를 더한 뒤 선택된 전문가들 사이에서 점수를 정규화한다. 출력은 라우티드 전문가 출력에 sqrt(A) 스케일을 곱한 값과, 모든 토큰에 항상 적용되는 shared expert 출력의 합이다. 라우티드·shared 전문가 모두 중간 차원을 줄인 FFN이라 합산 활성 파라미터가 dense 기준선과 비슷해진다. 로드 밸런싱은 보조 손실 없이 bias를 직접 갱신하는 방식으로 처리한다. 배치의 전문가 부하 분포 F와 균등 목표 Q의 차이를 RMS로 정규화해 학습률 0.001로 빼주는 식이다. 중요한 점은 라우터가 토큰 hidden state만 보고 결정하며, timestep이나 클래스 같은 명시적 의미 정보를 받지 않는다는 것이다.

무엇과 다른가

여기에 Pixel-Space Hierarchical Flow Matching이 붙는다. 픽셀 공간은 리샘플링만으로 의미 있는 다해상도 이미지를 얻을 수 있다는 성질을 이용해, 디노이징 궤적을 따라 공간 해상도를 바꾼다. Stage 0은 저해상도에서 부분 디노이징만 수행해 전역 구조를 잡고, Stage 1은 전체 해상도에서 세부 디테일을 다듬는다. 두 단계는 같은 PixelMoE 백본과 DDT(Decoupled Diffusion Transformer) 예측 헤드를 공유한다. Stage 0은 2배 bilinear 다운샘플된 이미지를 목표로 하고, 256x256에서는 노이즈 스케일 1, 512x512에서는 2를 쓴다. 각 단계 안에서는 단계 상대 timestep σ를 두고 시작점과 끝점을 선형 보간하며, 모델은 각 단계의 끝점을 예측한 뒤 이를 velocity로 환산하는 손실로 학습한다. σ는 0.05로 클램프해 수치 안정성을 확보한다. 샘플링 시에는 Stage 0의 업샘플된 끝점을 그대로 Stage 1 시작 상태로 쓸 수 없어(신호 계수와 노이즈 공분산이 다르다) 재스케일하고 블록 상관 가우시안 노이즈를 주입해 통계를 맞춘다. timestep 샘플링도 단계 배분(전역 LogitNormal 0.8, 0.8)과 단계 내 상대 timestep(로컬 LogitNormal 0.6, 0.8)을 분리한다.

어떻게 쓰나

실험은 ImageNet-1K 클래스 조건부 생성에서 수행했다. 기존 방법과의 비교 모델은 320 에폭, 그 외 실험은 100 에폭 학습했고, FID·IS·Precision·Recall을 측정하며 5만 장을 Heun 솔버 50스텝, CFG로 샘플링했다. 256x256에서 PixelMoE-B-E32A4는 활성 파라미터 225M로 FID 1.88, L-E32A4는 1.54, H-E32A4는 1.47을 기록했다. 512x512에서는 각각 2.05, 1.58, 1.53이다. 이 수치는 표현 정렬(representation alignment)이나 perceptual loss 없이 얻은 것이다. 희소 스케일링 자체의 효과도 분리해 측정했는데, 활성 파라미터를 632M로 고정한 채 라우티드 전문가 수 E를 8에서 16, 32로 늘리면 FID가 2.28에서 2.03, 1.78로 개선됐다. 수렴 속도도 달라서, PixelMoE-L-E32A4는 Dense-L이 200 에폭에 도달한 FID를 약 71 에폭 만에 달성해 필요한 학습량을 약 2.8배 줄였다. 패치 크기를 줄여 토큰을 세분화할수록 dense 대비 개선폭이 더 컸는데, 저자들은 토큰 표현이 더 국소화되고 라우팅 결정 수가 늘어나기 때문이라고 설명한다.

전제와 한계

라우팅 분석 결과도 흥미롭다. PixelMoE-L-E32A4의 검증셋 라우팅을 보면 클래스별 전문가 선호는 얕은 층에서 약하다가 중간층에서 뚜렷해지고 특히 12~20층에서 두드러진다. 반대로 디노이징 궤적에 따른 선호는 얕은 층에서 이미 강하고 중간층에서 약해졌다가 마지막 층에서 다시 강해진다. 두 계층적 단계 사이 전환에서도 패턴이 바뀐다. 그런데 전역 전문가 부하는 모든 분석 층에서 균등 분포에 가깝다. 즉 특정 전문가 몇 개가 계속 과점해서 생긴 패턴이 아니라, 균형을 유지하면서 입력 적응적 전문화가 일어난 것이다. 절제 실험도 이를 뒷받침한다. shared expert를 넣으면 FID가 2.22에서 1.78로 좋아지면서 총 파라미터는 2.00B에서 1.65B로 줄었다. 전문가 구성을 E8A1에서 E32A4로 바꾸면 가능한 조합이 8개에서 35,960개로 늘어나고 FID는 1.95에서 1.78로 개선됐다. bias 기반 밸런싱은 보조 로드밸런싱 손실 대비 2.05에서 1.78로 앞섰고, Token-Choice 라우팅은 Expert-Choice를 쓰는 DiffMoE 구현보다 FID가 크게 낮았다. timestep이나 의미 정보를 라우터에 명시적으로 넣는 변형은 오히려 기본 라우터를 넘지 못했다. 계층적 flow는 2.54에서 2.46, 단계 인식 timestep 샘플링은 2.65에서 2.46, DDT 예측 헤드는 3.76에서 2.46으로 각각 기여했다.

개발자 입장에서 이 논문의 실용적 의미는 두 가지다. 첫째, 픽셀 공간 생성에서도 MoE가 dense 대비 확실한 스케일링 이득을 준다는 통제된 비교 결과를 제시한다. 총 파라미터를 늘리되 활성 파라미터는 고정하는 방식으로 품질을 사려면 참고할 만하다. 둘째, 확산 전용 라우팅 트릭이 필요 없다는 주장이다. 토큰 hidden state만으로 라우팅해도 클래스와 노이즈 레벨 정보가 이미 충분히 담겨 있다는 것이 저자들의 해석이다. 다만 실제 도입 전에 확인할 것은, 이득이 ImageNet-1K 클래스 조건부 생성과 256·512 해상도에 한정해 측정됐다는 점, 그리고 512 해상도에서는 노이즈 스케일 같은 하이퍼파라미터를 다르게 쓴다는 점이다. 텍스트-이미지처럼 조건이 복잡한 태스크나 다른 해상도로 옮길 때 라우팅 전문화 패턴이 유지되는지는 별도 검증이 필요하다.

한계와 전제를 정리하면 이렇다. 원문 본문에는 별도의 한계 절이 없고, 저자들이 명시한 전제는 실험 범위와 비교 조건에 집중돼 있다. 모든 스케일링 비교는 '동일한 활성 파라미터' 조건에서 이뤄졌고, 절제 실험은 PixelMoE-L-E32A4와 PixelMoE-B-E32A4 두 구성에만 적용됐다. 데이터셋은 ImageNet-1K 클래스 조건부 생성 하나뿐이며, 평가는 FID·IS·Precision·Recall로 제한된다. 또한 저자들은 기존 픽셀 공간 MoE인 JiTMoE가 dense JiT 대비 미미한 이득만 보고했다는 점을 지적하며 자신들의 결과가 다르다고 주장하는데, 이 차이가 라우팅 설계 때문인지 계층적 flow나 학습 레시피 때문인지는 완전히 분리되지 않았다. 마지막으로 코드와 모델은 '곧 공개'한다고만 밝혀 현재로서는 재현이 불가능하다.