Technical Summary: Tracing the Oracle (TrO)

Problem Statement

Pretrained diffusion models offer significant potential for solving highly ill-posed 3D computed tomography (CT) inverse problems, such as sparse-view and limited-angle reconstruction. However, their practical application is hindered by two primary limitations:

1. Computational Overhead: The inference process typically requires thousands of sequential neural function evaluations (NFEs) to achieve high fidelity, imposing a substantial computational burden.
2. Inefficient Timestep Scheduling: Existing methods often rely on uniform or heuristic timestep schedules (e.g., Uniform-t, Cosine, EDM). These schedules fail to capture the non-uniform evolution of the reverse conditional diffusion stochastic differential equation (SDE). Consequently, they introduce substantial truncation errors, particularly when operating under strict computational budgets (e.g., fewer than 10 sampling steps). Furthermore, standard optimization strategies designed for unconditional generation (like GITS) are ineffective for conditional inverse problems because measurement constraints disrupt the unconditional generation manifold, and the stochastic nature of conditional solvers invalidates deterministic trajectory assumptions.

Methodology

The authors propose Tracing the Oracle (TrO), a plug-and-play framework designed to optimize timestep scheduling for 3D CT reconstruction. The core philosophy is to treat a densely sampled numerical integration trajectory as a high-fidelity "oracle" and extract a sparse, optimized subset of timesteps that minimizes the cumulative error relative to this oracle.

The methodology consists of three key components:

1. Oracle Trajectory Construction

A reference trajectory ($T_{oracle}$) is generated using the Decomposed Diffusion Sampling (DDS) framework with a large number of evaluations ($N=200$) on a subset of data. This dense path serves as the ground truth for the reverse diffusion process under specific measurement conditions.

2. Noise Reuse Mechanism

A critical challenge in quantifying the error between a sparse approximation and the oracle in a stochastic setting is the variance introduced by independent noise sampling. If the approximation and the oracle sample noise independently, the geometric distance between them fluctuates, obscuring the true deterministic truncation error.
To resolve this, TrO employs a noise reuse technique. When evaluating the transition cost from a source timestep $t_i$ to a target timestep $t_j$, the method reuses the exact stochastic noise vector ($\epsilon_{j-1}$) originally injected during the oracle's transition from $t_{j-1}$ to $t_j$. This isolates the deterministic truncation error from stochastic variance, allowing for precise error quantification.

3. Discrete Shortest-Path Optimization

The scheduling task is formulated as a discrete shortest-path problem. The goal is to select a subset of $L$ timesteps ($L \ll N$) that minimizes the cumulative transition cost.

• Cost Function: The cost $C_{i,j}$ between timesteps $t_i$ and $t_j$ is defined as the sum of the geometric distance (truncation error) and a regularization term:
  $$C_{i,j} = \|x_{t_j|t_i} - x_{t_j}\|_2 + \kappa \left( (t_i - t_j) - \frac{T}{L} \right)^2$$
  The regularization term penalizes large deviations from a uniform temporal stride to prevent the schedule from forcing the network to operate outside its trained conditional manifold.
• Solver: Dynamic Programming (DP) is employed to find the global optimum of this discrete path problem. Unlike greedy approaches, DP ensures the extracted schedule is the global minimum of the cumulative error.

Once computed offline on a calibration set, the optimized schedule is applied during the inference stage for 3D volume reconstruction without introducing additional computational overhead.

Key Contributions

1. Plug-and-Play Framework: The authors introduce TrO, which formulates timestep scheduling for measurement-conditioned diffusion as a discrete shortest-path problem to minimize cumulative trajectory errors in 3D inverse problems.
2. Noise Reuse Mechanism: A novel technique is introduced to isolate deterministic truncation errors from stochastic variance within SDE-based sampling trajectories, enabling accurate error quantification in conditional settings where standard deterministic metrics fail.
3. Performance Gains: Extensive evaluations demonstrate that TrO consistently improves reconstruction fidelity (PSNR and SSIM) over predefined heuristic schedules, particularly under strict computational budgets (e.g., 8–15 NFEs), without requiring additional inference time.

Experimental Results

The framework was evaluated on the AAPM dataset for Sparse-View CT (SV-CT) and Limited-Angle CT (LA-CT) tasks, integrated with the state-of-the-art DDS solver.

• Quantitative Performance: TrO outperformed various baselines (Uniform-t, Uniform-$\lambda$, Quadratic, EDM, Cosine) across different NFE budgets (8, 10, and 15 steps).
  • In the 8-view SV-CT task with 8 NFEs, TrO achieved a PSNR of 38.53 (Axial) and 39.78 (Coronal), surpassing the next best method (Quadratic) by a significant margin.
  • The performance gap between TrO and baselines widened as the NFE budget decreased, highlighting its efficiency in few-step regimes.
  • TrO also demonstrated superior performance compared to other diffusion-based solvers (DiffMBIR, Score-Med, MCG) that typically require thousands of NFEs to achieve competitive results.
• Qualitative Analysis: Visual comparisons showed enhanced artifact suppression and detail preservation, particularly in regions prone to reconstruction errors.
• Adaptive Behavior: Analysis of the optimized schedules revealed that TrO adapts to the specific measurement operator. For LA-CT, where low-frequency information is unreliable, TrO allocated more steps to the early stages of the reverse process to synthesize missing structures, whereas for SV-CT, it focused more on later stages for high-frequency refinement.

Significance and Claims

The paper claims that TrO addresses a critical bottleneck in applying diffusion models to 3D medical imaging: the trade-off between computational cost and reconstruction quality. By treating timestep scheduling as an optimization problem grounded in a high-fidelity oracle and resolving the stochastic error quantification issue via noise reuse, TrO enables:

• High-Fidelity Few-Step Reconstruction: Achieving reconstruction quality comparable to high-NFE solvers with a fraction of the computational cost (e.g., <15 steps).
• Task-Specific Adaptation: The ability to dynamically allocate computational resources to critical evolution stages based on the specific degradation characteristics of the inverse problem (e.g., sparse views vs. limited angles).
• Zero Overhead: The optimization is performed offline, meaning the inference stage incurs no additional computational burden beyond the standard sampling process.

The authors conclude that TrO facilitates the efficient and robust deployment of diffusion models in 3D medical imaging applications by narrowing the performance gap between few-step sampling and computationally expensive high-NFE solvers.