🤖 machine learning

TMPDiff: Temporal Mixed-Precision for Diffusion Models

The paper introduces TMPDiff, a temporal mixed-precision framework that adaptively assigns varying numeric precisions to different denoising timesteps in diffusion models to significantly reduce inference latency while maintaining or improving perceptual quality compared to uniform-precision baselines.

Original authors: Basile Lewandowski, Simon Kurz, Aditya Shankar, Robert Birke, Jian-Jia Chen, Lydia Y. Chen

Published 2026-03-17

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Basile Lewandowski, Simon Kurz, Aditya Shankar, Robert Birke, Jian-Jia Chen, Lydia Y. Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a master chef trying to bake a perfect, complex cake (generating an image) using a very expensive, slow oven (a massive AI model).

The problem is that this oven takes a long time to bake the cake because it has to go through 20 distinct stages of baking, checking the batter, and adjusting the temperature at every single step. If you try to speed things up by turning down the heat (using lower precision math), the cake might come out burnt or lumpy. If you keep the heat high (full precision) the whole time, it takes forever.

TMPDiff is a new "smart oven timer" that solves this by realizing: You don't need the high-heat setting for every single minute of baking.

Here is the breakdown of how it works, using simple analogies:

1. The Problem: The "All-or-Nothing" Approach

Currently, most AI image generators are like chefs who decide: "I will either use the super-precise, slow oven for the entire 20 minutes, OR I will use the fast, rough oven for the entire 20 minutes."

The Slow Oven (Full Precision): Makes a beautiful cake, but takes hours.
The Fast Oven (Quantization): Makes the cake quickly, but it often looks a bit "off" or blurry because the rough oven made mistakes early on that ruined the final result.

2. The Insight: Not All Steps Are Created Equal

The researchers discovered something fascinating about how these "baking" steps work:

The Early Steps (The Foundation): The first few minutes of baking are critical. This is when the cake's shape and structure are formed. If you mess up here, the whole cake is ruined. These steps need the High-Precision Oven.
The Late Steps (The Details): The last few minutes are just about smoothing out the frosting or adding a little color. If you use the Fast Oven here, the cake still looks great because the hard work was already done.

The old way ignored this. It treated every minute of baking as equally important. TMPDiff says, "Let's use the fancy oven for the critical first steps, and the fast oven for the easy finishing steps."

3. The Magic Trick: The "Additive Error" Theory

How does the AI know which steps are critical?

The researchers came up with a theory called the Additive Error Model. Imagine that every time you use the "Fast Oven," you drop a tiny crumb of dirt into the batter.

If you drop a crumb in the beginning, it gets mixed into the whole cake and ruins the texture.
If you drop a crumb at the very end, it just sits on top and is easy to wipe off.

They proved that these "crumbs" (errors) add up in a predictable way. By measuring how much damage a "crumb" causes at each specific minute, they can create a Priority List. They found that the errors from the early steps are huge, while errors from the late steps are tiny.

4. The Solution: TMPDiff (The Smart Scheduler)

Instead of trying every possible combination of "Fast vs. Slow" ovens (which would take a million years to calculate), TMPDiff uses a clever Adaptive Bisection strategy.

Think of it like finding a needle in a haystack, but you know the needle is heavy.

You don't check every single piece of hay.
You check the middle, the start, and the end.
You see where the "weight" (the potential for error) is highest.
You zoom in on those heavy areas to find the exact spots that need the "Slow Oven."
You leave the light areas for the "Fast Oven."

This allows the system to figure out the perfect schedule in seconds, rather than days.

5. The Result: The Best of Both Worlds

When they tested this on famous AI models (like FLUX and Stable Diffusion):

Speed: They got images 2.5 times faster than the high-quality version.
Quality: The images were 10-20% better than the standard "fast" version.
The Sweet Spot: On the FLUX model, they got an image that was 90% as perfect as the slow version, but in less than half the time.

Summary

TMPDiff is like a smart traffic controller for AI image generation. Instead of treating every second of the process the same, it knows exactly when to slow down and pay attention (High Precision) and when to speed up and coast (Low Precision). This lets us get high-quality images much faster without the computer getting confused or the image looking bad.

1. Problem Definition

Diffusion models, particularly for Text-to-Image (T2I) generation, suffer from high inference latency due to their iterative denoising processes involving billions of parameters. While quantization (reducing bitwidth from FP16/BF16 to INT4/INT8) is a standard method to accelerate inference, existing approaches apply a fixed precision across all denoising timesteps.

The authors identify a critical gap:

Spatial Mixed-Precision (varying precision across model layers) is well-explored but static at runtime.
Temporal Sensitivity is ignored. The paper hypothesizes that different denoising timesteps have varying sensitivities to quantization errors. Early steps (establishing global layout) may be more sensitive than later steps (refining details), or vice versa.
The Challenge: Finding the optimal assignment of precision (Full-Precision vs. Quantized) for each of the $T$ timesteps is a combinatorial search problem ( $2^T$ ), which is computationally intractable for large $T$ .

2. Methodology: TMPDiff Framework

The authors propose TMPDiff, a framework that assigns different numeric precisions to different denoising timesteps based on an Additive Error Model.

A. Theoretical Foundation: Additive Error Model

The core theoretical contribution is the derivation that quantization errors accumulate approximately additively across timesteps.

Error Propagation: Using a first-order Taylor expansion of the DDIM update rule, the authors derive a closed-form expression for the final accumulated deviation ( $\delta_0$ ) in the latent space.
The Ansatz: They propose that for a mixed-precision schedule $Z$ (where $z_t=1$ is full-precision and $z_t=0$ is quantized), the final error is a weighted sum of per-timestep errors:
$\boldsymbol{\delta}_0(Z) \approx \sum_{t=1}^{T} \left( \prod_{j=1}^{t-1} \mathbf{A}_j \right) B_t \boldsymbol{\epsilon}_t (1 - z_t)$
This implies that the impact of quantizing a specific timestep is independent of the precision choices of other timesteps, allowing the problem to be decomposed.

B. Empirical Validation

The authors validated this additive assumption experimentally:

Single-Timestep Analysis: They measured the "upcasting gain" (reducing error by switching one step to FP) and "downcasting loss" (increasing error by switching one step to INT). They found a strong linear correlation ( $R^2 > 0.88$ , Pearson $r > 0.94$ ) between these metrics, confirming that errors accumulate additively.
Multi-Timestep Prediction: They demonstrated that the total error of a complex schedule can be predicted by summing the individual gains/losses of the selected timesteps, with high rank-order agreement (Spearman $\rho > 0.99$ ).

C. The TMPDiff Algorithm (Adaptive Bisection)

To solve the optimization problem (minimizing error for a fixed latency budget $K$ full-precision steps) without exhaustive search:

Latency Budget: Determine the number of full-precision steps ( $K$ ) allowed based on the target speedup using Amdahl's Law.
Sparse Sampling: Instead of evaluating all $T$ $T$ timesteps, the algorithm uses an adaptive bisection strategy:
- Start with anchor points (start, middle, end).
- Iteratively identify the segment with the highest average potential gain.
- Sample the midpoint of that segment.
- Repeat until a sampling budget is exhausted.
Greedy Selection: Interpolate the gains for unsampled timesteps and select the top $K$ timesteps with the highest error-reduction potential to run in full precision.
Complexity: This reduces the search from exponential ( $2^T$ ) to linear time ( $O(T)$ or $O(B)$ where $B$ is the sampling budget).

3. Key Contributions

Additive Error Model: Theoretical derivation and empirical validation that quantization errors in diffusion models accumulate additively, enabling the decoupling of timestep sensitivity.
TMPDiff Framework: A novel temporal mixed-precision scheduling framework that dynamically switches between full-precision and quantized models during inference.
Efficient Optimization: An adaptive bisection algorithm that finds the optimal precision schedule with linear evaluation complexity, avoiding brute-force search.
Generality: Demonstrated that the temporal sensitivity scoring applies not just to quantization, but also to mixing models of different sizes (e.g., using a small model for insensitive steps and a large model for sensitive steps).

4. Experimental Results

The framework was evaluated on four state-of-the-art models (FLUX.1-dev, PixArt- $\Sigma$ , Sana, StableDiffusion XL) and three datasets (COCO, DCI, MJHQ).

Performance vs. Uniform Baselines:
- At matched speedups, TMPDiff consistently outperformed uniform-precision baselines (both full-precision and uniformly quantized).
- Quality Improvement: Achieved 10% to 20% improvement in perceptual quality metrics (SSIM, PSNR, LPIPS).
- FLUX.1-dev Results: Achieved 90% SSIM relative to the full-precision model while delivering a 2.5 $\times$ speedup over 16-bit inference. This represents a 60% reduction in latency with minimal quality degradation.
- SDXL Results: Outperformed the TensorRT 8-bit baseline significantly, achieving 0.85 SSIM vs. 0.42 SSIM at 1.33 $\times$ speedup.
Pareto Frontier: TMPDiff successfully fills the gap between the high-quality/slow full-precision pipeline and the low-quality/fast quantized pipeline, offering a tunable trade-off.
Model Mixing: When applied to mixing Sana (1.6B) and Sana (600M) models, TMPDiff outperformed heuristic and evolutionary search methods (DDSM) with lower computational cost.

5. Significance and Impact

New Optimization Axis: TMPDiff introduces "temporal mixed-precision" as a complementary optimization axis to existing "spatial mixed-precision" (layer-wise) techniques.
Practical Deployment: By reducing the search complexity from exponential to linear, TMPDiff makes dynamic precision scheduling feasible for large-scale diffusion models in production.
Efficiency: It allows users to achieve near-full-precision quality with significantly reduced inference latency, addressing the primary bottleneck of diffusion model adoption in real-time applications.
Future Work: The authors note that efficient kernel-level integration for switching precisions and weight offloading strategies are necessary for further deployment optimization.

In summary, TMPDiff provides a principled, mathematically grounded, and empirically validated method to accelerate diffusion models by intelligently allocating computational resources (precision) across the temporal dimension of the generation process.