A generative flow matching model that unifies image generation and editing
type
Post
thumbnail
category
📑 Paper
updatedAt
Jul 11, 2025 08:22 AM
Abstract
FLUX.1 Kontext는 간단한 sequence concatenation 방식으로 a single unified architecture 내에서 local edit과 in-context 생성 작업이 모두 가능합니다.
FLUX.1 Kontext는 objects와 character의 보존 성능이 향상되어, 반복적인 워크플로우에도 견고합니다.
개선 사항들을 검증하기 위해, 종합적인 벤치마크인 KontextBench도 소개합니다.
1. Introduction
생성 모델을 사용한 콘텐츠 합성의 핵심적인 2가지 보완 기능입니다.
Local editing: 주변 context를 온전히 유지하면서 local에서 제한적으로 수정
Generative editing: 시각적 개념(예: 특정 대상)을 추출한 후, 새로운 context에 맞게 합성
Shortcomings of recent approaches
(i) 합성 쌍에 대해 훈련된 instruction 기반 방법은 생성 파이프라인의 단점을 그대로 이어받아 다양성과 realism 성능을 제한합니다.
(ii) 여러 번의 편집에서 characters와 objects의 정확한 모양을 유지하는 것은 여전히 해결되지 않은 문제입니다.
(iii) 대규모 멀티모달에 통합된 autoregression 모델은 denoising 기반 방식에 비해 품질이 낮을 뿐만 아니라 긴 런타임이 발생하여 대화가 끊기는 경우도 많습니다
Kontext Solution
FLUX.1 Kontext는 context와 instruction tokens이 concatenate된 sequence에 대해 속도 예측만을 사용하여 학습된 간단한 flow matching 모델입니다.
Character consistency: FLUX.1 Kontext는 여러 번의 반복 편집 과정을 거쳐도 character 보존에 탁월합니다.
Interactive speed: FLUX.1 Kontext 는 빠릅니다. text-to-image와 image-to-image 모두 1024 × 1024 이미지 합성에 걸리는 시간은 3~5초 입니다.
Iterative application: 빠른 inference와 robust 일관성은 최소한의 시각적인 drift로 이미지 편집을 여러 번 연속해서 가능하게 합니다.
2. FLUX.1
FLUX.1은 이미지 autoencoder의 latent space에서 학습된 rectified flow transformer입니다. 그리고 적대적 목표를 가진 convolutional autoencoder를 처음부터 학습 시킵니다. 학습 컴퓨팅 규모를 확장하고 16개의 latent channels를 사용하여 다른 모델에 비해 재구성 기능을 향상 시켰습니다.
Double stream blocks
FLUX.1은 double stream과 single stream blocks 혼합하여 만들어졌습니다. Double stream blocks는 image와 text tokens에 대해 분리된 weights를 가지고 있으며, 그 둘의 혼합은 concatenate된 tokens에 attention 연산을 적용하여 수행됩니다.
(i) Feedforward block 내의 modulation parameter의 수를 절반으로 줄였습니다.
(ii) Attention input으로 들어가는 projection()과 MLP의 첫 번째 linear layer를 하나의 linear layer로 융합하고 (다시 split하여 사용), 마찬가지로 output도 concat 후 하나의 linear layer에서 처리 합니다. 이를 통해 더 큰 행렬-벡터 곱을 할 수 있게 됩니다 (순차적으로 계산이 아닌 병렬 처리 가능해 짐).
Positional embedding에는 인수분해된 3차원 Rotary Positional Embeddings (3D RoPE)를 활용합니다. 모든 latent token은 시공간 좌표로 index 됩니다 .
3. FLUX.1 Kontext
Kontext의 목표는 모델을 text prompt와 reference image를 함께 조건으로 하는 이미지를 생성할 수 있도록 학습 시키는 것입니다.
, 는 target image이고, 는 추가적인 conditional image이며, 는 text prompt입니다.
Conditional distribution 를 모델링하여 동일한 network가 일 때 local edits와 in-context 생성을 수행하고, 일 땐 T2I 작업을 수행합니다.
(i) → image-driven edits
(ii) → T2I
학습은 FLUX.1 T2I checkpoint에서 시작합니다. 그리고 optimization을 수행을 위해, 수많은 쌍을 수집했습니다.
Token sequence construction
이미지들은 frozen FLUX autoencode에 의해 latent로 encoding됩니다. 그 후에 context image tokens 는 image tokens 에 더해지고 img stream에 입력으로 들어갑니다.
이런 간단한 sequence concatenation만으로,
(i) 다양한 입력/출력 해상도 및 종횡비를 지원합니다.
Pixel space가 아니라 Latent space에 encoding되기 때문 + Transformer architecture의 유연성(고정된 크기의 그리드 입력에 의존하는 CNN과 달리 가변 길이의 sequence를 처리할 수 있음)
(ii) 여러 reference images로 쉽게 확장 가능합니다. ()
만약 하나의 sequence로 concatenate하지 않고 cross attention 방식을 선택했다면, 확장하기 위해서는 architecture 구조가 복잡해지게 될 것입니다.
게다가 cross-attn은 q가 image sequence에서 오고, k, v는 txt에서 오게 되는데 (txt의 q로 img의 k, v를 attn하는 것은, 텍스트가 이미지의 특정 공간적 특징을 직접적으로 탐색하고 지시해야 한다는 의미이기 때문에 모델이 이해하기 어려울 수 있음), 이는 비대칭적인 정보 흐름을 가질 수 있습니다.
※ Sequence concatenation vs Channel-wise concatenation of and
x와 y의 channel-wise concatenation도 테스트 해봤지만, 초기 실험에서는 성능이 별로 좋지 않았습니다.
Positional 정보는 3D RoPE embeddings을 사용해서 encode했습니다. 모든 context tokens y에 해당하는 embeddings은 상수 값의 offset을 갖습니다. 그 offset을 가상 time step으로 활용해서 내부 공간적 구조는 건들지 않은 채로 context와 target block을 확실하게 구분하도록 했습니다.
Constant offset
constant offset은 y tokens이 x tokens과 동일한 “기준 위치"에서 시작하는 것처럼 보이게 만들면서, 그들만의 상대적인 위치 정보는 유지하도록 하는 트릭입니다.
context y tokens가 x token sequence 뒤에 이어 붙여지더라도 단순히 sequence가 길어지는 것(선형적으로 이어짐)으로 인해 원래 가지고 있던 공간적 위치 정보가 왜곡되지 않고,
x : [0 ~ N-1]
y : [N ~ N+M-1]
x와는 별개의, 그러나 x와 동등한 기준점을 가진 독립적인 공간적 context를 제공한다는 것을 모델에게 알려주기 위해 constant offset을 적용함.
x : [0 ~ N-1]
y : [N-C ~ N+M-C-1]
A token position is denoted by the triplet , for the target tokens, for context tokens.
Rectified-flow objective
Flow matching 정리
- (1) 식
- (2) 식
Notation
: 시간 에 따라 데이터 가 어떻게 노이즈 와 섞여서 중간 상태 가 되는지를 나타내는 함수
: 원본 데이터 (target image)
: 에서 샘플링된 무작위 noise vector
: scaling, scaling
: (maps to) 을 입력으로 받아 라는 형태의 로 변환한다는 의미
- (3) 식
Notation
: 주어진 noise 일 때, 특정 상태에서의 조건부 벡터 필드 (벡터 필드는 각 점 에서 이 점이 앞으로 어떤 방향으로 움직여야 하는지, 변화율을 가르키는 벡터)
: ‘정의’
: 경로 함수 를 에 대해 미분한 것 (의 작은 변화가 에 얼마나 영향을 미치는지 즉, 경로함수의 속도)
: 경로함수 의 역함수. 노이즈가 섞인 상태 와 노이즈 이 주어졌을 때, 원래의 데이터 가 무엇이었는지 추론하는 함수.
⇒ 전체 식의 의미: Flow matching에서 사용되는 CFM(Conditional Flow Matching)의 핵심적인 정의 중 하나로, 목표는 에서 으로 가는 흐름을 학습하는 것입니다. 이를 위해 각 와 쌍에 대해 고유한 직선 경로를 정의하고, 이 경로를 따라가는 벡터 필드를 로 정의합니다. 이 벡터필드는 가 에 따라 어떻게 움직여야 하는지를 지시합니다. 모델은 나중에 를 예측하도록 학습합니다.
- (4) 식
Flow matching:
Conditional Flow matching:
Conditional (조건부) → marginal (무조건부) 전환
식에서 보면, 학습 시킬 신경망인 는 조건의 사전 분포에서 샘플링 되는 를 모르는데, 목표 값인 을 아는 것은 불가능함.
그래서 를 가 아니라 으로 정의함.
⇒
그런데 결합 분포 정의 에 따라,
에서 샘플링하고, 그 이 주어졌을 때 에서 샘플링하는 것과 동일함.
즉, 원래 CFM 방식 자체가 () 쌍을 결합 분포에서 샘플링하는 것과 동일함.
따라서 을 정의하면
속도 예측 → 노이즈 예측 전환
와 를 3번 째 식에 넣으면,
⇒
여기에 다시 log signal-to-noise ratio 와 를 사용하면,
⇒
이걸 에 대입하면,
- (5) 식
최종 (※ )
⇒
정리
Flow matching:
결정 경로를 따르지만, 직선이 아닐 수도 있어서 많은 스텝 수가 필요함.
조건부 → 무조건부 flow matching
Sequence concatenation을 통해 CFM을 FM형태로 학습이 가능함 (조건 유/무에 따라 손실 함수 변경하지 않아도 학습 가능)
Rectified flow matching:
결정론적 결정 경로를 여러 번의 ‘정류(rectify)’과정을 통해 직선 경로로 변함.
경로가 직선에 가까워지면, 노이즈에서 데이터까지 도달하는데 필요한 샘플링 수를 극적으로 줄일 수 있음. (ex: 50~250 → 1~4 steps)
이 학습과정의 손실 함수를 수학적으로 변환하면, 결과적으로 모델이 단순한 정규분포를 따르는 노이즈를 예측하는 형태와 동일해짐.
직선 경로를 따라가는 속도 벡터 필드를 마치 노이즈를 예측하는 것처럼 표현할 수 있게 됨.
⇒ 노이즈 은 즉, 깔끔한 벡터임. 모델이 이러한 단순한 노이즈를 예측하는 것이 복잡한 속도 벡터 필드를 직접 예측하는 것보다 학습하기 쉬울 뿐만 아니라, diffusion 모델과 유사해져서 가중치 스케일링, 샘플링 전략 등을 적용하기도 쉬워짐.
We use a logit normal shift schedule for , where we change the mode depending on the resolution of the data during training. When sampling pure text-image pairs () we omit all tokens , preserving the text-to-image generation capability of the model.
Expressing shifting of the timestep schedule via the Logit-Normal Distribution
학습 중에 데이터의 해상도에 따라 모드 를 변경하는 에 대해 logit normal shift schedule을 사용합니다. 순수한 텍스트-이미지 쌍()을 샘플링할 때는 모든 토큰 를 생략하여 모델의 텍스트-이미지 생성 기능을 유지합니다.
logit normal shift schedule
학습 과정에서 시간 t를 어떻게 샘플링할지 t의 분포를 제어함. 특정 noise-level 구간에 모델의 학습을 집중 시키기 위해 사용.
ex) noise가 많을 땐 전체적인 형태와 같은 큰 특징을, noise가 적을 땐 질감, 패턴과 같은 디테일에 대한 분포를 학습하도록 도움.
값()을 조절함으로써, 특정 해상도 이미지를 학습할 때 모델이 더 많은 시간을 특정 t 구간(특정 noise-level)의 샘플을 보고 학습하도록 유도함.
요약하자면, 타임스텝을 1000중에 어디를 세세히 쪼갤 지를 정하는 것
Adversarial Diffusion Distillation
flow matching 모델의 샘플링에는 일반적으로 50~250개의 guided network evaluation를 사용하여 일반(ODE) 또는 확률 미분 방정식(SDE)을 푸는 작업이 포함됩니다.
여기에는 몇 가지 잠재적인 단점이 있습니다.
(i) such multi-step sampling is slow, rendering model-serving at scale expensive and hindering low-latency, interactive applications.
(ii) guidance may occasionally introduce visual artifacts such as over-saturated samples.
저희는 latent adversarial diffusion distillation (LADD)를 사용하여 샘플링 단계 수를 줄이면서 적대적 훈련을 통해 샘플의 품질을 높이는 두 가지 challenge를 모두 해결합니다.
Implementation details
T2I 체크포인트에서 시작하여 I2I 및 T2I 작업에서 모델을 jointly 미세 조정합니다. 이 방식은 여러 reference image를 포함할 수 있지만, 현재는 컨디셔닝을 위해 단일 context image에 중점을 둡니다.
FLUX.1 Kontext [pro]는 flow objective에 이어 LADD로 훈련됩니다. 12B DiT로 guidance-distillation를 통해 FLUX.1 Kontext [dev]를 얻습니다. 편집 작업에서 FLUX.1 Kontext [dev]의 성능을 최적화하기 위해 I2I 훈련에만 집중하며, FLUX.1 Kontext [dev]는 T2I 작업은 훈련하지 않습니다.
all-gather 연산은 bfloat16으로 수행되는 반면, gradient reduce-scatter은 수치 안정성을 높이기 위해 float32를 사용하는 등 FSDP2를 혼합 정밀도와 함께 사용합니다. selective activation checkpointing을 사용하여 최대 VRAM 사용량을 줄입니다. 처리량을 개선하기 위해 Flash Attention 3와 개별 transformer block의 지역 컴파일을 사용합니다.
4. Evaluations & Applications
Kontext Bench - Crowd-sourced Real-World Benchmark for In-Context Tasks
모델 편집을 위한 기존 벤치마크는 real-world 사용 현황을 파악하는 데 한계가 있는 경우가 많습니다.
(i) InstructPix2Pix, relies on synthetic Stable Diffusion samples and GPT-generated instructions, creating inherent bias.
(ii) MagicBrush, is constrained by DALLE-2’s capabilities during data collection.
(iii) Emu-Edit, uses lower-resolution images with unrealistic distributions and focus solely on editing tasks.
(iv) DreamBench, lacks broad coverage
(v) GEdit-bench, does not represent the full scope of modern multimodal models.
(vi) IntelligentBench, remains unavailable with only 300 examples of uncertain task coverage.
이러한 격차를 해소하기 위해 크라우드 소싱된 실제 사용 사례에서 KontextBench를 컴파일했습니다. 이 벤치마크는 개인 사진, CC 라이선스 아트, 퍼블릭 도메인 이미지, AI 생성 콘텐츠 등 108개의 기본 이미지에서 파생된 1026개의 고유한 이미지-프롬프트 쌍으로 구성되어 있습니다.
5가지 핵심 tasks을 다루고 있습니다.
(i) Local instruction editing (416 examples)
(ii) Global instruction editing (262)
(iii) Text editing (92)
(iv) Style reference (63)
(v) Character reference (193)
State-of-the-Art Comparison
FLUX.1 Kontext는 T2I와 I2I 합성을 모두 수행하도록 설계되었습니다. 위에서 언급했듯이 [dev]의 경우 I2I 작업에만 집중했습니다. 또한, 생성 성능을 개선하기 위해 더 많은 컴퓨팅을 사용하는 FLUX.1 Kontext [max]도 추가했습니다.
Image-to-Image Results
이미지 편집 evaluation의 경우 다양한 성능을 평가합니다.
(i) Image quality
(ii) Local editing
(iii) Character reference (CREF)
(iv) Style reference (SREF)
(v) Text editing
(vi) Computational efficiency
전반적으로 FLUX.1 Kontext는 character 일관성 유지 성능에서 SOTA 입니다.
Text-to-Image Results
현재 T2I 벤치마크는 주로 일반적인 선호도에 초점을 맞추고 있습니다. ("어떤 이미지를 선호하십니까?") 이러한 광범위한 평가 기준은 종종 채도가 높은 색상, 중앙에 위치한 subjects에 대한 과도한 초점, 뚜렷한 bokeh 효과, 동질적인 스타일로의 수렴을 의미하는 특징적인 'AI 미학'을 선호한다는 것을 관찰할 수 있습니다.
→ “bakeyness”
이러한 한계를 해결하기 위해 T2I 평가 지표를 5가지로 세분화했습니다.
(i) Prompt following
(ii) Aesthetic (” Which image do you find more aesthetically pleasing?”)
(iii) Realism (”Which image looks more real?”)
(iv) Typography accuracy
(v) Inference speed
이 벤치마크를 Internel-T2I-Bench라고 부르며, GenAI 벤치에 대한 additional evaluation로 이 벤치마크를 보완합니다.
경쟁 모델이 특정 영역에서는 우수하지만 다른 영역에서는 좋지 못한 경우가 많습니다. FLUX.1 Kontext는 평가 카테고리 전반에서 균형 잡힌 성능을 보여줍니다.
Iterative Workflows
Multiple edits에서 character와 object의 일관성을 유지하는 것은 스토리텔링이 중요한 애플리케이션의 경우 매우 중요합니다. 입력 이미지와 연속 편집을 통해 생성된 이미지 사이의 AuraFace embeddings의 cosine 유사도를 추가로 계산해보니, 다른 모델에 비해 FLUX.1 Kontext의 일관성 유지가 더 잘 됩니다.
Specialized Applications
FLUX.1 Kontext는 standard 생성, 스타일 참조(SREF) 외에도 여러 애플리케이션을 지원합니다. 특히 빨간색 타원과 같은 기하학적 마커에 반응하여 목표 수정을 안내하는 시각적 단서를 통해 직관적인 편집을 지원합니다.
5. Discussion
FLUX.1 Kontext는 context에 맞는 이미지 생성과 편집을 단일 프레임워크에 결합한 flow matching 모델입니다. 간단한 sequence concatenation 및 학습 방법을 통해, FLUX.1 Kontext는 state-of-the-art 성능을 달성했고 동시에 multi-turn edit 시 character 변화, 느린 inference 속도, 낮은 출력 품질과 같은 주요 제한 사항을 해결했습니다.
주요 contributions은 multiple processing tasks를 다루는 하나의 통합된 architecture, 반복 작업 전반에서 뛰어난 character 일관성, interactive 속도, and KontextBench 입니다.
Limitations
과도한 multi-turn editing은 이미지 품질을 저하시키는 시각적 아티팩트를 유발할 수 있습니다.
Future work
여러 reference image 입력이 가능하도록 확장하고, inference latency를 줄여 real-time 애플리케이션이 가능하도록 집중해야 합니다. 가장 중요한 것은 multi-turn editing 중 성능 저하(열화 현상)를 줄이면 무한히 유동적인 콘텐츠 제작이 가능하다는 것입니다.
