Source-linked AI summary
Koopman Observers for Diffusion Acceleration: Correcting Feature Forecasts with Shallow Measurements
Hanru Bai, Yuanchao Xu, Fengyi Li
TL;DR
Feature forecasts based only on past activations cannot directly incorporate changes in the current denoising state. The paper introduces an observation-corrected Koopman framework that uses fresh shallow features to update deep-feature predictions while keeping the denoiser frozen. It reports lower Inception-feature error than channelwise affine prediction on both evaluated datasets under a matched schedule.
Problem
Past-activation forecasts cannot directly incorporate changes in the current denoising state, motivating a way to use fresh shallow features to correct deep-feature forecasts.
Source: §1 · p. 1Method
A time-dependent Koopman approximation predicts joint shallow–deep feature increments, and shallow-feature innovations correct deep-feature estimates without denoiser retraining.
Source: §1 · pp. 1–2Results
19.9% on CIFAR-10 and 11.9% on a ten-class ImageNet subset lower final-image Inception-feature mean squared error than channelwise affine prediction at the four-partial-step budget.
Source: §1 · p. 2Takeaways & Limitations
The experiments report that observation correction reduces error relative to the full sampler beyond coupled open-loop prediction, at about 1.5–1.6× speedups on two class-conditional U-Nets.
Source: §6 · p. 8Evaluation covers two low-resolution U-Net checkpoints; transfer to another setting requires checking prediction accuracy and may require recalibration.
Source: §6 · p. 8
Abstract
from arXiv · showhide
Feature caching accelerates diffusion sampling by replacing expensive network evaluations with predictions from previously computed activations. However, forecasts based only on past features cannot directly incorporate changes in the current denoising state. We investigate whether inexpensive, freshly computed features can serve as observations for correcting these predictions. We introduce an observation-corrected Koopman framework for accelerating frozen diffusion models. Using calibration trajectories, we identify finite-dimensional, time-dependent Koopman approximations that jointly describe the increments of shallow and deep network features. During accelerated sampling, these operators predict the evolution of expensive deep features, while innovations in the observed shallow features correct the predicted state. Periodic full evaluations refresh the observer, and all generative-model parameters remain unchanged. This formulation enables controlled comparisons of temporal prediction and observation correction. Across three 10,000-image runs per dataset, our method reduces paired Inception-feature MSE by $19.9\%$ on CIFAR-10 and $11.9\%$ on a ten-class ImageNet subset relative to channelwise affine prediction under the same four-partial-step schedule. Matched ablations attribute additional reductions of $4.54\%$ and $4.67\%$ to observation correction. The observer achieves $1.89\times$ and $1.85\times$ measured speedups over DDIM-50, supporting improved reference-sampler fidelity without retraining the denoiser.
1 INTRODUCTION
Past-feature forecasts cannot directly use changes in the current denoising state, while partial evaluations still provide fresh shallow features. The paper uses those observations to correct deep-feature forecasts and reports lower Inception-feature error than channelwise affine prediction under a matched schedule.
1 INTRODUCTION
- Figure 1 presents selected 64×64 ImageNet improvement cases, not average performance; its rows compare full-model, open-loop, and observation-corrected generations.Examples share initial noise and class label, and the accelerated variants share the full-evaluation schedule.Source: Figure 1 · p. 2
- A calibrated gain maps the discrepancy between predicted and observed shallow features into a correction of the estimated deep features, without accessing skipped deep features or retraining the denoiser.Periodic full evaluations refresh the observer.Source: §1 · p. 2
- 19.9% on CIFAR-10 and 11.9% on a ten-class ImageNet subset lower final-image Inception-feature mean squared error than the channelwise affine baseline at the four-partial-step budget.The comparison covers three sampling runs with 10,000 images per dataset per run.Source: §1 · p. 2
- Matched correction ablations show additional Inception-feature mean squared error reductions of 4.54% on CIFAR-10 and 4.67% on the ten-class ImageNet subset over disabled correction.The channelwise, open-loop, and observation-corrected predictors use the same full-evaluation schedule.Source: §1 · p. 2
- The evaluated observer implementations achieve measured speedups of 1.89× on CIFAR-10 and 1.85× on the ten-class ImageNet subset over full 50-step DDIM sampling.Source: §1 · p. 2
2 RELATED WORK
Prior work reduces diffusion sampling cost through fewer denoising evaluations, feature caching, and temporal feature forecasting. This paper distinguishes its approach by using current shallow-feature observations to update estimated deep features.
Diffusion acceleration and feature forecasting
- DDIM and DPM-Solver reduce sampling evaluations, while DeepCache reuses deep representations and TaylorSeer, Spectrum, L2P, and FastCache forecast feature evolution through temporal extrapolation, spectral expansions, learned historical coefficients, or linear approximations.The cited works are Ho et al. (2020), Song et al. (2021), Lu et al. (2022), Ma et al. (2024), Liu et al. (2025b), Han et al. (2026), Shen et al. (2026), and Liu et al. (2025a), respectively.Source: §2 · p. 2
Adaptive caching and feedback
- HarmoniCa addresses trajectory dependencies and training–inference mismatch, RACER reallocates full evaluations using forecast disagreement, and WorldDynCache combines a lifted latent surrogate with risk-controlled exact transitions.These works establish adaptive control and dynamics-inspired approximation as existing directions.Source: §2 · pp. 2–3
Koopman modeling and state estimation
- The paper identifies feature-transition matrices and corrects deep-state estimates with shallow-feature innovations, without assuming an exact invariant observable subspace or guaranteed observer convergence.Unlike WorldDynCache’s local-memory surrogate, its approach explicitly identifies feature-transition matrices; its practical value is assessed against same-information regression controls.Source: §2 · p. 3
3 PRELIMINARIES AND PROBLEM FORMULATION
The formulation keeps the sampler grid and denoiser parameters fixed while estimating deep features skipped during partial evaluations. It treats shallow features as measurements and uses data-calibrated, finite-dimensional dynamics to predict and correct feature evolution.
3.1 Diffusion Sampling
- For deterministic sampling such as DDIM with η = 0, the prescribed sampler update advances the sample along a decreasing diffusion-time grid; the objective keeps the grid and denoiser parameters fixed while reducing network computation.Source: §3.1 · p. 3
3.2 Partial Evaluation and Feature Estimation
- The denoiser is partitioned into shallow computation, expensive deep computation, and a final readout, with the original frozen network supplying each component.The shallow features provide retained skip connections, while the deep representation enters the final decoder.Source: §3.2 · p. 3
- Shallow features are computed at every step, but deep features are evaluated only at full-evaluation indices; otherwise, a causal estimator uses prior observations, cached features, and estimated states.At skipped evaluations the estimator cannot access the current deep feature.Source: §3.2 · p. 3
- The estimated deep feature is defined at the accelerated sampling state and can differ from the full-sampling reference because earlier approximations alter subsequent inputs.Source: §3.2 · p. 3
3.3 Finite-Dimensional Koopman Approximation
- A time-dependent Koopman operator acts linearly on observable functions even when the underlying dynamics are nonlinear; its state may include sample, conditioning, time, and history.Source: §3.3 · p. 3
- A finite observable dictionary yields a data-identified affine approximation whose residual captures finite-dictionary closure and identification errors, and is not assumed to vanish.Source: §3.3 · p. 3
- The method builds joint shallow–deep feature-increment observables, then calibrates bases, transition matrices, and correction gains offline without updating the denoiser parameters.Approximation fidelity and sampling cost are evaluated separately from distributional generation quality.Source: §3.3 · p. 4
4 METHOD
The method accelerates a frozen diffusion model by predicting skipped deep-feature increments and correcting them with current shallow measurements. Offline calibration identifies the operators and gains, while periodic full evaluations refresh the observer.
4.1 Joint Increment Observables
- Joint shallow–deep feature increments are represented in orthonormal channel bases and scaled by the root mean square of projected increments.Shallow increments remain observable at partial evaluations; measuring deep increments requires consecutive full evaluations.Source: §4.1 · p. 4
4.2 Offline Operator and Gain Identification
- Affine EDMD-style regression identifies time-dependent operators and correction gains from calibration data, with rank and regularization selected on validation trajectories.Complete trajectories are split into calibration, validation, and held-out diagnostic sets; all patches and conditioning branches from a trajectory remain together.Source: §4.2 · p. 4
4.3 Observation-Corrected Prediction
- At each partial evaluation, a forecast of the joint increment state is corrected using the current shallow observation, then the decoder receives the estimated deep feature and current shallow skips.No current deep feature is computed during this operation.Source: §4.3 · p. 4
4.4 Refreshes and Computational Cost
- Two consecutive full evaluations initialize the observer, and later full-evaluation pairs refresh its feature and increment estimates at the current accelerated states.A refresh does not reset the sample to the full-sampling reference trajectory; matching full-evaluation counts alone does not match runtime.Source: §4.4 · pp. 4–5, Figure 2 · p. 5
4.5 Error Propagation
- With the unregularized gain, correction minimizes expected deep-residual error; after reinitialization it does not increase the next-step mean squared deep error, with later steps conditioned on calibration-like error moments.The shallow residual predicts the fraction of deep residual measured by the stated observability ratio.Source: §4.5 · p. 5
- Episode error accumulates increment errors and drift outside the retained deep subspace, which vanishes when r = Cd; frequent reinitialization limits the horizon.The paper makes no claim of contraction, and for block-diagonal dynamics shallow correction acts on the residual alone.Source: §4.5 · pp. 5–6
5 EXPERIMENTS
The experiments compare feature-forecasting fidelity and distributional quality, isolate the effect of current shallow observations, and test how correction and temporal modeling vary across sampling stages and predictor configurations.
5.1 Experimental Setup
- The evaluation uses two frozen class-conditional U-Nets, a deterministic DDIM-50 reference, and three 10,000-image runs per dataset, with paired Inception-feature MSE measuring agreement with the reference.The observer uses offline calibration and does not update denoiser parameters; baseline configurations are selected by independently measured latency, with a 5% gap defining approximate matching.Source: §5.1 · p. 6
5.2 Efficiency, Fidelity, and Generation Quality
- At the middle budget, the observer reaches 1.73× and 1.70× acceleration and reduces Eϕ by 36.5% and 32.0% versus channelwise affine prediction on CIFAR-10 and ImageNet-10, respectively.Its FID is 13.515 ± 0.079 and 14.133 ± 0.097, compared with full DDIM-50's 13.393 ± 0.049 and 13.724 ± 0.083.Source: §5.2 · p. 7, Figure 3 · p. 7
5.3 Do Current Observations Improve Prediction?
- Current observations reduce Eϕ by 4.54% on CIFAR-10 and 4.67% on ImageNet-10 versus disabled correction, and by 12.78% and 17.44% versus delayed observations.The fixed-predictor ablation uses the same s = 4 refresh positions; measured latency overhead over disabled correction is 0.63% and 0.66%, respectively.Source: §5.3 · p. 7
5.4 Dependence on Sampling Stage and Prediction Horizon
- Correction reduces local deep-feature error at all 23 evaluated CIFAR-10 positions but only 7 of 23 ImageNet-10 positions; these diagnostics use reference trajectories and omit evolving-image feedback.At the earliest ImageNet-10 anchor, reductions range from 13.6% to 19.1%, while most middle and late positions show small error increases.Source: §5.4 · p. 7, Figure 4 · p. 8
5.5 The Role of Temporal Dynamics
- On the same s = 4 schedule, static regression has Eϕ of 0.012275 on CIFAR-10 and 0.034979 on ImageNet-10, versus 0.002713 and 0.017323 for the rank-256 observer.The full-trajectory diagnostics show the same aggregate ordering of linear extrapolation, fitted open-loop dynamics, and observation-corrected dynamics.Source: §5.5 · p. 8
- Increasing observer rank from 64 to 256 reduces paired feature error by 10.7% and 11.0%, while dense-versus-diagonal dynamics improve error by 12.6% on CIFAR-10 but only 0.5% on ImageNet-10 at rank 64.These comparisons support temporal state modeling in the evaluated predictor family, while showing dataset-dependent benefits from cross-channel coupling.Source: §5.5 · p. 8
6 CONCLUSION
The conclusion frames frozen-diffusion acceleration as state estimation, with shallow features measuring an unobserved deep-feature state. Experiments report lower full-sampler error than coupled open-loop prediction, while the authors limit the scope of the result.
6 CONCLUSION
- The formulation treats skipped deep features as an unobserved state, computed shallow features as measurements, and time-dependent affine Koopman approximations of increments as the dynamics.The observer corrects its prediction using shallow-measurement innovations, with identified coupling playing the role of an output matrix.Source: §6 · p. 8
- In experiments on two class-conditional U-Nets, correction reduces error relative to the full sampler beyond coupled open-loop prediction at speedups of about 1.5–1.6×.Source: §6 · p. 8
- The one-step guarantee applies to the calibration distribution; the paper establishes no closed-loop contraction, and episode error follows the stated error recursions.Source: §6 · p. 8
- The fitted predictor is tied to its checkpoint, sampling grid, and conditioning, while evaluation covers two low-resolution U-Net checkpoints rather than diffusion transformers, text-to-image, or video models.Transfer requires checking prediction accuracy and possibly recalibrating.Source: §6 · p. 8
A PROOF OF PROPOSITION 1
The appendix proof derives the correction-gain expression and error recursions, then describes the real-image references used for evaluation.
A PROOF OF PROPOSITION 1
- For γ = 0, the proof identifies the minimizing correction matrix by setting the convex quadratic loss gradient to zero, yielding L = ΣuyΣ−1 yy.The proof also establishes F(0) − F(L) ≥ 0 for the stated regularized expression.Source: A PROOF OF PROPOSITION 1 · p. 10
- The proof derives the observer error recursion from the prediction and correction relations, then obtains the cumulative error expression by summing per-step increments.Source: A PROOF OF PROPOSITION 1 · p. 10
B.1 Datasets and Reference Distributions
- The real-image references comprise 50,000 CIFAR-10 training images and 13,000 images from ten ImageNet checkpoint classes, not an ImageNet-1k evaluation.Images are RGB-converted, resized, and center-cropped to 64 × 64 for the ImageNet-10 reference.Source: §B.1 · p. 10
B.2 Offline Identification and Data Splits
Offline identification uses full-sampling trajectories from frozen checkpoints, with trajectory-level data splits separating calibration, validation, diagnostics, and formal evaluation. The fitted observer has fixed rank-256 shallow and deep components and does not claim exact finite-dimensional Koopman invariance.
B.2 Offline Identification and Data Splits
- Calibration, validation, and diagnostic splits keep complete initial-noise trajectories together, while new 10,000-image sampling runs are separate.The collector retains 16 fixed spatial locations per trajectory; Table 4 reports three formal sampling runs per dataset.Source: §B.2 · p. 10, Table 4 · p. 11
- The observer uses rank-256 shallow and deep increments, yielding a joint state dimension of 512, with identified parameters fixed during evaluation.The normalized-coordinate dynamics ridge parameter is 1.0 on both datasets; static rank-256 regression uses ridge 0.03.Source: §B.2 · p. 10
- The calibrated operator models network-feature increments, not the checkpoint’s native bottleneck operator, and requires no exact finite-dimensional Koopman-invariance claim.Offline predictor fitting is distinct from training the denoiser.Source: §B.2 · p. 11
B.3 Refresh Schedules and Baseline Implementations
The appendix specifies refresh schedules and baseline configurations, defines matched observation interventions, and reports diagnostic feature-prediction improvements on reference trajectories. Its controls distinguish current shallow measurements from delayed or disabled correction.
B.3 Refresh Schedules and Baseline Implementations
- The schedule uses 28/22, 23/27, or 20/30 full/partial steps for maximum partial-run lengths 2, 3, or 4, and baselines are selected by measured latency.At the middle budget, all seven baseline families are within 3.81% of observer latency; at s = 4, closest DeepCache and Spectrum candidates fall outside the 5% matching tolerance.Source: §B.3 · p. 11
B.4 Observation Interventions and Predictor Controls
- The current, delayed, and disabled interventions share a rank-256 base artifact, but delayed correction uses the preceding step’s shallow increment and disabling correction suppresses both state components.The open-loop control reproduces disabled-correction outputs, so it is not independent evidence; same-fit interventions are the direct correction test, not proof of superiority over every direct affine predictor.Source: §B.4 · p. 12
B.5 Feature Prediction on Reference Trajectories
- Observation correction reduces reference-trajectory aggregate diagnostic error from 0.3829 to 0.3664 on CIFAR-10 and from 0.3907 to 0.3763 on ImageNet-10.These diagnostics use shallow observations on full reference trajectories and aggregate horizons of one, two, and four steps; formal generation results assess accumulated effects separately.Source: §B.5 · p. 12, Table 6 · p. 13
B.6 Paired Errors and Uncertainty
The section defines paired-error and uncertainty procedures, reports classwise correction patterns, and gives batch-size timing results alongside the boundary that offline identification cost was not measured.
B.6 Paired Errors and Uncertainty
- Paired-bootstrap intervals resample paired image indices within each run, and describe image variation conditional on the checkpoint and fitted predictor, not pooled three-run uncertainty.Each comparison uses 1,000 resamples and reports the 2.5th and 97.5th percentiles separately for each run.Source: §B.6 · p. 13, Table 8 · p. 14
- Current observations improve on disabled correction in every class of both datasets, with relative gains of 0.37%–9.07% on CIFAR-10 and 0.77%–6.51% on ImageNet-10.Against delayed correction, current observations improve nine of ten ImageNet classes; airliner error increases 5.20%, and classwise averages do not guarantee gains for every image or run.Source: §B.6 · p. 13
B.7 Batch-Size Sensitivity and Offline Cost
- At s = 4, observer speedups range from 1.87× to 2.12× on CIFAR-10 and 1.80× to 1.92× on ImageNet-10 across batch sizes.Formal quality evaluation uses only batch sizes 50 and 20, respectively; the completed GPU suite reused calibrated predictors and did not measure end-to-end offline identification cost.Source: §B.7 · p. 14
- The serialized base observer is approximately 73.37 MiB per dataset, while the fused deployment artifact is approximately 36.05 MiB, excluding the denoiser.Serialized file size is not peak accelerator memory usage.Source: §B.7 · p. 14
B.8 Evaluation Records and Reproducibility
Evaluation records retain per-run and per-image outcomes, configuration details, and correctness checks, while qualitative previews use fixed indices and the introduction’s teaser examples have a narrower scope.
B.8 Evaluation Records and Reproducibility
- The evaluation retains all configurations and records per-run FID and KID, paired errors, classwise summaries, timing repetitions, hashes, and completion markers.Across two datasets, three runs, and 10,000 images per run, configurations share initial noise across methods, so generation counts are not independent initial conditions.Source: §B.8 · p. 14
- Correctness checks report zero maximum output differences for split full-forward paths versus native DDIM and DeepCache full-forward branches versus native execution.These checks validate the tested execution paths.Source: §B.8 · p. 14
B.9 Qualitative Evaluation and Scope
- Formal previews use the same first 40 indices per run across methods, whereas introduction teaser examples come from an earlier development cohort and are selected improvement cases.The teaser examples should not be treated as representative random samples.Source: §B.9 · p. 14