Tracing the Oracle: Improving Diffusion Timestep Scheduling for 3D CT Reconstruction

Abstract: Pretrained diffusion models demonstrate impressive potential in solving highly ill-posed 3D computed tomography (CT) inverse problems, while the inference process suffers from significant computational overhead. Furthermore, existing uniform timestep schedules fail to capture the non-uniform evolution of the reverse conditional diffusion stochastic differential equation, thereby introducing substantial truncation errors. To overcome this limitation, we propose Tracing the Oracle (TrO), a plug-and-play framework for improved timestep scheduling. Specifically, we treat densely sampled numerical integration trajectories on a few samples as the reference oracle. The optimized schedule is extracted by leveraging dynamic programming to globally minimize the cumulative error between the few-step approximation and the oracle. This mechanism precisely allocates the limited sampling steps to critical evolution stages that are highly susceptible to truncation errors. Our extensive experiments on the AAPM dataset across multiple 3D CT reconstruction tasks demonstrate that, when combined with the state-of-the-art 3D CT reconstruction method DDS, our optimized timesteps significantly improve reconstruction fidelity and computational efficiency compared to existing heuristic schedules, especially under a strict budget of no more than 10 sampling steps.

• The paper introduces TrO, which adaptively selects timesteps via discrete optimization to minimize error in 3D CT reconstruction.
• It leverages dynamic programming and noise reuse to allocate computational steps effectively under stringent evaluation budgets.
• Empirical results show TrO significantly improves PSNR and SSIM compared to uniform schedules, enabling high-quality reconstructions with fewer steps.

Tracing the Oracle: Adaptive Diffusion Timestep Scheduling for 3D CT Reconstruction

Overview

This work introduces Tracing the Oracle (TrO), a framework for non-uniform timestep scheduling in diffusion-based conditional inverse problems, with a primary focus on accelerating high-fidelity 3D computed tomography (CT) reconstruction. TrO formulates timestep selection as a discrete optimization by globally minimizing the approximation error between a sparse evaluation trajectory and a densely sampled high-fidelity "oracle" path. The methodology is motivated by the limitations of uniform and heuristic timestep schedules, particularly when computational constraints prevent the use of large neural function evaluation (NFE) budgets. TrO leverages dynamic programming, noise-reuse, and explicit discretization error modeling to allocate timesteps to stages of the reverse stochastic differential equation (SDE) process that are most susceptible to truncation error.

Methodology

The study starts from the premise that numerical inference in diffusion-based inverse problems is dominated by the discretization of the reverse SDE, and that uniform step allocation is suboptimal, especially for complex measurement-conditioned processes such as 3D CT. To address these, the authors introduce several key innovations:

• Oracle Trajectory Construction: An oracle trajectory is generated for a calibration sample using the Decomposed Diffusion Sampler (DDS) with a high NFE (typically $N=200$). The path samples fine-grained dense states of the learned denoising diffusion process, with each reverse step incorporating measurement data and gradient-based refinement using CG-accelerated ADMM.
• Plug-and-Play Discrete Optimization: Given an allowed budget $L \ll N$ of steps, the problem reduces to selecting a sparse subsequence of timesteps that minimizes the $\ell_2$ norm difference between the sparse and oracle trajectories.
• Noise Reuse Mechanism: To isolate discretization error and avoid stochastic variance, TrO aligns noise injections at each step by reusing the exact noise vectors from the oracle trajectory, enabling deterministic error estimation even in the presence of stochastic SDE sampling.
• Cost Matrix and Dynamic Programming Search: The total error for skipping intermediate steps is tabulated as a cost matrix, including a regularization penalty to prevent excessively non-uniform allocations that may misalign with the trained model's operational regime. The shortest-path solution over this matrix—solved via dynamic programming—produces a globally optimal (under the cost function) schedule.
• Generalization: Diffusion trajectories for samples from the same anatomical distribution display strong structural similarity in discretization error, allowing the optimized schedule to be computed offline on a calibration sample and reused for all subsequent inference without further overhead.

  Figure 1: Geometric interpretation of the Tracing the Oracle (TrO) framework, contrasting heuristic uniform scheduling with the proposed cumulative error-minimizing adaptive timestep allocation strategy.

Empirical Evaluation

Experimental Setup

The framework is evaluated on the AAPM dataset for two primary CT scenarios: sparse-view CT (SV-CT) and limited-angle CT (LA-CT). The DDS pipeline is used as the baseline diffusion inverse solver, and fixed (uniform, quadratic, logarithmic, cosine, EDM) and learned (TrO) schedules are compared under tight NFE budgets (8, 10, 15). State-of-the-art fast and baseline solvers (DiffMBIR, Score-Med, MCG) are included for reference, typically evaluated with four orders-of-magnitude larger sampling budgets.

Quantitative Results

TrO consistently outperforms all fixed schedules across SV-CT and LA-CT setups under severe computational constraints. The improvements in PSNR (e.g., more than +1 dB at NFE=8 vs. the best fixed schedule) and SSIM are substantial, especially for the lower NFE regime where truncation dominates over model capacity. The benefit is robust across all image planes (axial, coronal, sagittal). Compared to baseline methods, TrO-equipped DDS at 10-15 NFEs surpasses solvers utilizing 4000 NFEs.

Figure 2: Qualitative comparison of 3D CT reconstructions, illustrating enhanced high-frequency detail preservation and suppression of artifacts with TrO's adaptive scheduling (see highlighted regions).

Timestep Analysis

Schedule visualization reveals that optimal allocations discovered by TrO are highly problem-dependent: for SV-CT, more steps are allocated to later reverse time stages (refining high-frequency details), whereas for LA-CT, the schedule front-loads the early stages (synthesizing missing low-frequency structures in poorly measured regions).

Figure 3: Visualization of timestep allocations over discrete evaluation steps for various scheduling strategies, showing that TrO dynamically adjusts according to measurement operator and budget.

Ablation and Hyperparameter Study

An ablation over the schedule regularization parameter $\kappa$ shows optimal performance is achieved with moderate non-uniformity; extreme allocations can degrade fidelity, presumably by deviating from the network's training-time operational regime.

Theoretical and Practical Implications

Theoretical Insights: TrO establishes that, for measurement-conditioned SDE solvers, task- and operator-adaptive timestep allocation outperforms all uniform strategies. The global optimization and explicit separation of truncation error from stochastic variance enables practical schedule discovery in previously intractable conditional settings. The strong dependency of optimal schedules on inverse problem geometry underscores the need for domain-aware, data-driven adaptation in numerical generative inference.

Practical Impact: By closing the gap to oracle-level reconstructions at 10–15 steps, TrO makes high-quality diffusion-based reconstruction feasible under clinically relevant constraints. This enables broader deployment in resource-constrained and real-time medical imaging systems, improving accessibility without loss of fidelity. Since the schedule discovery is performed offline, the approach does not introduce any extra runtime or computational burden during inference.

Future Directions

Potential avenues for future research include:

• Extension to Other Modalities: Broadening TrO’s scheduler optimization to other inverse problems in medical and industrial imaging (e.g., MRI, PET, X-ray, seismic).
• Joint Model-Scheduler Training: Co-optimizing network parameters and schedules via meta-learning or bilevel optimization, potentially enabling even greater acceleration.
• Schedule Generalization and Transfer: Studying robustness when the measurement operator or anatomical distribution shifts; learning more data- and operator-agnostic schedule priors.
• Integration with Advanced Solvers: Combining TrO with fast ODE/SDE-based solvers or neural acceleration methods to approach real-time, high-dimensional generative inference.

Conclusion

Tracing the Oracle (TrO) provides a robust, theoretically grounded, and practically effective mechanism for optimizing timestep schedules in measurement-conditioned diffusion SDE solvers. By directly minimizing the discretization error on oracle reference trajectories and leveraging a noise-reuse mechanism, TrO allocates sparse computational budget where it is most impactful. The method produces non-uniform, task-specific schedules that significantly improve 3D CT reconstruction fidelity and efficiency under stringent computational budgets, enabling scalable deployment of diffusion-based generative models for complex inverse problems in medical imaging.