Calibri: Enhancing Diffusion Transformers via Parameter-Efficient Calibration

Abstract: 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.

• The paper demonstrates that calibrating Diffusion Transformer blocks using only ~10^2 scalar parameters significantly enhances generative image quality.
• It employs CMA-ES for efficient black-box optimization, revealing uneven block contributions and achieving measurable gains in metrics like HPSv3.
• The ensemble calibration approach boosts model robustness and accelerates inference, delivering competitive performance with minimal computational overhead.

Calibri: Parameter-Efficient Calibration for Diffusion Transformers

Introduction and Motivation

The paradigm shift from U-Net-based backbones to Diffusion Transformers (DiT) has materially altered the landscape of visual content generation, encompassing tasks from text-to-image synthesis to multimodal and video generation. Despite architectural homogeneity, empirical evidence suggests DiT blocks contribute unequally to overall model performance—some block outputs, when ablated or reweighted, can even yield superior generation performance compared to the original intact model.

This motivates the central premise of the Calibri framework: post-hoc, parameter-efficient calibration of block contributions in a pretrained DiT can substantially increase generative quality. By learning only $\sim 10^2$ scalar calibration coefficients, the method eschews the computational burden of full-model fine-tuning, yet delivers measurable performance gains. The calibration is formulated as a black-box optimization—effectively solved via CMA-ES—targeting reward maximization as scored by an external model of preference or alignment.

Diffusion Transformer Architecture and Block Calibration

Diffusion Transformers (DiTs) comprise sequential DiT blocks, each with MHSA and feed-forward layers, often further extended in multimodal settings to MM-DiT where modality-specific tokens interact selectively.

Figure 1: DiT block scheme illustrating the architecture used for blockwise calibration.

Motivational ablation experiments reveal that: (1) disabling specific blocks can sometimes improve sample quality, and (2) introducing a learnable scalar per block further improves model outputs. Experiments indicate that standard DiT weighting schemes fail to optimally allocate representational capacity across layers.

Figure 2: DiT block ablation demonstrates uneven block importance, and blockwise scaling demonstrates empirically optimal coefficient values differing from pre-trained settings.

Calibri: Methodology and Optimization

Calibri introduces output scaling at three levels of granularity: block, layer, and gate (the latter specifically for MM-DiT architectures enabling multimodal interaction tuning). These coefficients are attached to pre-trained blocks, and only the calibration parameters are optimized.

The calibration task is cast as a black-box optimization, maximizing a reward function $R(c)$ induced by human-preference models (e.g., HPSv3, PickScore). To circumvent differentiability and non-smoothness in reward functions, Calibri employs the Covariance Matrix Adaptation Evolution Strategy (CMA-ES). Candidate parameter sets are sampled, evaluated via generated outputs, and iteratively refined.

Figure 3: Illustration of the calibration parameter search, highlighting the CMA-ES evolutionary optimization loop.

Comparative results demonstrate that finer scaling granularity (gate, layer vs. block) offers slight performance improvements but at the cost of increased computation due to larger search spaces. Calibri remains highly parameter-efficient across all settings.

Figure 4: Quantitative comparison of block, layer, and gate scaling on Flux, highlighting consistent improvements with minimal parameters.

Ensemble Calibration and Inference Acceleration

Calibri introduces an ensemble extension: separate calibrations are combined linearly at inference, with weights also subject to optimization. This enables leveraging diversity across differently calibrated instances, improving robustness and aggregate performance. Importantly, the ensemble calibration also generalizes Skip Layer Guidance and autoguidance strategies by training-time tuning rather than post-hoc heuristics.

Empirical findings show that Calibri ensemble not only outperforms single calibrated or original models on reward metrics, but also achieves equivalent or better performance with fewer inference steps, thus reducing computational load.

Figure 5: Calibri Ensemble outperforms the original model and achieves higher HPSv3 reward for fewer inference steps.

Empirical Evaluation

Calibri was evaluated on several leading open-source DiT-based T2I models (FLUX, Stable Diffusion 3.5M, Qwen-Image), consistently yielding improvements in HPSv3, ImageReward, and Q-Align metrics. Notably, calibrated models required substantially fewer inference steps without a loss in output quality—Flux, for example, achieves higher metrics with 15 steps versus 30 in the baseline.

Figure 6: Qualitative improvements of Calibri over baseline models for equivalent or lower inference cost.

User studies on Flux.1-dev and Qwen-Image corroborate quantitative findings, showing significant preference for Calibri samples in both overall quality and text alignment.

Integration with Alignment Methods and Reward Generalization

Calibri can complement or supplement other alignment techniques such as Flow-GRPO. Experiments show that applying Calibri to models already optimized for one reward further boosts target and auxiliary metrics, underscoring a degree of orthogonality with respect to reward-specific full-model fine-tuning. Impressively, Calibri attains competitive performance to full fine-tuning (e.g., Flow-GRPO) while modifying $10^5$ fewer parameters.

Figure 7: Qualitative comparison between Calibri and Flow-GRPO, illustrating Calibri's competitiveness for orders-of-magnitude fewer parameters.

Calibri demonstrates effective generalization across a variety of reward functions, not limited to the reward used for calibration.

Figure 8: Illustration of Calibri applied using different reward objective functions, showing improvement generalization.

Calibration Efficiency and Practicality

Calibri's optimization process is highly efficient: block-wise calibration of strong models such as Flux or Qwen-Image generally converges in under 1000 CMA-ES iterations and consumes only tens to hundreds of H100 GPU-hours as a fixed cost. Once calibrated, inference is permanently accelerated—up to 2x faster—for comparable or improved performance.

Analysis confirms stable and convergent optimization dynamics when using CMA-ES, with layer scaling parameters plateauing as expected and sigma decreasing accordingly.

Figure 9: CMA-ES provides much faster and more stable reward maximization compared to a gradient-based optimizer, Flow-GRPO.

Figure 10: Sigma decrease during CMA-ES optimization indicating reliable convergence behavior and training termination heuristics.

Limitations and Future Directions

A critical bottleneck remains the quality of external reward models used for calibration. Existing reward models exhibit insensitivity to subtle or rare visual errors (e.g., anatomical artifacts), potentially limiting the calibration effectiveness.

Figure 11: Reward model insensitivity to subtle visual artifacts can limit the optimality of learned calibration parameters.

Moreover, as calibration pursues the reward model objective, generative diversity can be affected; Calibri, however, tends to preserve diversity better than full fine-tuning approaches at comparable reward gains.

Future work may focus on: (i) advancing or ensembling reward models for artifact awareness, (ii) structured regularization of calibration coefficients for diversity stabilization, and (iii) application to non-text-to-image or multimodal diffusion settings.

Conclusion

Calibri represents a principled, parameter-efficient post-hoc calibration strategy for DiT-based generative models, achieving consistent empirical gains with only $\sim 10^2$ optimized parameters and a fraction of the cost of conventional fine-tuning. Its minimal computational and implementation burden, broad compatibility, and ability to accelerate inference without compromising quality position it as a pragmatic augmentation for state-of-the-art image generation. Its flexible ensemble framework, robust generalization across reward objectives, and complementarity with existing alignment approaches further expand its practical applicability.