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

arXiv: https://arxiv.org/abs/2606.06236 · HTML: https://mlsanity.com/paper/2606.06236/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Medical imaging and segmentation, one of 14 directions in Diffusion models.
- Builds on 35 of its 45 references that are in the map.
- First posted on arXiv 4 Jun 2026.

## Placed in

- Diffusion models / [Medical imaging and segmentation](https://mlsanity.com/diffusion/medical-imaging-and-segmentation/)

## Builds on (most cited first, 30 of 35)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [Image Super-Resolution Using Deep Convolutional Networks](https://mlsanity.com/paper/1501.00092/) (2014)
- [Multi-Stage Progressive Image Restoration](https://mlsanity.com/paper/2102.02808/) (2021)
- [Solving Inverse Problems in Medical Imaging with Score-Based Generative Models](https://mlsanity.com/paper/2111.08005/) (2021)
- [Improving Diffusion Models for Inverse Problems using Manifold Constraints](https://mlsanity.com/paper/2206.00941/) (2022)
- [UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models](https://mlsanity.com/paper/2302.04867/) (2023)
- [Denoising Diffusion Models for Plug-and-Play Image Restoration](https://mlsanity.com/paper/2305.08995/) (2023)
- [Compressed Sensing using Generative Models](https://mlsanity.com/paper/1703.03208/) (2017)
- [Inversion by Direct Iteration: An Alternative to Denoising Diffusion for Image Restoration](https://mlsanity.com/paper/2303.11435/) (2023)
- [Improving Diffusion Inverse Problem Solving with Decoupled Noise Annealing](https://mlsanity.com/paper/2407.01521/) (2024)
- [Solving 3D Inverse Problems using Pre-trained 2D Diffusion Models](https://mlsanity.com/paper/2211.10655/) (2022)
- [Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems](https://mlsanity.com/paper/2303.05754/) (2023)
- [Align Your Steps: Optimizing Sampling Schedules in Diffusion Models](https://mlsanity.com/paper/2404.14507/) (2024)
- [Fast ODE-based Sampling for Diffusion Models in Around 5 Steps](https://mlsanity.com/paper/2312.00094/) (2023)
- [Intermediate Layer Optimization for Inverse Problems using Deep Generative Models](https://mlsanity.com/paper/2102.07364/) (2021)
- [Accelerating Diffusion Sampling with Optimized Time Steps](https://mlsanity.com/paper/2402.17376/) (2024)
- [On the Trajectory Regularity of ODE-based Diffusion Sampling](https://mlsanity.com/paper/2405.11326/) (2024)
- [Improving Diffusion Models for Inverse Problems Using Optimal Posterior Covariance](https://mlsanity.com/paper/2402.02149/) (2024)
- [DiffusionBlend: Learning 3D Image Prior through Position-aware Diffusion Score Blending for 3D Computed Tomography Reconstruction](https://mlsanity.com/paper/2406.10211/) (2024)
- [S4S: Solving for a Diffusion Model Solver](https://mlsanity.com/paper/2502.17423/) (2025)
- [Provable Diffusion Posterior Sampling for Bayesian Inversion](https://mlsanity.com/paper/2512.08022/) (2025)
- [Learnable Sampler Distillation for Discrete Diffusion Models](https://mlsanity.com/paper/2509.19962/) (2025)
- [Misspecified Phase Retrieval with Generative Priors](https://mlsanity.com/paper/2210.05571/) (2022)
- [A Unified Framework for Uniform Signal Recovery in Nonlinear Generative Compressed Sensing](https://mlsanity.com/paper/2310.03758/) (2023)
- [Unleashing the Denoising Capability of Diffusion Prior for Solving Inverse Problems](https://mlsanity.com/paper/2406.06959/) (2024)

## Abstract (arXiv)

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.