Mage-Flow - An Efficient Native-Resolution Foundation Model for Image Generation and Editing

X. Zhang, P. Zhang, S. Zheng, J. Guo, et al., "Mage-Flow: An Efficient Native-Resolution Foundation Model for Image Generation and Editing," arXiv:2607.19064, 2026.

이미지 생성 모델의 최근 흐름을 한 마디로 요약하면, 더 크게 만드는 것입니다. Z-Image가 6B, Qwen-Image가 20B, FLUX.2가 32B, HunyuanImage 3.0이 80B입니다. 품질은 확실히 올라갔지만, 그만큼 이 모델들을 직접 고치고 실험하고 배포하는 일은 점점 어려워졌습니다. 가중치가 공개되어도 연구자가 만질 수 없으면 열린 모델이라 부르기 어렵습니다.

Mage-Flow는 반대 방향을 택했습니다. 백본을 4B에 고정해 두고, 대신 그 주변을 다시 설계했습니다. 결과적으로 GenEval에서 비교 대상 전체 중 최고점을 받으면서 A100 한 장에서 1024² 이미지를 0.59초에 생성합니다.

저자

논문에는 개인 이름이 아니라 Microsoft Mage Team이라는 팀 이름이 먼저 나옵니다. 기여자 명단은 본문 맨 뒤에 24명으로 붙어 있습니다.

프로젝트 리드는 장신제입니다. 홍콩과학기술대학교에서 박사 과정을 밟으며 신경 압축을 연구했고, 이 논문의 출발점인 "VAE를 학습된 이미지 코덱으로 본다"는 관점이 그 연구 배경에서 곧장 나옵니다. 공동 1저자 중 한 명인 선이페이는 확산 모델 가이던스 이론을 연구해 온 인물로, 4스텝 Turbo 변형의 증류 설계와 맞닿아 있습니다.

시니어 저자 리빈과 루옌은 이 팀의 성격을 설명해 줍니다. 두 사람은 DCVC 계열 신경 영상 압축 연구를 함께 이끌어 왔습니다. 생성 모델 팀이 아니라 코덱 팀이 이미지 생성 논문을 쓴 셈이고, 그래서 이 논문의 문제 설정이 "어떻게 더 잘 그리나"가 아니라 "어떻게 더 싸게 표현하나"에 가깝습니다.

배경

잠재 확산 파이프라인에서 사람들이 보통 최적화하는 대상은 확산 백본입니다. 그런데 토크나이저는 학습할 때도, 추론할 때도, 편집할 때마다 반복해서 호출됩니다. 그리고 그 비용은 해상도에 따라 빠르게 늘어납니다.

문제는 지금 쓰이는 공개 VAE 대부분이 256×256 이미지를 전제로 설계된 VQGAN·LDM 오토인코더 구조를 그대로 물려받았다는 점입니다. 전역 어텐션 블록처럼 고해상도에서 특히 비싼 구성 요소가 들어 있습니다. 논문이 든 예가 구체적입니다. FLUX.2-Klein-4B로 1K 이미지를 4스텝 생성할 때, VAE 디코딩이 전체 생성 시간의 14%를 차지합니다.

스텝 수를 줄일수록 이 비율은 커집니다. 확산 백본을 아무리 빠르게 만들어도 토크나이저가 병목으로 남는 구간이 생기는 것입니다.

어떻게 만들었나

Mage-Flow 스택은 두 개의 모델 요소와 하나의 시스템 요소로 이루어집니다. 경량 토크나이저 Mage-VAE, 4B 규모의 Native-Resolution MMDiT, 그리고 CUDA 커널 융합 학습 인프라입니다.

Mage-VAE

디코더 설계는 같은 팀의 이전 작업인 CoD-Lite에서 왔습니다. 압축 지향으로 확산 사전학습을 한 뒤 1스텝으로 증류하면, 무거운 DiT나 U-Net 디코더 없이도 강한 재구성 품질이 나온다는 결과입니다. 이 원리를 따라 Mage-VAE 디코더는 1스텝 픽셀 확산 모델이 됩니다. 합성곱 확산 블록을 쌓고 분리된 픽셀 확산 헤드로 RGB를 직접 복원하는 완전 합성곱 구조라, 전역 어텐션이 없고 디코딩 비용이 해상도에 거의 선형으로 붙습니다.

인코더는 디코더의 구조적 쌍대(dual)로 설계했습니다. 디코더가 잠재에 조건화된 픽셀 생성기라면, 인코더는 픽셀에 조건화된 잠재 생성기로 볼 수 있다는 발상입니다. 그래서 인코더도 1스텝 확산 모델로 만들어, 잠재 추출을 픽셀 복원만큼 가볍게 만들었습니다.

세 번째가 KL 정규화입니다. 보통의 VAE는 사후분포를 표준 가우시안 사전분포에 맞추지만, Mage-VAE는 FLUX.2-VAE가 유도하는 앵커 잠재 분포 쪽으로 정규화합니다. FLUX.2-VAE가 32채널에 8배 공간 축소를 쓰고 이 잠재가 확산 트랜스포머에 들어가기 전 보통 2배 패치화되므로, Mage-VAE는 아예 128채널로 16배 축소된 잠재를 직접 내놓아 패치화를 내부에 흡수했습니다. 생성에 친화적인 잠재 구조를 물려받으면서 트랜스포머가 바로 먹을 수 있는 형태를 만드는 것입니다.

학습은 3단계입니다. 1단계에서 인코더와 디코더를 각각 다단계 확산 모델로 사전학습하고, 2단계에서 디코더를 재구성 손실과 DMD 손실, DINOv2 투영 GAN 손실로 1스텝 증류하고, 3단계에서 앵커 잠재 KL을 걸어 둘을 함께 미세조정합니다.

Native-Resolution MMDiT

백본은 SD3가 도입한 MMDiT 블록 설계를 따르는 4B 트랜스포머입니다. 텍스트와 이미지 토큰을 이어 붙여 결합 자기어텐션으로 처리하되, 모달리티별 정규화·투영 층으로 두 스트림의 통계를 분리해 둡니다.

표준 MMDiT 학습과 갈리는 지점은 해상도 처리입니다. 관행적인 버킷 기반 학습에서는 이미지를 미리 정해 둔 해상도·종횡비 버킷에 배정하고, 한 스텝은 한 버킷에서만 샘플을 뽑습니다. 배칭은 간단해지지만 연속적인 해상도 분포가 이산화되고, 버킷 양자화 불일치가 생기고, 한 업데이트가 보는 종횡비 다양성이 줄어듭니다. 아주 넓거나 아주 긴 출력은 해당 버킷을 명시적으로 추가하지 않는 한 지원하기 어렵습니다.

Mage-Flow는 네이티브 해상도 시퀀스로 바로 학습합니다. 서로 다른 해상도의 이미지를 Mage-VAE로 인코딩해 가변 길이 잠재 토큰 열로 펼치고, 고정 토큰 예산 아래 하나의 연속 배치로 패킹합니다. 텍스트 조건도 최대 길이로 패딩하지 않고 같은 방식으로 패킹합니다. FlashAttention의 가변 길이 커널과 샘플별 누적 오프셋을 쓰면 명시적인 블록 대각 마스크를 만들지 않고도 어텐션을 샘플 안으로 제한할 수 있습니다.

효과는 두 가지입니다. 학습 쪽에서는 한 업데이트가 이질적인 이미지 크기와 종횡비를 함께 보게 되어, 체크포인트 하나가 512에서 2048까지의 높이·너비와 512×2048 같은 극단적 종횡비로 자연스럽게 일반화됩니다. 추론 쪽에서는 classifier-free guidance의 조건부·무조건부 분기를 한 배치로 묶어 한 번의 순전파로 평가할 수 있습니다. 원래의 디노이징 궤적을 보존하면서 CFG 오버헤드만 줄이는 방식이고, 측정된 가속은 \(1.09\times\)에서 \(1.15\times\)입니다.

텍스트 조건은 얼린 Qwen3-VL-4B-Instruct 인코더가 임베딩 \(\tau\)로 만듭니다. 이미지 \(x\)는 Mage-VAE가 잠재 \(z\)로 인코딩하고, 학습은 rectified flow matching 목적함수를 씁니다.

\[\mathcal{L}(\theta) = \mathbb{E}_{(x,\tau),t,\epsilon}\left[\left\lVert v_\theta(z_t, t, \tau) - (z - \epsilon)\right\rVert_2^2\right]\]

여기서 \(z_t = (1-t)z + t\epsilon\), \(\epsilon \sim \mathcal{N}(0, I)\) 입니다.

편집 모델은 같은 목적함수와 같은 백본을 쓰되 조건화 입력만 바꿉니다. Qwen3-VL이 편집 지시문과 원본 이미지를 함께 인코딩하고, Mage-VAE가 원본과 목표를 각각 \(z^{\text{src}}\), \(z^{\text{tgt}}\)로 인코딩합니다. 원본 토큰과 목표 토큰을 구분하기 위해 2D RoPE에 프레임 차원을 하나 더해 각 시각 토큰에 \((h, w, f)\) 위치 인덱스를 부여합니다. 손실은 목표 토큰에서만 계산하므로, 편집 전용 모듈을 따로 붙이지 않고 생성 모델에서 곧장 초기화할 수 있습니다.

커널 융합

여기가 논문에서 시스템 이야기가 나오는 부분입니다. 반복 실행되는 세 모듈, 즉 Mage-VAE의 합성곱 확산 블록과 얼린 Qwen3-VL 트랜스포머 블록과 4B NR-MMDiT 블록에는 주 연산 외에도 메모리 바운드 연산 체인이 잔뜩 붙어 있습니다. 정규화, 적응적 변조, RoPE 적용, 게이팅, 활성화, 잔차 덧셈입니다. 이들을 별도 CUDA 커널로 띄우면 큰 활성값 텐서를 반복해서 읽고 쓰게 되어 메모리 트래픽과 커널 실행 오버헤드가 쌓입니다.

이 체인들을 융합해 중간값을 온칩 메모리에 두고 최종 출력만 되쓰게 했습니다. 블록이 순전파·역전파마다 여러 번 실행되므로, 국소적인 연산자 융합이 스택 전체 처리량 이득으로 번역됩니다.

B200 8GPU 노드, 전역 배치 8, GPU당 패킹 시퀀스 50,000 토큰 기준입니다.

읽는 방식이 중요합니다. 토크나이저 교체만으로 이미 \(1.41\times\)가 나오는데, 이건 백본을 건드리지 않고 얻은 이득입니다. 반면 VAE와 텍스트 인코더 융합은 각각 소폭에 그칩니다. 가장 큰 몫은 마지막 줄, 4B 확산 백본 블록 융합에서 나옵니다. 학습 스텝을 지배하는 것이 백본이기 때문입니다. 초록이 말하는 약 \(2.5\times\)는 이 표의 \(2.48\times\)가 정본입니다.

결과

생성 쪽 주요 벤치마크입니다. Mage-Flow 계열은 전부 4B입니다.

GenEval에서 Mage-Flow가 0.90으로 비교 대상 전체 1위입니다. 80B인 HunyuanImage-3.0이 0.72, 32B인 FLUX.2-dev가 0.87이라는 점을 같이 보면 파라미터 수와 이 지표의 상관이 생각보다 약하다는 이야기가 됩니다. 세부 항목에서는 위치 관계(position)가 0.93으로 특히 두드러지는데, 다른 모델 대부분이 0.2에서 0.8 사이에 흩어져 있는 항목입니다.

다만 DPG-Bench와 OneIG-EN에서는 Mage-Flow가 1위가 아닙니다. 오히려 DPG-Bench는 RL 정렬 전인 Base가 86.26으로 정렬 후 모델과 비슷합니다. 모든 축에서 이긴 것이 아니라 프롬프트 준수와 텍스트 렌더링 축에서 강하다고 읽는 편이 정확합니다.

편집 쪽입니다.

GEdit 중국어에서 오픈소스 1위, 영어에서 근소한 2위입니다. 눈여겨볼 부분은 스텝 수입니다. 4스텝짜리 Turbo가 50스텝짜리 20B 모델들과 같은 자리에 있습니다. ImgEdit에서는 4.38로 FireRed의 4.56에 밀리는데, 세부 항목을 보면 Extract에서 3.94 대 4.34로 벌어지는 것이 주 원인입니다.

속도와 메모리는 이 논문의 본진입니다. A100 한 장 기준으로 Mage-Flow가 4.37초, 4스텝 Mage-Flow-Turbo가 0.59초, 30스텝 Mage-Flow-Edit이 10.55초, Mage-Flow-Edit-Turbo가 1.02초입니다. 최대 GPU 메모리는 생성·편집 모두 18에서 20GB 사이로, 비교 모델 중 가장 낮습니다. FLUX.2-dev가 179.63GB를 써서 A100 두 장이 필요한 것과 대비됩니다.

Mage-VAE 자체 성능도 정리해 둘 만합니다. CLIC 2020 기준 PSNR 36.61로 FLUX.2-VAE의 36.88에 근소하게 못 미치지만 SSIM은 0.9450 대 0.9447로 앞섭니다. FFHQ에서는 PSNR·SSIM·LPIPS 모두 최고입니다. 그러면서 인코딩 MACs가 픽셀당 173 대 2134, 디코딩이 215 대 4798입니다. 논문 표현으로 약 \(12.3\times\)와 \(22.3\times\) 절감입니다. 4096×4096에서는 비교 대상 VAE 대부분이 OOM으로 죽는데 Mage-VAE는 동작합니다.

회고

토크나이저를 바꿔 끼워 본 실험이 흥미롭습니다. 백본을 고정한 채 Mage-VAE와 FLUX.2-VAE를 맞바꿔 보면, Mage-Flow-Turbo 백본에서는 TIIF-Short가 83.58 대 81.18로 Mage-VAE가 낫습니다. 그런데 FLUX.2-Klein-4B 백본에 Mage-VAE를 꽂으면 오히려 소폭 떨어집니다. 토크나이저 단독 성능이 아니라 백본과의 공동 설계가 이득의 출처라는 뜻이고, 이건 "Mage-VAE만 가져다 쓰면 된다"는 결론을 막는 결과입니다.

저자들이 인정한 한계도 있습니다. LongText 중국어 분할은 영어만큼 강하지 않고, 데이터 보강 여지가 남아 있다고 본문에서 직접 적었습니다. OneIG의 Diversity 항목은 Base 0.159, RL 정렬 후 0.124, Turbo 0.105로 단계마다 떨어집니다. 정렬과 증류가 선호도를 올리는 대신 출력 다양성을 깎는 흔한 대가가 여기서도 그대로 나타납니다.

논문이 시스템 보고서 성격이라 소거 실험이 균일하지 않다는 점도 짚어 둘 만합니다. 커널 융합 표는 잘 쪼개져 있지만, 앵커 잠재 KL을 표준 가우시안 KL로 되돌렸을 때 어떻게 되는지 같은 핵심 설계 선택의 소거는 본문에 없습니다. 향후 과제로는 더 견고한 다중 이미지 편집, 더 강한 다국어 장문 렌더링, 에이전트형 시각 창작 워크플로와의 통합을 들었습니다.

정리

첫째, 확산 파이프라인에서 토크나이저는 곁다리가 아니라 예산 항목입니다. 스텝 수를 줄일수록 비중이 커지므로, 백본만 최적화하면 어느 지점에서 멈춥니다.

둘째, 버킷을 없애고 네이티브 해상도로 패킹하면 해상도 다양성이 제약이 아니라 학습 신호가 됩니다. 같은 메커니즘이 추론 단계에서 CFG 비용도 줄여 줍니다.

셋째, 4B로 32B·80B와 겨루려면 모델과 학습 시스템을 함께 설계해야 합니다. MFU 13.88%에서 29.28%로 가는 구간은 알고리즘이 아니라 커널에서 나왔습니다.