PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Rectified Flow Transformers for High-Resolution Image Synthesis

Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion English, Kyle Lacey, Alex Goodwin, Yannik Marek, Robin Rombach

arXiv:2403.03206v1cs.CV

TL;DR

Rectified flow offers a simple straight path from data to noise but lacks decisive evidence in high-resolution text-to-image synthesis. This paper introduces perceptually biased timestep sampling and a bidirectional multimodal transformer, finding that its largest models outperform state-of-the-art open and closed-source models.

  • Problem

    Rectified flow has not been decisively established for high-resolution text-to-image synthesis, with prior empirical evidence mostly limited to smaller class-conditional models.

  • Method

    The paper reweights rectified-flow noise scales toward perceptually relevant timesteps and combines this with a scalable transformer using bidirectional image–text token mixing.

  • Results

    The largest models outperform leading open and closed-source models in prompt-understanding evaluations and human preference ratings.

  • Takeaways & Limitations

    Validation-loss improvements correlate with stronger text-to-image benchmarks and human preferences, while modified timestep sampling improves rectified-flow performance over uniform sampling.

  • Takeaways & Limitations

    The high-resolution finetuning analysis assumes constant images, an assumption the paper notes is not realistic.

Abstract

from arXiv · show

Diffusion models create data from noise by inverting the forward paths of data towards noise and have emerged as a powerful generative modeling technique for high-dimensional, perceptual data such as images and videos. Rectified flow is a recent generative model formulation that connects data and noise in a straight line. Despite its better theoretical properties and conceptual simplicity, it is not yet decisively established as standard practice. In this work, we improve existing noise sampling techniques for training rectified flow models by biasing them towards perceptually relevant scales. Through a large-scale study, we demonstrate the superior performance of this approach compared to established diffusion formulations for high-resolution text-to-image synthesis. Additionally, we present a novel transformer-based architecture for text-to-image generation that uses separate weights for the two modalities and enables a bidirectional flow of information between image and text tokens, improving text comprehension, typography, and human preference ratings. We demonstrate that this architecture follows predictable scaling trends and correlates lower validation loss to improved text-to-image synthesis as measured by various metrics and human evaluations. Our largest models outperform state-of-the-art models, and we will make our experimental data, code, and model weights publicly available.

1. Introduction

The introduction motivates rectified flow as a theoretically favorable but not yet established alternative to diffusion, while emphasizing that forward-path choices affect sampling and learned generation. The work contributes improved rectified-flow noise samplers and a scalable bidirectional text-image architecture, with large models outperforming established open and closed-source systems.

  • Diffusion models generate new data by learning to invert paths from data toward random noise, proving effective for high-dimensional perceptual data.
  • Forward-path choices affect sampling, training-test alignment, artifacts, and the learned backward process.A forward process that fails to remove all noise can produce gray image samples.
  • Rectified flow connects data and noise on a straight line and has better theoretical properties, but remains insufficiently established in practice.Prior empirical demonstrations were mostly limited to small and medium-sized class-conditional experiments.
  • The proposed architecture replaces fixed text representations with learnable image and text streams that support bidirectional information flow and improve text-to-image modeling.The architecture is combined with an improved rectified-flow formulation and exhibits predictable validation-loss scaling, with lower loss correlating strongly with improved automated evaluations.
  • The work introduces new rectified-flow noise samplers and systematically studies diffusion and rectified-flow formulations for scalable text-to-image synthesis.Its largest models outperform SDXL, SDXL-Turbo, Pixart-α, and DALL-E 3 in prompt-understanding evaluation and human preference ratings.

2. Simulation-Free Training of Flows

The section formulates generative modeling as an ODE mapping noise to data and motivates simulation-free training by regressing vector fields instead of solving the ODE directly. Conditional flow matching makes this regression tractable, while time-dependent weighting changes optimization without changing the objective’s optimum.

  • ODE formulation: Generative models define an ODE mapping noise samples x1 from p1 to data samples x0 from p0, with velocity v parameterized by neural-network weights Θ.The velocity is written as vΘ(yt, t).
  • Simulation-free training: Directly solving the ODE with differentiable solvers is computationally expensive, especially for large networks parameterizing vΘ(yt, t).The method therefore avoids simulation during training.
  • Simulation-free training: Instead, the method directly regresses a vector field ut that generates a probability path between the data distribution p0 and noise distribution p1.The forward process is constructed so its marginals are consistent with p0 and p1.
  • Flow matching: A marginal vector field can be constructed from conditional vector fields, allowing the conditional dynamics to generate the marginal probability paths.This construction uses conditional vector fields ut(·|ϵ) to represent the marginal field ut.
  • Flow matching: Direct Flow Matching regression is intractable because of marginalization, whereas Conditional Flow Matching provides an equivalent tractable objective using conditional vector fields.The equivalence preserves the desired training solution while avoiding explicit marginalization.
  • Weighted objectives: Time-dependent weighting leaves the objective optimum unchanged but can alter the optimization trajectory, enabling a unified analysis of weighted losses and diffusion formulations.The objective can be reparameterized as a noise-prediction objective before introducing these weights.

3. Flow Trajectories

The section compares Rectified Flow with EDM, Cosine, and (LDM-)Linear diffusion trajectories, then introduces timestep-sampling strategies that emphasize intermediate steps where velocity prediction is more difficult. These strategies include logit-normal, mode sampling with heavy tails, and a CosMap mapping aligned with the cosine schedule.

  • Flow trajectories: Rectified Flows connect the data distribution and a standard normal distribution through straight paths, with the network directly parameterizing velocity vΘ.The RF loss trains this velocity uniformly over timesteps in [0, 1].
  • Flow trajectories: EDM, Cosine, and (LDM-)Linear define alternative forward processes using quantile-based, cosine, and variance-preserving schedules, respectively.Cosine weighting depends on the parameterization, while DDPM and LDM specify discrete diffusion coefficients βt.
  • Timestep sampling: Intermediate timesteps receive more emphasis because the RF target ϵ −x0 is more difficult to predict in the middle of [0, 1].Changing the timestep distribution from uniform U(t) to density π(t) is equivalent to using a weighted loss Lwπ.
  • Timestep sampling: Logit-normal sampling biases training toward data p0 or noise p1 through location parameter m and controls distribution width through scale parameter s.Negative m favors data p0, positive m favors noise p1, and s controls how wide the distribution is.
  • Timestep sampling: Mode sampling with heavy tails maintains strictly positive density on [0, 1], favoring the midpoint for positive s and endpoints for negative s.At s = 0, it reduces to the uniform weighting πmode(t; s = 0) = U(t).

4. Text-to-Image Architecture

The text-to-image architecture combines pretrained text and image representations in a DiT-based multimodal transformer. Separate modality-specific weights preserve distinct representation spaces while joint attention enables bidirectional interaction between text and image tokens.

  • Multimodal Diffusion Backbone: The backbone builds on DiT and uses timestep and pooled-text embeddings in its modulation mechanism, while richer text information is retained through token-level representations.The pooled text representation alone retains only coarse-grained information about the text input.
  • Multimodal Diffusion Backbone: The model encodes text and images, flattens 2 × 2 latent-image patches, concatenates both sequences, and applies modulated attention and MLP blocks.Image patches form a sequence of length 1/2 · h · 1/2 · w before concatenation with text embeddings.
  • Multimodal Diffusion Backbone: Separate weight sets for text and image modalities let each representation operate in its own space while joint attention incorporates information from the other modality.The design is equivalent to two independent transformers whose modality sequences are joined for attention.

5. Experiments

The experiments compare 61 flow and diffusion formulations under controlled training and evaluation conditions, finding that modified rectified-flow timestep sampling performs especially well. Further experiments show that the proposed architecture and scaling strategy improve convergence, evaluation metrics, and human preference.

  • Experimental setup: The study trains 61 formulations on ImageNet and CC12M, evaluating validation loss, CLIP, and FID for EMA and non-EMA weights.Runs are compared under controlled optimization, architecture, datasets, and samplers.
  • Formulation comparison: Across 24 control settings, non-dominated ranking shows rf/lognorm(0.00, 1.00) consistently performs well and outperforms uniform timestep sampling.Only rectified-flow variants with modified timestep sampling outperform the LDM-Linear baseline among evaluated formulations.
  • Formulation comparison: rf/lognorm(0.50, 0.60) ranks best with 50 sampling steps but reaches average rank 8.5 with 5 steps, demonstrating sensitivity to sampling settings.Rectified-flow formulations generally degrade less than other formulations when reducing the number of sampling steps.
  • Architecture comparison: MM-DiT significantly outperforms vanilla and cross-attention variants, while three parameter sets provide only a small gain over two.UViT initially learns faster than CrossDiT, but CrossDiT ultimately performs better than UViT.
  • Scaling results: Validation loss smoothly decreases with model size and training steps, correlates strongly with comprehensive metrics and human preference, and shows no saturation for image or video models.The largest model excels in most GenEval categories, outperforms DALL·E 3 overall, and the d = 38 model surpasses current proprietary and open models in human preference evaluation.

6. Conclusion

The work analyzes scaling in rectified flow models for text-to-image synthesis, introducing improved timestep sampling and demonstrating a multimodal transformer architecture. These contributions improve upon prior diffusion training formulations while preserving favorable few-step sampling properties.

  • The study presents a scaling analysis of rectified flow models for text-to-image synthesis.
  • The proposed timestep sampling improves over previous diffusion training formulations for latent diffusion models.
  • The timestep sampling retains rectified flow’s favorable properties in the few-step sampling regime.
  • The transformer-based MM-DiT architecture accounts for the multimodal nature of text-to-image synthesis.

Broader Impact

The work aims to advance machine learning broadly and image synthesis specifically, while acknowledging potential societal consequences without highlighting any particular one.

  • The paper seeks to advance machine learning and image synthesis, acknowledges many potential societal consequences, and refers readers to Po et al. (2023) for broader discussion of diffusion-model ramifications.The authors do not identify a specific consequence as requiring emphasis.

Supplementary · A. Background · B. On Flow Matching

The supplementary background positions diffusion models as the standard approach for image and video generation, while introducing rectified flows as an ODE-based alternative. It also motivates transformer architectures for scaling diffusion models and outlines rectified flows’ relationship to continuous normalizing flows.

  • A. Background: Diffusion models generate data by approximating the reverse ODE of a stochastic process that transforms data into noise.
  • A. Background: Diffusion models have become the standard approach for generative modeling of images and videos.
  • B. On Flow Matching: Rectified flow models construct a transport map between two distributions through an ordinary differential equation.
  • B. On Flow Matching: Rectified flows are closely connected to continuous normalizing flows and diffusion models.
  • B. On Flow Matching: Unlike continuous normalizing flows, rectified flows and stochastic interpolants do not require ODE simulation during training.
  • A. Background: Transformers are established for scaling in natural-language processing and computer vision, whereas U-Net architectures have dominated diffusion-model backbones.

B.1. Details on Simulation-Free Training of Flows · B.2. Details on Image and Text Representations

The appendix establishes simulation-free flow training through the continuity equation and the equivalence of flow-matching objectives. It also specifies latent image and text representations, including downsampled, normalized image latents and concatenated pretrained text-model features.

  • B.1. Details on Simulation-Free Training of Flows: The continuity equation is a necessary and sufficient condition for a velocity field v_t to generate the probability-density path p_t.The condition is stated as d/dt p_t(x) + ∇·[p_t(x)v_t(x)] = 0.
  • B.1. Details on Simulation-Free Training of Flows: The derivation shows that the conditional velocity field u_t(z|ϵ) generates the conditional density path p_t(z|ϵ), using the continuity equation and the defined interpolation.These steps establish the sufficient condition needed for the simulation-free construction.
  • B.1. Details on Simulation-Free Training of Flows: The flow-matching and conditional-flow-matching objectives are equivalent, up to terms independent of the model parameters Θ.The appendix denotes this equivalence as L_FM ⇋ L_CFM.
  • B.2. Details on Image and Text Representations: RGB images X ∈ R^{H×W×3} are encoded by a pretrained autoencoder into smaller latent representations x = E(X) ∈ R^{h×w×d}.The spatial downsampling factor is 8, giving h = H/8 and w = W/8.
  • B.2. Details on Image and Text Representations: The forward process and flow sampling operate in latent space, after which sampled latents are decoded to pixels and normalized using globally computed training-data statistics.Latents are normalized by their mean and standard deviation over a subset of the training data.
  • B.2. Details on Image and Text Representations: Text conditioning is encoded with pretrained, frozen models combining CLIP, OpenCLIP, and a T5-v1.1-XXL encoder-decoder text model.The specified CLIP encoders are CLIP L/14 and OpenCLIP bigG/14.
  • B.2. Details on Image and Text Representations: At depth d = 22, the FID gap between 8-chn and 16-chn autoencoder spaces becomes negligible, while the 16-channel space requires more capacity at smaller model sizes.The authors select the 16-channel model for scaling to much larger models.
  • B.2. Details on Image and Text Representations: The CLIP and T5 representations are dimension-matched, concatenated along the sequence axis, and combined into a final context representation c_ctxt ∈ R^{154×4096}.The resulting caption representations c_vec and c_ctxt are used in different ways.

B.3. Preliminaries for the Experiments in Section 5.1.

The experiments use captioned ImageNet and CC12M data, standardized optimization and evaluation procedures, and stratified loss and sampler analyses. Models are assessed on COCO-2014 validation data and generated samples under varied guidance and sampling-step settings.

  • Optimization: All models are trained with AdamW, a global batch size of 1024, a 10^-4 learning rate, and 1000 linear warmup steps.Training uses mixed precision and maintains exponential moving average weights updated every 100 batches with decay 0.99.
  • Evaluation: Models are evaluated during training using CLIP scores, FID, and validation losses on the COCO-2014 validation split.These metrics are monitored regularly as described for the Section 5.1 experiments.
  • Evaluation: Losses are evaluated at eight equally spaced timestep values in [0, 1] because their magnitude and variance vary substantially across timesteps.This stratified evaluation is intended to make timestep-dependent loss behavior comparable.
  • Sampling analysis: For sampler analysis, 1000 samples per sampler are generated under varying guidance scales and sampling-step counts, using Euler discretization.Samples are evaluated with CLIP L/14 scores and FID computed from CLIP L/14 image features against the validation images.

B.4. Improving SNR Samplers for Rectified Flow Models

The section introduces novel timestep densities for training rectified flow models, including mode and logit-normal samplers. The logit-normal sampler outperforms uniform rectified flow and established diffusion baselines.

  • Sampler design: Novel timestep densities bias rectified flow training toward selected sampling distributions, including the mode and logit-normal samplers.Figure 11 visualizes both distributions for biasing training-timestep sampling.
  • Performance: The logit-normal sampler outperforms the classic uniform rectified flow formulation and established diffusion baselines including EDM and LDM-Linear.The cited baselines are uniform rectified flow, EDM, and LDM-Linear.
  • Sampler design: The explored samplers bias the distribution of training timesteps rather than changing the rectified flow formulation itself.The mode and logit-normal distributions are presented specifically as mechanisms for biasing training-timestep sampling.

C. Direct Preference Optimization · D. Finetuning for instruction-based image editing

The paper evaluates Direct Preference Optimization on 2B and 8B base models and finds that DPO improves aesthetic quality, spelling, prompt following, and general quality. It also applies channel-wise input-target latent concatenation to instruction-based image editing, achieving results despite no text-manipulation training data, unlike an SDXL-based reproduction.

  • C. Direct Preference Optimization: DPO-finetuned models generally produce more aesthetically pleasing samples with better spelling than the base models.Figure 13 compares base and DPO-finetuned outputs, including lakeside, sauropod, and typography prompts.
  • C. Direct Preference Optimization: DPO is applied to the 2B and 8B parameter base models for preference optimization.The method follows Wallace et al. (2023) and avoids finetuning the entire model by introducing learnable Lo…
  • D. Finetuning for instruction-based image editing: Instruction-based image editing concatenates input and target latents along the channel dimension before patching.This follows a common approach for instruction-based editing and image-to-image diffusion models, adapted here to the proposed model.
  • C. Direct Preference Optimization: Human evaluators prefer DPO-finetuned models for both prompt following and general quality.This preference is reported in the human evaluation comparing base and DPO-finetuned models.
  • D. Finetuning for instruction-based image editing: The proposed model supports instruction-based image editing even though the training data contained no text-manipulation tasks.The result demonstrates applicability beyond the tasks explicitly represented in the training data.
  • D. Finetuning for instruction-based image editing: The authors could not reproduce similar editing results when training an SDXL-based model on the same data.This comparison uses an SDXL-based editing model as the baseline.

E. Data Preprocessing for Large-Scale Text-to-Image Training … E.3. Assessing the Efficacy of our Deduplication Efforts

The paper preprocesses frozen encoder outputs to reduce training-time memory and compute, while accepting reduced augmentation flexibility and increased storage costs. It prevents and evaluates memorization by removing near-duplicate training images with SSCD-based clustering and testing extracted generations before and after deduplication.

  • E.1. Precomputing Image and Text Embeddings: Precomputing autoencoder latents and text encoder representations removes encoder GPU requirements and skips forward encoding during training after the first epoch.The representations are constant during training, reducing memory and saving time and total compute after precomputation.
  • E.1. Precomputing Image and Text Embeddings: Precomputation prevents per-epoch random augmentation and requires square-center cropping, or resizing and cropping into aspect-ratio buckets for higher-resolution finetuning.Dense text-encoder outputs also increase storage costs and loading times.
  • E.2. Preventing Image Memorization: The training dataset is scanned for duplicated examples and deduplicated to reduce the risk that trained models produce verbatim copies of images.The motivation is that memorization of training samples can create several issues in generative image models.
  • E.2. Preventing Image Memorization: SSCD provides the backbone for large-scale near-duplicate detection, generating image embeddings suitable for clustering and downstream tasks.The deduplication procedure follows methods from Carlini et al. (2023) and Somepalli et al. (2023a), while the number of clusters N follows Nichol (2022).
  • E.2. Preventing Image Memorization: autofaiss and FAISS index-factory functionality enable efficient clustering of high-dimensional image embeddings with a predefined number of centroids.The implementation trains a custom index for large-scale clustering.
  • E.2. Preventing Image Memorization: Different SSCD thresholds are evaluated for their data-removal effects, and four thresholds are selected for the final deduplication run.Algorithm 1 describes the deduplication approach, with results shown in Figure 16b and the final selection in Figure 16a.
  • E.3. Assessing the Efficacy of our Deduplication Efforts: Memorization is assessed by extracting samples from small purpose-trained models and comparing generations before and after SSCD-based deduplication.The procedure generates examples from known prompts and applies membership inference to distinguish novel generations from memorized ones.
  • E.3. Assessing the Efficacy of our Deduplication Efforts: 350,000 most-duplicated examples selected with an SSCD threshold of 0.5 each receive 500 candidate generations per text prompt to increase the likelihood of detecting memorization.The detection procedure compares generated samples across random seeds using similarity and identifies sufficiently large cliques as memorized.
Loading 2403.03206v1…