Diffusion Active Learning: Towards Data-Driven Experimental Design in Computed Tomography

Authors: Luis Barba, Johannes Kirschner, Tomas Aidukas, Manuel Guizar-Sicairos, Benjamín Béjar
Primary area: applications to computer vision, audio, language, and other modalities
Keywords: Active Learning, Diffusion, Tomography, Computer Vision, Experimental Design

Abstract:
We introduce _Diffusion Active Learning_, a novel approach that integrates a generative diffusion model with sequential experimental design to adaptively acquire data for solving inverse problems in imaging. We first pre-train an unconditional diffusion model on domain-specific data. The diffusion model is aimed to capture the structure of the underlying data distribution, which is then leveraged in the active learning process. During the active learning loop, we use the forward model of the inverse problem together with the diffusion model to generate conditional data samples from the posterior distribution, all consistent with the current measurements. Based on the generated samples we quantify the uncertainty in the current estimate in order to select the most informative next measurement. We showcase the proposed approach for its application in X-ray computed tomography imaging. Our results demonstrate significant reductions in data acquisition requirements (_i.e._, lower X-ray dose) and improved image reconstruction quality across several real-world tomography datasets.

Decision: reject
Reviewer scores: 6, 5, 8, 5

Reviews and discussion:
<thread>
creation date:
1729782262466

thread_type:
official_review

<thread_object>
creator: reviewer_coa4

<content>
summary:
The paper proposes a framework for adaptive-sampling in the context of limited-angle X-ray/computed tomography (CT). Using measurements collected from a subset of angles, a diffusion model is used to generate approximate posterior samples. Then, the forward model is applied to each posterior sample to obtain the corresponding measurement at different angles. The uncertainty is represented as the variation in the measurements at each angle. Finally, the angle with the largest variation is selected as the next angle to collect measurements for. The paper shows that the active learning approach provides higher PSNR with fewer measurement steps compared to uniform sampling.

strengths:
- The paper tackles an important problem in the field of limited-angle CT. Long-scan times and high radiation doses clearly pose hurdles in all applications from medical tests to chip analysis.
- The solution is well-motivated. By identifying the angles with the most uncertainty, the proposed method promises to select the next angle with the most information. 
- The experiments demonstrate notable gains in PSNR using the active learning approach versus the uniform sampling.

weaknesses:
- The contributions of the paper were not explicitly clear to me. From the experiments, there were two independent variables that were changed: 1) the method used to generate samples and 2) the use of the active learning procedure. Is the main contribution the use of a diffusion model for the sampling procedure? Or is the main contribution the active learning procedure? Or is the combination of the two the main contribution? The diffusion sampling is based on an existing approach (Song et al. 2023), and it seems like the active sampling approach is based on existing uncertainty sampling. Thus, it is difficult to see where the novelty/contribution of the paper lies. It would be helpful if you could explicitly stated the contributions in a set of bullet points in the introduction. 
- The structure of the experimental section is confusing, particularly section 4.2. There is not any context as to what the methods (SWAG, Bootstrap, etc). are used for. Before introducing them, it would be helpful to identify where they are utilized in the framework. It was not clear to me until the results section that they would be substituted in for the diffusion sampling. Also, "Comparison Methods" would be a better suited title for the subsection.
- In a similarly light, the paragraph from lines 468-473 lacks context. It is unclear which Table/Figure the analysis is discussing.
- I'm not fully convinced about the practical advantage of the active sampling with the diffusion approach. As stated in the conclusion, diffusion models are inherently computationally heavy and slow. Thus, while you may need fewer measurements overall, the collection of each measurement would take much longer. For example, if it takes x times as long to choose the next angle than it does to just sample the next uniform angle, then you would want to show that your method allows you to collect at least x times fewer samples.
- In Line 300-301, it would be useful to use a different variable rather than t in order to avoid confusion with the diffusion time steps.

questions:
- In Eq 3, why do you take the mean of the posterior samples first and then apply the forward operator? Would it make more sense to take the mean of the measurements (i.e. apply the forward operator first and then take the mean)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730568732845

thread_type:
official_review

<thread_object>
creator: reviewer_fmc9

<content>
summary:
This paper utilizes a generative model which is then used in the active learning process to choose the next, most informative measurement. First, the authors train an unconditional diffusion model on a specific-dataset. In the second step, samples are generated whereby the diffusion model is conditioned on the measurements. With these samples the next measurement angle is chosen which has the highest posterior variance. According to this the total dose and acquisition time can be reduced.

strengths:
* The paper contains a really good explanation of the novel approach. 
* The approaches, results and limitations of already existing work is well discussed. 
* Reducing the dose or measurement time during the CT measurement is an essential problem.
* Novel combination of diffusion models with active learning.
* The results are well discussed and compared to different baselines.

weaknesses:
* Pre-training of the diffusion model is necessary. Further steps depend on this.
* The diffusion model is highly dependent on the trained data.
* The diffusion model could introduce undesirable biases. 
* How to get the posterior distribution could be discussed in more detail.

questions:
* Why are medical images not suited for this approach? In the paper it is stated because they are acquired very fast and therefore sparse but the goal is to have fewer measurements while keeping the resolution high?
* How does this model perform with samples that are slightly out of the distribution the diffusion model was trained on?
* Samples can be destroyed when a high dose or a long-time measurement is taken. How would this approach reconstruct the image? Would it automatically reduce the distortions which could be undesirable?
* Why are the smaller images first cropped and then rescaled? The distribution changes when rescaling images.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730639262414

thread_type:
official_review

<thread_object>
creator: reviewer_8grq

<content>
summary:
The paper proposed Diffusion Active Learning that integrates a generative diffusion model with active learning to select projection angles. Based on the pretrained  unconditional diffusion model,  the proposed model using the sampled images to select the most informative next measurement.

strengths:
The proposed method use half or less measurement to achieve the same performance with the compared methods.

weaknesses:
(1)	Comparison with activate learning method is preferred, such as the method proposed in [1]. Please compare the reconstruction result and inference time  with use the same number of projection angle.
(2)	The experiment is performed with parallel radon transform. More complex setting, such as fan-beam or 3D Cone beam, can verify the effectiveness of the Proposed method.  
(3)	The inference time is a huge disadvantage for you need n round k times sampling and n times full view projection. During, the n\times k sampling, the inversion problem cannot be avoided.  Please discuss potential ways to mitigate the computational cost. Please give  a more detailed analysis of the trade-off between computational cost and reconstruction quality.
(4)	 The improvement of the result may come from the diffusion model.  Comparison with DPS or proposed method without active loop using sparse view projection data, i.e. uniform projection angles (27,15,18 angles), is necessary.  This can  help give the explanation of the benefits of  proposed method  from the active learning component or the diffusion model

[1]Ce Wang, Kun Shang, Haimiao Zhang, Shang Zhao, Dong Liang, S. Kevin Zhou. Active CT Reconstruction with a Learned Sampling Policy, Proceedings of the 31st ACM International Conference on Multimedia, October, Pages 7226–7235, 2023

questions:
(1)	Plot the distribution of selected angle of different datasets.
(2)	The shape of the objection has influence of the selected angel,  
(3)	The setting of projection geometry must be given.
(4)	The definition of the notation must be given such as x^* in algorithm 1.
(5)	Testing on real projection data can verify the value of the proposed method.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730648460118

thread_type:
official_review

<thread_object>
creator: reviewer_hvxt

<content>
summary:
Context for those unfamiliar: Computed tomography (CT) acquires multiple X-ray _projection_ images of an object to reconstruct the 3D object. Due to ionizing radiation, there are significant risks associated with acquiring multiple X-ray viewing angles, leading to an undersampled ill-posed inverse problem. Many lines of work aim to reconstruct 3D CT using as few X-ray projections as possible.

Submission 10594 presents an active learning strategy to adaptively sample viewing angles most informative to the reconstruction, to reduce overall X-ray dosage. It first pretrains a diffusion model on fully sampled CTs from the same domain. Then, during inference, it uses the uncertainty of the posterior samples of the diffusion model to adaptively sample new angles.

Experiments are presented on three simulated datasets, where the proposed diffusion-based method compares favorably to other generative models.

strengths:
- The submission tackles an important yet rarely-trodden inverse imaging problem. 
- The submission is very open with its limitations which is an absolute breath of fresh air in modern papers. For example, L078 gives a much needed disclaimer about the risk of hallucinations from generative models in ill-posed medical image reconstruction problems. The submission’s discussion does a great job of listing limitations as well.
- Overall, the submission is very clearly and straightforwardly presented and was a very easy read.

weaknesses:
I am open to changing my score and look forward to the rebuttal. As of now I see the following areas that should be addressed,

## 1. The same method was presented in Elata, et al ECCV 2024

The submission has the same idea, methods, and subject matter as [Elata, et al ECCV 2024](https://arxiv.org/abs/2407.08256). **This overlap does not affect my rating** as ICLR’s reviewer guide states that papers that came online after Jul 1 count as contemporaneous and Elata et al first appeared on Jul 11. 

However, could the authors please enumerate the technical differences between the works such that readers can have clear takeaways from this paper? 

For example, the acquisition function is different between the two papers, but their covariance-based acquisition function does seem to be inadvertently benchmarked in the Appendix of this submission as well and they perform identically.

## 2. Limited experiments

My biggest reservation is w.r.t. the submission’s limited experimental depth from the following aspects.

### 2.1. Missing Active CT baselines

While somewhat niche, active learning for CT reconstruction has been studied by previous works as well. For example,
- https://arxiv.org/abs/2006.02420
- https://arxiv.org/abs/2211.01670
- https://dl.acm.org/doi/10.1145/3503161.3548204

Could the authors please describe why these works were not discussed and/or benchmarked against in this submission? If it is feasible, it would be good to see experiments comparing the submission against them. Of course, it is understandable if this is not feasible given the limited discussion period.

### 2.2. Only CT experiments

As the submission itself states, nothing in the submission is particularly specific to CT and it could just as well be used for other sensor-domain reconstruction problems such as MRI. As MRI is widely used, has a clear case for acceleration (patient comfort, time costs, etc.), and MRI active learning is more widely studied than CT active learning, is there a specific reason why it is not studied in this submission?

Further, there are several reinforcement learning methods cited in the paper for MRI active learning. Could any of them be also adapted for CT active learning to form benchmarks for this submission?

### 2.3. No low-dose / sparse-view baseline(s)

The submission motivates itself by potentially reducing CT dosage. Low-dose and/or sparse-view CT reconstruction are immensely popular topics with both learned and hand-crafted priors used. However, the paper does not benchmark against any of the work within this field and instead only benchmarks against other sampling-based methods specifically constructed for this submission. 

While I understand that sampling view prediction and low-dose reconstruction are somewhat orthogonal and can be combined, the method in this paper _requires_ the use of a diffusion model. This then precludes the use of useful low-dose reconstruction methods based on priors such as total variation. 

Could the authors please discuss the differences between the proposed method and existing methods for low-dose reconstruction and whether regularizers such as TV can also be used in the proposed setup?

### 2.4. Only simulated data

While this is endemic across the field, the submission uses _only_ simulated synthetic X-ray projection data in its experiments, simulating it using the same exact forward model as it does in its model. As per the “inverse crime” phenomenon, this can create highly optimistic results and exaggerate differences between methods.

Within CT, there is a small set of datasets that provide both CT and raw _measured_ projection data. For example, please see:
- https://www.cancerimagingarchive.net/collection/ldct-and-projection-data/ (they provide scripts to rebin to fanbeam if necessary)
- https://www.nature.com/articles/s41597-019-0235-y
- https://www.nature.com/articles/s41597-023-02484-6

As detailed above, the paper could have also used active learning baselines for MRI and there are large datasets of real k-space measurements for MRI.

Could the authors please detail why the experiments only use simulated projections?

## 3. Technical contribution

Reductively speaking, the paper can be viewed as a combination of Hard Data Consistency (Song et al 2023) and uncertainty sampling. The submission instead proposes to use “soft” data consistency which is hard DC + early stopping, but it does not perform an ablation of this choice (please correct me if I missed it). As this is the primary technical delta, please perform an ablation if possible.

## 4. Minor
- Runtime requirements are not reported at all. As the paper is motivated by accelerating scans, it should quantify what the additional computational overhead boils down to.
- L462: “on pair” → “on par”

questions:
- Could the authors please enumerate the technical differences between the submission and Elata24 such that readers can have clear takeaways from this paper? 
- Could the authors please describe why active CT baselines were not discussed and/or benchmarked against in this submission? 
- Why are the experiments limited to just CT if the method is generically applicable?
- Could the authors please discuss the differences between the proposed method and existing methods for low-dose reconstruction and whether regularizers such as TV can also be used in the proposed setup?
- Could the authors please detail why the experiments only use simulated projections?
- An ablation from hard to soft data consistency would be nice.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732742027554

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers,

We would like to thank you for your time and the valuable feedback. Several reviewers had raised questions about the contributions and benchmarks, which we answer below. Please refer also to the updated pdf for additional analysis of our experimental evaluation, including a RL baseline and evaluation on the fastMRI dataset. We also comment on the concurrent work by Elata et al (2024).

## Summary of Contributions

* We introduce diffusion active learning, a novel approach that **combines** diffusion posterior sampling and active learning for angle selection in CT (and row selection in MRI, see below)
* To the best of our knowledge, this is **one of the first works to demonstrate the benefit of using diffusion models posterior sampling in combination with active learning**. In particular, we demonstrate that the advantage is **not** purely additive (i.e. cannot be attributed to better angle selection or better reconstruction alone), but the best result is obtained precisely by the combination of both (see Figure 11 and Section C.3 in the updated PDF).
* Moreover, our work highlights how “hallucinations” of the **diffusion model captures the variance in the estimation**, which is essential for the active learning process in the early stages; at the same time,  **diffusion posterior sampling ensures data-consistency** so that with enough measurements, the final reconstruction contains no unwanted hallucinations. 
* We demonstrate that the **efficacy of active learning is dataset dependent**: Our evaluation shows clear gains on the composite and chip data, and no gains on the lung data set with the CT model. This on its own is an important observation and has not been emphasized in earlier works. It should not be overlooked as the potential gains for any active learning method are tied to the particular data distribution, and the forward process governing the observations. In some scenarios, such acquisition (e.g., CT on lung data) does not allow for efficient adaptive sampling, and one should not expect to have a method that universally works on general data distributions under constrained acquisition setups.
* We demonstrate a **significant computational advantage of using soft-data consistency** instead of hard data consistency as a diffusion posterior sampling approach (see additional experimental evaluation below, and Figure 6 in the updated PDF).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732742118430

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
## Additional Experimental Evaluation

### Better hyperparameters for diffusion posterior sampling with soft-data consistency

While re-evaluating our experimental results and performing additional ablations, we found that our implementation of soft-data consistency was using suboptimal hyperparameters which caused the PSNR to plateau after around 30-50 steps at a PSNR of around 32-35; the pdf reports now the **updated results showing larger gains for the diffusion active learning**, in particular for a large number of measurement angles. The conclusions remain the same, and the performance for very sparse data as well. The technical difference is to ensure that the data consistency is performed as a last step in the denoising update (previously, the last step in the denoising pipeline was a diffusion denoising step, adding unnecessary noise to the final diffusion samples). 


### Additional baselines on reported experiments (e.g., RL)

Several reviewers where asking about RL based baselines; we initially did not consider RL methods as baselines for several reasons: First, training RL algorithms requires a full forward simulation of the data pipeline and is specific to the reconstruction method used; in addition, training an RL policy in combination with the diffusion model leads to a computational overhead, rendering the approach effectively unfeasible. In comparison, diffusion models can be trained “offline” on existing reconstructions, and do not require access to the forward model. Second, training RL algorithms is generally considered to be very sensitive to hyper-parameter choices, while diffusion models for CT reconstructions are much better understood. Third, RL methods are “black box” as they just output an angle, whereas posterior samples from the diffusion model can be more easily interpreted.

Most importantly, however, **the goal of our work is to demonstrate that diffusion models can provide a sufficiently structured prior that can be exploited for sequential data acquisition** (the exact acquisition function, e.g. variance, entropy of committee-based, is secondary, see our comparison of different acquisition functions). RL instead, arguably, addresses the orthogonal problem of learning an acquisition function that is effective for a specific reconstruction method. One could therefore ask if using RL to select angles for diffusion posterior reconstruction leads to better reconstructions than using uncertainty sampling; in our opinion this would likely lead to marginal gains at most; with a significant computational overhead and increased complexity of the overall approach. Diffusion-posterior sampling strategies are still orders of magnitude slower than FPB and even SART, which would blow up the training time of a RL policy trained that uses a diffusion model in its reconstruction step. 

That said, we use the author implementation of Shen et al (2020, https://arxiv.org/abs/2006.02420) to add a **RL-based baseline** to our evaluation (using SART reconstruction). **Preliminary results are shown in Figure 5 for the first 50 steps**; we will provide a complete evaluation for all datasets for the final version of our paper.
Note that when training RL the chip dataset with a **fixed rotation**, the reconstructions initially show a small advantage w.r.t. diffusion active learning, but saturate quickly below the diffusion approach. When training the RL approach on arbitrarily rotated images (as we did for the diffusion approach), the RL approach is significantly worse than diffusion active learning.
This shows that this RL strategy is better at learning a global sampling strategy that works for all samples, instead of a sample-dependent strategy as done with DAL.


We are aware of further baselines that are reported for MRI reconstructions, however due to the limited time we could not include them for the rebuttal (e.g. https://arxiv.org/abs/2211.01670 does not provide code). Other baselines such as https://dl.acm.org/doi/10.1145/3503161.3548204 learn a global sampling distribution jointly for all the samples, i.e., the sampling strategy is not sample-dependent. This is good for distributions with lower variability and that have always the same orientation. In many tomographic settings however (like synchrotron nano-tomography), the orientation is completely random, so learning a fixed distribution that is not sample-dependent will suffer even when presented with the same sample with an arbitrary rotation.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733312902510

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers,

We'd like to briefly summarize the additional evaluation that we are providing based on the feedback of the reviews. We will add these to the final version of our paper.

Additional baselines and ablation:
* Active CT [Wang et al; 23] and RL baseline [Shen et al (2020]: Preliminary results are here: https://drive.proton.me/urls/Y4H5V360QR#o6THDhGTNA1j  The final version of our paper will contain evaluation on all datasets.
* Evaluation using fan beam geometry: https://drive.proton.me/urls/DV7FBVGQ3G#O68PZXnzeajY
* Evaluation fastMRI single-coil knee data: Appendix B
* Ablation study for soft data consistency:  Figure 6 and subsection "Soft Data Consistency and Early Stopping" in Appendix A
* Ablation angle selection vs reconstruction: Appendix C.3
* Visualization of selected angles: Figures 13-15 in the Appendix

We will also revise the text to better highlight our contributions, and discuss the concurrent work of Elata et al (for a detailed discussion of differences, see our earlier response below).
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper presents a new framework for CT angle selection. It is the first to combine diffusion models with active learning for this task, which has a certain degree of innovativeness. However, as the reviewers pointed out, the paper does not clearly describe its contributions and technical details, and the description of the experiments is also unclear, with some design lacked reasonable explanations. The presentation of experimental results also appears insufficient. 

Therefore, I believe the paper is not yet ready for publication at ICLR.