PaperlayerView exact revision on arXiv

Source-linked AI summary

One Transformer Fits All Distributions in Multi-Modal Diffusion at Scale

Fan Bao, Shen Nie, Kaiwen Xue, Chongxuan Li, Shi Pu, Yaole Wang, Gang Yue, Yue Cao, Hang Su, Jun Zhu

arXiv:2303.06555v2cs.LGcs.CV

TL;DR

Existing multimodal generative models are typically bespoke and limited to a single task. UniDiffuser unifies marginal, conditional, and joint diffusion modeling in one model through modality-specific timesteps, achieving strong performance across generation tasks and comparisons with specialized systems.

  • Problem

    Existing generative models are bespoke systems that allow only a single multimodal generation task, motivating a unified framework covering varied conditioning types.

  • Method

    UniDiffuser unifies diffusion learning for marginal, conditional, and joint distributions by predicting perturbed-data noise with independently configurable timesteps for each modality.

  • Results

    UniDiffuser outperforms Versatile Diffusion across text-to-image and image-to-text metrics and guidance scales, while its text-to-image FID is comparable to Stable Diffusion and outperforms DALL·E 2.

  • Takeaways & Limitations

    One UniDiffuser model supports multiple multimodal generation tasks with minimal modification to diffusion models and no additional training or inference overhead.

  • Takeaways & Limitations

    The unified marginal-generation interpretation assumes that sufficiently late noisy data are negligibly different from standard Gaussian noise.

Abstract

from arXiv · show

This paper proposes a unified diffusion framework (dubbed UniDiffuser) to fit all distributions relevant to a set of multi-modal data in one model. Our key insight is -- learning diffusion models for marginal, conditional, and joint distributions can be unified as predicting the noise in the perturbed data, where the perturbation levels (i.e. timesteps) can be different for different modalities. Inspired by the unified view, UniDiffuser learns all distributions simultaneously with a minimal modification to the original diffusion model -- perturbs data in all modalities instead of a single modality, inputs individual timesteps in different modalities, and predicts the noise of all modalities instead of a single modality. UniDiffuser is parameterized by a transformer for diffusion models to handle input types of different modalities. Implemented on large-scale paired image-text data, UniDiffuser is able to perform image, text, text-to-image, image-to-text, and image-text pair generation by setting proper timesteps without additional overhead. In particular, UniDiffuser is able to produce perceptually realistic samples in all tasks and its quantitative results (e.g., the FID and CLIP score) are not only superior to existing general-purpose models but also comparable to the bespoken models (e.g., Stable Diffusion and DALL-E 2) in representative tasks (e.g., text-to-image generation).

1. Introduction

Existing multimodal generative models typically support single tasks, motivating a unified framework that fits multiple distributions in one model. UniDiffuser uses modality-specific timesteps and a transformer to perform diverse generation tasks without additional overhead.

  • Motivation: Multimodal generative models are often bespoke systems limited to a single task, whereas a general system should cover arbitrary conditioning types.Examples include text-conditioned image creation and image captioning.
  • UniDiffuser: UniDiffuser explicitly fits all relevant distributions in one diffusion model without additional training or inference overhead.Its central insight is to unify marginal, conditional, and joint modeling as noise prediction at modality-specific perturbation levels.
  • UniDiffuser: The framework modifies diffusion by perturbing every modality, providing separate timesteps, and predicting noise for every modality.A zero timestep conditions on a modality, while a maximum timestep ignores it for unconditional generation.
  • Architecture: A transformer backbone handles inputs from different modalities and captures their interactions.The architecture is motivated by transformer applications in multimodal generative modeling.
  • Capabilities: UniDiffuser performs image, text, text-to-image, image-to-text, and image-text pair generation by setting appropriate timesteps.The implementation uses latent-space diffusion with a CLIP image encoder and GPT-2 text decoder on large-scale paired image-text data.
  • Results: UniDiffuser produces perceptually realistic samples across tasks, with quantitative results superior to existing general-purpose models and comparable to bespoke models on representative tasks.The introduction highlights FID and CLIP score results, including text-to-image generation.

2. Background

Diffusion models generate data by adding noise through a forward process and learning to reverse it with a noise-prediction network. Conditional diffusion extends this procedure by conditioning the reverse process on paired data.

  • Unconditional diffusion: Diffusion models gradually inject noise into data through a Markov chain.The perturbed sample is formed using a noise schedule that controls the corruption level.
  • Noise prediction: The noise schedule uses βt for the injected noise level and αt = 1 − βt for the retained signal component.These coefficients determine how x0 and Gaussian noise combine to form xt.
  • Unconditional diffusion: The reverse diffusion process is approximated by Gaussian transitions whose mean depends on the noisy input.The reverse model reconstructs less-noisy samples from progressively noisier states.
  • Noise prediction: The noise-prediction network estimates the conditional expectation of injected noise from the perturbed sample and timestep.Under the regression objective, the optimal predictor equals E[ϵx|xt].
  • Conditional diffusion: Conditional generation models q(x0|y0) using paired data and a reverse Gaussian process whose noise predictor is conditioned on y0.The conditional predictor estimates E[ϵx|xt, y0].
  • Classifier-free guidance: Classifier-free guidance improves conditional sample quality by linearly combining conditional and unconditional predictions.The models share parameters through a null token, and s denotes the guidance scale.

3. Method

UniDiffuser uses one joint diffusion model to capture marginal, conditional, and joint distributions by varying modality-specific timesteps. Its simple noise-prediction objective supports multiple generation tasks and free classifier-free guidance without additional training overhead.

  • UniDiffuser targets marginal, conditional, and joint distributions from two-modal data within one diffusion model.
  • Different modality-specific timesteps unify the required noise expectations, with maximum noise marginalizing a modality and zero noise conditioning on it.The model learns E[ϵx, ϵy|xtx, yty] across all timestep pairs.
  • A joint noise-prediction network perturbs and predicts noise for both modalities using independently sampled timesteps.The objective concatenates the perturbed modalities and trains on their Gaussian noise.
  • One parameter update requires a single forward-backward calculation for multiple tasks, matching DDPM efficiency despite slightly higher gradient variance.The authors do not observe slower convergence from the two independent timesteps.
  • Proper timestep settings make UniDiffuser’s inference procedure the same as bespoke models while covering unconditional, conditional, and joint sampling.
  • Classifier-free guidance applies directly to conditional and joint sampling because UniDiffuser already models marginal distributions and needs no null token.The joint model can be viewed through conditional scores when both modalities share the same timestep.

4. UniDiffuser on Images and Texts

UniDiffuser is implemented for image-text data through a two-stage latent-space pipeline: modality-specific encoding and transformer-based joint diffusion training. The design combines reconstruction and semantic embeddings for images with compressed CLIP text representations and GPT-2 decoding.

  • The implementation first encodes images and texts into continuous latent embeddings, then trains a transformer-parameterized UniDiffuser on those embeddings.
  • Image encoder-decoder: Image embeddings concatenate Stable Diffusion autoencoder features with 512-dimensional CLIP semantic features.The autoencoder component supports reconstruction, while the CLIP component supports image semantics for image-to-text generation.
  • Text encoder-decoder: The text encoder uses Stable Diffusion’s CLIP representation, reduces each of 77 vectors from 768 to 64 dimensions, and supplies the result to a GPT-2 decoder.GPT-2 receives the text latent as a prefix embedding and reconstructs text autoregressively.
  • Image and text latent embeddings are concentrated within approximately [−2, 2] with comparable mean and variance ranges, so no additional normalization is applied.The reported means are 0.0269 for images and 0.0127 for text; standard deviations are 0.7919 and 0.5957, respectively.
  • Transformer backbone: The joint noise-prediction network uses a modified U-ViT transformer backbone to handle inputs from different modalities.U-ViT treats data, conditions, and timesteps as tokens and uses long skip connections.

5. Related Work

Prior multi-modal generative systems commonly learn conditional distributions, while Versatile Diffusion extends multitask generation through a multi-flow architecture. UniDiffuser instead uses a unified diffusion formulation with a single training loss and supports broader generation capabilities.

  • Prior multi-modal generative work commonly models conditional distributions for applications including text-to-image, text-to-video, text-to-speech, and image-to-text generation.
  • Versatile Diffusion uses a multi-flow architecture trained in a traditional multitask framework, requiring multiple feed-forward passes and tuned gradient multipliers.
  • UniDiffuser replaces this multitask setup with a single training loss and one forward-backward pass per update.
  • UniDiffuser handles more tasks than Versatile Diffusion, including joint sampling, without complex training tricks.
  • UniDiffuser outperforms Versatile Diffusion in image-to-text and text-to-image generation on the reported FID and CLIP scores.
  • Transformers are used to fuse modalities in representation learning, while diffusion models have demonstrated high-quality generation across images, audio, video, point clouds, and molecular conformations.

6. Experiments

The experiments evaluate UniDiffuser on large-scale image-text data across generation tasks and compare it with general-purpose and bespoke systems. UniDiffuser consistently outperforms Versatile Diffusion, performs comparably to Stable Diffusion and better than DALL·E 2 on text-to-image generation, and supports variation and joint-generation applications.

  • 6. Experiments: Experiments assess multiple generation tasks, compare UniDiffuser with existing large models, and demonstrate data variation, blocked Gibbs sampling, and image interpolation.
  • 6.1. Setup: The evaluation uses three LAION-5B subsets, including approximately 2B English image-text pairs, 170M high-resolution multilingual pairs, and 193M filtered high-quality pairs.
  • 6.1. Setup: Figure 5 compares UniDiffuser and Versatile Diffusion in text-to-image generation using matched classifier-free guidance scales, with CLIP score on the horizontal axis and FID on the vertical axis.
  • 6.1. Setup: Figure 6 compares the models in image-to-text generation across matched classifier-free guidance scales using CLIP score on the vertical axis.
  • 6.1. Setup: Table 1 reports zero-shot FID on MS-COCO, using classifier-free guidance scale 3 for UniDiffuser and Versatile Diffusion.
  • 6.2. Main Results: UniDiffuser consistently outperforms Versatile Diffusion across both text-to-image and image-to-text tasks, all reported metrics, and all guidance scales.
  • 6.2. Main Results: On text-to-image generation, UniDiffuser is comparable to Stable Diffusion and outperforms DALL·E 2 despite supporting multiple generation tasks.
  • 6.2. Main Results: Examples demonstrate joint, conditional, and unconditional generation, plus image and text variation, blocked Gibbs sampling, and interpolation between images in the wild.

7. Conclusion

UniDiffuser provides one general-purpose framework for multiple multimodal generation tasks, while its current implementation has limitations in text smoothness and future-work scope.

  • 7. Conclusion: The implementation’s generated text is not yet smooth, mainly because the text data is noisy.The authors identify smoother text generation as a current limitation rather than a resolved property.
  • 7. Conclusion: Semi-supervised learning and extension to more modalities are identified as future work.These capabilities are presented as potential extensions rather than results established by the current implementation.
  • 7. Conclusion: UniDiffuser fits multiple generation tasks with one transformer network, enabling simultaneous improvements across tasks through scaling and optimization.The paper identifies this shared-network design as easier to improve and maintain under large-scale pretraining.
  • 7. Conclusion: The authors warn that large-scale multimodal generation may have consequences such as deepfakes and state that sampled images are watermarked.They also plan a systematic protocol to mitigate this issue before releasing the code and model.
  • 7. Conclusion: The reported samples cover joint, conditional, unconditional, variation, blocked Gibbs, and interpolation generation tasks.Examples include image-text joint generation, text-to-image, image-to-text, unconditional generation, modality variation, blocked Gibbs sampling, and image interpolation.

B. The Training and Sampling Algorithms

UniDiffuser trains by sampling paired data, modality-specific timesteps, and Gaussian noise, then perturbing the modalities for noise-prediction learning. Its conditional sampler is described as analogous to sampling the other modality given the first.

  • B. The Training and Sampling Algorithms: Training samples paired data, independent timesteps for each modality, and Gaussian noise before constructing perturbed variables.The algorithm initializes x0, y0 from q(x0, y0), samples tx and ty uniformly, and draws εx and εy from N(0, I).
  • B. The Training and Sampling Algorithms: The perturbed x modality is formed as xtx = √αtxx0 + √1 − αtxϵx.This step explicitly injects Gaussian noise according to the sampled timestep tx.
  • B. The Training and Sampling Algorithms: The conditional sampling algorithm generates x0 conditioned on y0 and states that sampling y0 conditioned on x0 is analogous.The supplied algorithm begins from xT sampled from N(0, I), iterates backward over timesteps, and returns x0.

C. Summary of Classifier-Free Guidance Models

The appendix organizes separate unconditional, conditional, and joint samplers and notes that classifier-free guidance models are summarized for conditional and joint generation.

  • C. Summary of Classifier-Free Guidance Models: UniDiffuser provides separate procedures for unconditional sampling of x0, conditional sampling of x0 given y0, and joint sampling of x0 and y0.The listed algorithms distinguish these sampling settings explicitly.
  • C. Summary of Classifier-Free Guidance Models: Unconditional x0 sampling starts from xT ∼ N(0, I), iterates from T down to 1, and returns x0.At each step, Gaussian variables are sampled when t > 1, with the terminal noise set to zero.
  • C. Summary of Classifier-Free Guidance Models: Joint sampling starts both modalities from standard Gaussian noise and returns the denoised pair x0, y0.The algorithm samples xT and yT, iterates backward jointly, and suppresses terminal sampling noise.
  • C. Summary of Classifier-Free Guidance Models: Classifier-free guidance is available for conditional and joint generation, while unconditional-sampling models are also included.The guidance scale s is identified in the table caption as the classifier-free guidance scale.
  • C. Summary of Classifier-Free Guidance Models: The joint noise-prediction network is implemented with a U-ViT.Its detailed configuration is referenced in Table 3.

E. Details of the GPT-2 Text Decoders

The text decoder compresses CLIP text representations into embeddings that GPT-2 reconstructs autoregressively, while interpolation uses spherical interpolation between modality embeddings and diffusion solvers.

  • E. Details of the GPT-2 Text Decoders: The GPT-2 decoder receives y0 as a prefix embedding and reconstructs the original text autoregressively.The text reconstruction pipeline is illustrated in Figure 17 and described as using a finetuned GPT-2 decoder.
  • E. Details of the GPT-2 Text Decoders: The text embedding is produced as y0 = Linear(CLIP(T)), with CLIP parameters frozen while the linear layer and GPT-2 are trained or finetuned.The CLIP output is projected to a lower-dimensional embedding before reconstruction.
  • E. Details of the GPT-2 Text Decoders: The 124M-parameter GPT-2 decoder is finetuned on 2.3B image-text pairs from LAION-2B-en.Training uses 235K steps with batch size 768; generation uses beam search with beam size 5 and maximum length 67.
  • E. Details of the GPT-2 Text Decoders: A 64-dimensional text embedding achieves BLEU-1 0.969 and BLEU-4 0.894 for reconstruction on the MS-COCO test set.The paper reports that the input texts are reconstructed very well in Figure 18.

G. Comparison of Examples

UniDiffuser produces examples that align better with text and images than Versatile Diffusion, while using fewer resources than comparable models. It also supports five generation tasks with only 10% extra parameters over Stable Diffusion.

  • Example Comparisons: UniDiffuser samples align better with texts than Versatile Diffusion on text-to-image generation.This comparison is shown in Figures 19 and 20.
  • Example Comparisons: UniDiffuser samples align better with images than Versatile Diffusion on image-to-text generation.This comparison is shown in Figure 21.
  • Efficiency Comparison: UniDiffuser is more efficient than Stable Diffusion and Versatile Diffusion in inference time and memory.The comparison uses generation of 10 samples with 25 denoising steps on one A100 80GB GPU.
  • Efficiency Comparison: 10% extra parameters let UniDiffuser support five generation tasks compared with Stable Diffusion at comparable training cost.The tasks are image, text, text-to-image, image-to-text, and image-text pair generation.
  • Efficiency Comparison: UniDiffuser has fewer parameters than Versatile Diffusion while achieving superior results.The passage reports this comparison against the general-purpose model Versatile Diffusion.
Loading 2303.06555v2…