PaperlayerView exact revision on arXiv

Source-linked AI summary

Calibri: Enhancing Diffusion Transformers via Parameter-Efficient Calibration

Danil Tokhchukov, Aysel Mirzoeva, Andrey Kuznetsov, Konstantin Sobolev

arXiv:2603.24800v2cs.CV

TL;DR

DiT blocks have uneven contributions, motivating a need to determine whether their standard weighting is optimal. Calibri calibrates DiT components by optimizing roughly 102 scaling coefficients with CMA-ES and achieves consistent gains across text-to-image models while reducing inference steps. Its effectiveness depends partly on the quality of the reward models used to select coefficients.

  • Problem

    Prior work indicates that DiT blocks contribute unevenly, leaving the optimal weighting of these components an open calibration question.

  • Method

    Calibri frames DiT calibration as black-box reward optimization over ∼102 parameters and uses the gradient-free CMA-ES evolutionary strategy to search scaling coefficients.

  • Results

    Calibri consistently improves performance across Flux, SD-3.5M, and Qwen-Image while using 15 inference steps instead of 30, 40, and 50, respectively.

  • Takeaways & Limitations

    Calibri provides a parameter-efficient way to improve generative quality while reducing inference-step requirements and computational costs.

  • Takeaways & Limitations

    Reward-model shortcomings can cause Calibri to select a suboptimal set of calibration coefficients.

Abstract

from arXiv · show

In this paper, we uncover the hidden potential of Diffusion Transformers (DiTs) to significantly enhance generative tasks. Through an in-depth analysis of the denoising process, we demonstrate that introducing a single learned scaling parameter can significantly improve the performance of DiT blocks. Building on this insight, we propose Calibri, a parameter-efficient approach that optimally calibrates DiT components to elevate generative quality. Calibri frames DiT calibration as a black-box reward optimization problem, which is efficiently solved using an evolutionary algorithm and modifies just ~100 parameters. Experimental results reveal that despite its lightweight design, Calibri consistently improves performance across various text-to-image models. Notably, Calibri also reduces the inference steps required for image generation, all while maintaining high-quality outputs.

1. Introduction

The paper argues that DiT blocks contribute unevenly and that post-hoc calibration can improve generation quality. Calibri uses lightweight black-box optimization to calibrate these components and reduce inference steps.

  • DiT blocks contribute unevenly to final generation, with some blocks whose removal significantly changes model outputs.
  • Selective block disabling can improve generation quality, indicating that some blocks may introduce detrimental artifacts.
  • A single learned scalar that re-weights each block’s output consistently enhances performance over the original model.
  • Calibri formulates calibration as black-box reward optimization over only ∼102 parameters and solves it with the gradient-free CMA-ES evolutionary strategy.
  • Calibri Ensemble combines multiple calibrated models, while Calibri reduces inference steps and improves efficiency and generation quality.

2. Related Work

Related work traces diffusion backbones from U-Net architectures toward scalable DiT-based designs and multimodal transformer variants. It also covers interpretability of diffusion components and alignment with human preferences.

  • Diffusion Models Backbones: Early diffusion models used U-Net backbones with residual blocks, pixelwise self-attention, and cross-attention for text-image conditioning.
  • Diffusion Models Backbones: Recent diffusion research has shifted toward DiT architectures, motivated in part by transformer scalability.
  • Diffusion Models Backbones: MM-DiT processes textual and visual inputs with distinct transformers before combining their sequences through unified attention operations.
  • Diffusion Model Backbone Interpretability: Interpretability studies use cross-attention maps to predict spatial locations of textual concepts for applications including image editing and layout control.
  • Visual Generative Model Alignment: Human-feedback alignment methods for diffusion and rectified-flow models include reward models, reward backpropagation, DPO, DDPO, and GRPO.

3. Method

Calibri analyzes uneven DiT-block contributions and calibrates them with lightweight scaling parameters optimized through reward-based evolutionary search. The method supports block, layer, gate, and ensemble calibration for standard and multimodal DiT architectures.

  • 3.1. Preliminaries: Standard DiT blocks combine Multi-Head Self-Attention and feed-forward layers, with LayerNorm and time-embedding modulation through α, β, and γ vectors.The block output is formed through residual updates around the attention and feed-forward operations.
  • 3.1. Preliminaries: MM-DiT processes concatenated textual and visual tokens in parallel, using separate modulation vectors and restricting inter-modal communication to multimodal attention.The visual and textual streams use distinct α, β, and γ parameters.
  • 3.2. Motivation: Layer ablations on 64 prompts, using FLUX images and five random seeds, show that removing some DiT layers can improve Image Reward rather than degrade it.Each layer is bypassed by setting its residual-output scaling γ to 0.
  • 3.2. Motivation: For every DiT block, some output scaling factor s improves performance over the original configuration, with s = 0 denoting ablation and s = 1 the original model.The tested coefficients are s ∈ {0, 0.25, 0.5, 0.75, 1.25, 1.5}.
  • 3.3. Calibri: Calibri calibrates a minimal parameter subset, framing coefficient selection as reward maximization over approximately 102 parameters.The calibrated model may adjust output-level weights and internal-layer parameters.
  • 3.3. Calibri: Block, layer, and gate scaling provide progressively finer calibration, with gate scaling targeting separate visual and textual pathways in MM-DiT.Block scaling shares one coefficient across attention and MLP outputs, whereas layer scaling uses distinct coefficients.
  • 3.3. Calibri: CMA-ES performs gradient-free search by sampling candidate coefficients from a Gaussian distribution and updating its mean and covariance toward successful directions.This procedure iteratively balances exploration and exploitation of the calibration search space.
  • 3.4. Calibri Ensemble: Calibri can calibrate an ensemble of N models simultaneously, leveraging model diversity for enhanced generative performance and robustness.The framework also applies to classifier-free guidance with N = 2 conditional and unconditional models.

4. Experiments

Experiments evaluate Calibri across multiple DiT-based text-to-image models, calibration granularities, ensemble settings, inference budgets, and alignment checkpoints. Calibri generally improves reported quality metrics while reducing inference steps, and remains effective when combined with alignment methods.

  • Experimental setup: Experiments cover FLUX.1-dev, SD-3.5M, Qwen-Image, and an SD-3.5M checkpoint fine-tuned with Flow-GRPO.Evaluation uses HPSv3, Q-Align, and ImageReward, with HPDv3 prompts for final metrics.
  • Calibration design decisions: Layer scaling provides more consistent improvements across multiple reward functions than gate scaling, despite gate scaling achieving the highest HPSv3 value.The three scaling schemes have relatively similar resulting performance but substantially different training speeds.
  • Ensembling: Calibri Ensemble consistently increases HPSv3 across inference steps and shifts the optimal sampling range from 30–50 baseline steps to 10–15 steps.This reduces the number of function evaluations required for strong performance.
  • Different backbones: Calibri improves generation metrics across Flux, SD-3.5M, and Qwen-Image while using 15 inference steps instead of 30, 40, and 50, respectively.Qualitative comparisons also report superior output quality across the baseline models.
  • Human evaluation: A 200-user study with 5,600 assessments found that evaluators preferred Calibri on Overall Preference and Text Alignment for Flux.1-dev and Qwen-Image.The calibrated models were also reported to be 2–3.3× faster than baselines.
  • Calibration cost: Calibration costs range from 32 to 356 NVIDIA H100 GPU-hours and are incurred once offline.Flux block calibration takes 32 H100 GPU-hours and yields an approximately 2× permanent inference speed-up.

5. Conclusion

Calibri improves DiT generative performance through lightweight calibration and consistently enhances quality across text-to-image diffusion models while reducing inference steps.

  • Calibri calibrates DiT components using a parameter-efficient approach that modifies approximately 102 parameters.The method uses CMA-ES to optimize calibration coefficients in a black-box reward-optimization framework.
  • Calibri Ensemble combines calibrated models at inference time to further enhance results.
  • Calibri consistently improves generation quality across a range of text-to-image diffusion models.
  • Calibri reduces the number of inference steps required for image generation while retaining high-quality outputs.The paper presents this efficiency improvement as relevant to computationally constrained applications.

Supplementary Material Structure

The supplementary document organizes additional analyses around methodology limitations, generation diversity, reward-model choice, and CMA-ES-based parameter search.

  • Section A analyzes limitations of the proposed methodology.
  • Section B compares generated-image diversity before and after incorporating Calibri.
  • Section C explains the rationale for the chosen reward model and discusses its impact on system performance.
  • Section D discusses why CMA-ES is used for parameter search and justifies its effectiveness.

A. Limitations

Calibri relies on reward models whose limited sensitivity to anatomical artifacts can lead to suboptimal calibration coefficients, motivating future improvements in reward modeling.

  • Reward models approximate user preferences for generated images and serve as Calibri’s optimization objective.
  • Current reward models often insufficiently detect anatomical inconsistencies such as extra limbs and distorted fingers.These limitations are illustrated in Figure 9.
  • Reward-model shortcomings can cause Calibri to select a suboptimal set of calibration coefficients.
  • Further advances in reward modeling are expected to improve artifact sensitivity and Calibri’s future performance.

B. Generation diversity

Calibri preserves generation diversity while reducing inference steps and improving generation quality, unlike Flow-GRPO in the reported comparison.

  • Calibri is evaluated for its effect on generation diversity because reward-model optimization can reduce diversity.
  • Calibri achieves diversity comparable to SD-3.5M while using 15 rather than 40 inference timesteps.
  • The Calibri-optimized model has significantly higher generation quality than the original model despite reduced inference time.
  • Flow-GRPO reduces generation diversity from 0.20 to 0.15 and does not accelerate inference time.
  • Applying Calibri to a model already optimized by Flow-GRPO produces no further change in generation diversity.

C. Different Rewards

Calibri is evaluated with multiple reward models as optimization objectives, with HPSv3 producing the strongest overall quality gains and PickScore also performing strongly. Improvements extend beyond the metric directly optimized, supporting Calibri as a general-purpose calibration method rather than reward hacking.

  • Calibri experiments evaluate layer-scale calibration on Flux across different reward models.The quantitative comparison is summarized in Table 7 and illustrated in Figure 10.
  • HPSv3 calibration achieves the most significant quality improvement across all metrics, while PickScore shows similarly strong performance.
  • The strongest reward model improves metrics beyond its directly targeted objective, indicating that Calibri is not tailored to reward hacking.The reported cross-metric gains motivate its characterization as a general-purpose technique for improving generation quality.

D. CMA-ES vs gradient-based parameter search

The paper compares CMA-ES with Flow-GRPO for optimizing Calibri layer scales on FLUX. CMA-ES is reported as more efficient, with coefficient convergence indicating when training can stop.

  • The comparison addresses gradient-based alignment challenges caused by reward models operating across noisy latent spaces and generated images.These challenges necessitate repeated inference steps for accurate reward computation.
  • CMA-ES is substantially more efficient than Flow-GRPO when training Calibri layer scales on FLUX.Both optimizers are evaluated on T2I-Compbench++ test prompts under the main experimental setup.
  • Calibri coefficients converge effectively under CMA-ES, with stabilized sigma and a plateau in the training curve.Training can be terminated after this convergence behavior appears to reduce computational resource use without compromising performance.
Loading 2603.24800v2…