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

2026-06-04

AI Summary

This work addresses the high computational cost of diffusion model inference in 3D CT reconstruction and the suboptimal performance of existing uniform timesteps, which fail to align with the non-uniform dynamics of the reverse conditional diffusion process, leading to significant truncation errors. To overcome these limitations, the authors propose a plug-and-play timestep scheduling optimization framework that introduces, for the first time, an Oracle-guided global error minimization mechanism. By leveraging dynamic programming, the method extracts an optimal sparse schedule from a densely sampled trajectory—serving as a reference Oracle—and allocates limited timesteps precisely to the most error-sensitive stages. This yields a task-adaptive, non-uniform schedule that transcends conventional heuristic strategies. Evaluated on the AAPM dataset, the approach achieves substantially improved reconstruction fidelity and efficiency with only ≤10 steps, outperforming current scheduling schemes.

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.

Problem

• diffusion timestep scheduling
• 3D CT reconstruction
• truncation error
• computational overhead
• ill-posed inverse problems

Innovation

• diffusion timestep scheduling
• 3D CT reconstruction
• dynamic programming
• oracle-guided optimization
• computational efficiency