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

Dear Reviewer, Please refer to the global response for a discussion of points that was shared across the reviews, including: * a discussion of contributions * ablation attributing the effect from better reconstruction and better angle selection Further details to your specific questions are below. *"The structure of the experimental section is confusing, particularly section 4.2. In a similarly light, the paragraph from lines 468-473 lacks context. It is unclear which Table/Figure the analysis is discussing."* * Thank you, will update the experimental section with additional experiments for the final version of our paper, and address your feedback. *"I'm not fully convinced about the practical advantage of the active sampling with the diffusion approach. "* * As mentioned above, the main application of our technique is in synchrotron X-ray facilities, where experiments can take several hours or days. In this case, the computational overhead of DAL is negligible, and the gains in quality from AL strategies have a significant impact in beam time and in costs. We agree however that these gains are less significant in medical imaging, which, however, was never intended as the primary application of our work. In fact, one of our contributions is to show that AL strategies do not work for the lung data set, where there is little to no advantage over uniform sampling. *“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.”* * Thank you for noting this clash in notation, we will change the notation to avoid confusion. *“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)?”* * The order in which the empirical expectation is computed makes no difference for linear forward models (CT, MRI). In non-linear forward models, the order makes a difference. Taking the expectation first and then applying the forward model corresponds more closely to a committee-base method, which aims to distinguish a candidate model (the mean estimate) from alternative models via the (Gaussian) KL divergence. Applying the forward model first and then taking the expectation resembles more closely the uncertainty sampling approach, using the variance of the observations to guide the acquisition.

Dear Reviewer, Thank you for the feedback and evaluation of our work. Please refer below for the answers to your questions. *Why are medical images not suited for this approach?* * Our results show that the benefit of Active Learning is dataset specific, i.e., it depends on the intrinsic structure of the samples. In the case of the medical images in the Lung dataset we tested, the benefit that AL has over uniform sampling is almost negligible. This is not the case for structured dataset like the composite and chip dataset studied in our paper. So while the diffusion-based reconstruction technique can be applied for medical images and will have state of the art reconstruction quality, our claim is that the results of the active acquisition will be almost indistinguishable from those of obtained with uniform sampling. For MRI the benefit of AL exists, but it is small. Besides the little advantage of active acquisitions in the medical datasets tested, there are also safety considerations in high-stakes medical settings that might preclude direct use of diffusion-based reconstructions. On the other hand, we claim these techniques are better suited for synchrotron facilities, where imaging times are long, and samples often have distinct structures such as the Chip and Composite data. *How does this model perform with samples that are slightly out of the distribution the diffusion model was trained on?* * If the samples are far from the trained distribution, the reconstructions of the diffusion model are not accurate and the method could fail. However, if samples are relatively close, even though the reconstructions might not have high PSNR at the beginning, the structure of the predicted image can guide the active learning process and still provide information on the directions to sample. In this case one would likely need more measurements to avoid hallucinations from out of distribution samples. Another line of research is to use the approach proposed by Barbano et al. (2023) (cited in our paper) which proposes a steerable conditional diffusion model designed to adapt to out-of-distribution scenarios in imaging inverse problems. In this case, one could trade-off some quality for robustness. This is a really interesting area of research that deserves further attention, and we believe that our paper can further spark interest in this problem. *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?* * We believe our approach can cope better with damage than other methods like FBP, as the final output is forced to lie within the training distribution. Whether this is desirable or not, depends heavily on the application and the goal of the experimenter; if the sample was damaged to the point it is no longer in the distribution, then we are back in the case discussed above. However, we have not benchmarked the capabilities of our model to deal with radiation damage. This is however orthogonal to the Active learning strategy proposed, and is indeed an interesting line of research of diffusion-based reconstruction methods for CT. *Why are the smaller images first cropped and then rescaled? The distribution changes when rescaling images.* * The datasets tested in our paper come from real tomographic reconstructions. However, the tomograms have a few thousand pixels per side, this is not a size with which we can do extensive benchmarking, and are at the edge of what can be done with modern diffusion models. Therefore, we used crops of these tomograms. The images at 512x512 have no scaling and reflect the real pixel resolution of underlying distribution. For 128x128, the 128x128 crops have a small field of view, so we took 256x256 images and downscaled them to 128x128. This indeed changes the distribution, but does so fairly for all benchmarked algorithms. This 128x128 size is mainly used to run extensive tests and have statistical significance.

Dear Reviewer, Thank you for the feedback on our work. Please see our response to all reviewers which addresses points shared by several reviewers, in particular: * comparison and discussion of further baselines * evaluation on real MRI data and discussion Additional details to specific questions are below. *“(1) Comparison with activate learning method is preferred, such as the method proposed in [1]. Please compare the reconstruction result and inference time with the same number of projection angle."* * The mentioned paper learns a global sampling distribution that should work 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 arbitrary, so learning a fixed distribution that is not sample-dependent will suffer even when presented with the same sample with an arbitrary rotation. For this reason, we chose to not add further non-adaptive baselines at this point. *“(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.”* * There are indeed several projection settings, and the proposed DAL strategy in principle applies to any forward model. In this paper, we decided to focus on parallel beam geometry given the importance in synchrotron X-ray experiments with parallel beam geometry, where time savings coming from AL are most dire. We have extended our experiments to work with MRI, and we could include fan-beam geometry in the final version of the paper. *“(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.”* * While the inference of a single image can take from 10 to 20 seconds, the inference of k samples can be done in almost the same time by batching the inference in the diffusion model. We achieved between 20 to 30 seconds for an entire loop of the DAL algorithm. For scientific imaging in synchrotron X-ray facilities, where repositioning the sample can take several minutes, our running times can be included in real-time experimental setups. For faster setups, one could mitigate the computational cost by taking fewer diffusion steps in DDIM sampling, or by taking fewer consistency steps. This provides a trade-off between time and quality. Our technique can be further accelerated by choosing more than one angle at each iteration of the AL loop.In the case of MRI, we propose a new accelerated version that can help further mitigate these issues. *“(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”* * In Figure 6 of the Appendix, we include a comparison with DPS and Hard data consistency, which shows that our method achieves the same or better PSNR for all tested sparsity settings. Since Figure 4 and 5 of our paper show already the significant advantage of DAL over taking uniform projection angles. These two remarks combined provide DAL with a clear advantage over DPS (or any other method) using uniform projection angles. *"(1) Plot the distribution of selected angle of different datasets. (2) The shape of the object has influence of the selected angel?"* * We added the visualization of the angles in Figures 13-15 in the Appendix. Note that diffusion active learning selects a non-uniform angle distribution for the chip and composite data, while choosing a close to uniform distribution on the lung data. *"(3) The setting of projection geometry must be given."* * We are currently using parallel beam geometry, and this is now explicitly mentioned in line 370 of the updated PDF. *"(4) The definition of the notation must be given such as x^* in algorithm 1. "* * Thank you, this is fixed. *"(5) Testing on real projection data can verify the value of the proposed method."* * See the additional experiments on the fastMRI dataset and the discussion above. [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

Dear Reviewer, We ran additional experiments on the chip, composite, and lung datasets using 2D fan-beam geometry. Please find an anonymous link to the results here: https://drive.proton.me/urls/DV7FBVGQ3G#O68PZXnzeajY We found that diffusion active learning consistently and significantly outperformed both adaptive and non-adaptive baselines on the chip and composite datasets, while showing little advantage on the lung dataset (as also for parallel beam geometry). In other words, the results for parallel-beam geometry carry over to the fan-beam geometry, with a slightly less pronounced advantage for AL over uniform. Likely due to the fact that the information of a specific direction is spread over a larger set of angles in fan-beam geometry. We'd like to emphasise that one advantage of our method is that we can evaluate different beam geometries and experimental settings without retraining the diffusion model. This is not the case for several prior works that require training for a fixed forward model and a fixed reconstruction method. As discussed in the shared response, we also added evaluation on yet another forward model, MRI, evaluated on real measurements of the fastMRI dataset. We strongly believe that **evaluation on three different forward models and four different CT/MRI datasets** (one including real measurements) is a sufficient validation of our approach. Furthermore, as comparison with another active learning method, we also added a Reinforcement learning benchmark for CT in the pdf as mentioned in the shared response. We are now looking into adding the active learning baseline that you suggested, and we will post an update here shortly.

Dear Reviewer, Thank you for your detailed evaluation. Please see our response to all reviewers which addresses the following points in the review: * Concurrent work by Elata, et al ECCV 2024 * Reinforcement learning baselines * fastMRI benchmark * simulated data vs real measurements Additional details to specific questions are below. *"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?"* * We would like to point out that in our method, when doing inference using the diffusion model, we compute gradients of the consistency loss (eq (2) of our paper) to guide the diffusion model for posterior sampling. Equation (2) however can be extended by using any regularization that is beneficial for the reconstruction. We have tried adding TV and it provides a marginal gain in our tested examples. Nevertheless, we decided not to include it in our write up, as its use is orthogonal to the main message of our paper. Other regularizations can also be incorporated. * For our method, we opted for the use of a gradient descent approach that iteratively refines the prediction \hat{x_0} by minimizing eq (2). However, one could envision using other iterative approaches like SART, SIRT or other regularized gradient descent methods to guide the diffusion process. *"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."* * Indeed, there is a plethora of sparse-view reconstruction methods used in the literature. However, Diffusion based approaches like DDRM (Kawar et al.), Diffusion Posterior Sampling (Chung et al. 2022) and Hard Data Consistency (Song et al 2023) have been already benchmarked heavily against other classical sparse-reconstruction methods, and they showed remarkable improvements in terms of reconstruction quality. Our technique builds on top of them and improves them further as can be seen in Figure 6 comparing the performance of our reconstruction method. Thus we retain by transitivity the advantage over classical sparse-view reconstruction methods. We will make sure to highlight this in the final version of the paper. *“Technical contribution” / “primary technical delta”* * Note that we added an ablation comparing soft- and hard data consistency demonstration. * As highlighted above, our work makes several contributions, in particular demonstrating the benefits of using a learned diffusion prior for active learning; and how this effect is dataset dependent. See the global response for a detailed discussion. *“Runtime requirements are not reported at all.”* * Figure 5 (right) in the original submission shows runtimes for all methods; note that this is only a qualitative assessment as performance improvements are likely possible for all methods. In the context of long data acquisition times of X-ray nano-tomography (up to several days), the cost of performing diffusion posterior sampling is offset by the improved sample efficiency, even for relatively expensive diffusion models. In addition, our proposed soft-data consistency posterior sampling is significantly faster compared to prior works (see additional experimental evaluation).

## Concurrent Work by Elata et al, ECCV 2024 We thank Reviewer hvxT for raising awareness about the concurrent work by Elata et al (ECCV 2024). We carefully went over this work, and indeed, the proposed method by Elata et al. is similar to ours, and coincides in the case of linear forward models. However, there are several technical differences, also in the presentation and evaluation: * Elata et al. motivate their approach via a PCA decomposition motivated by the linear inverse problem setup, whereas our paper derives the proposed algorithm as a general instance of uncertainty sampling. In the linear case, the resulting acquisition strategies effectively coincide. However, our proposed maximum variance acquisition function provides a perspective that also applies to non-linear forward models, while it is perhaps less clear how to extend the linear PCA to non-linear forward models. * Their work uses DDRM-based inference methods, which have been outperformed by new conditional sampling approaches for inverse problems like Diffusion Posterior Sampling (Chung et al. 2022) and Hard Data Consistency (Song et al 2023). The sampling method we propose in this paper, matches or surpasses both DPS and Hard Data Consistency in terms of speed and quality for CT reconstructions (see Figure 6 with additional experimental evaluation). * Our evaluation focuses on CT reconstruction, while we found the experiments of Elata et al on CT data to be very limited. The results reported by Elata et al on CT data only show a single example, partially contradicting our evaluation on a larger set of lung scans, where we found no gains of using active learning compared to a uniform allocation. However a single example cannot be taken as representative of the whole data distribution. But even in that single example, we remark that the results by Elata et al report only a small gain compared to the uniform allocation (which can easily vanish when evaluating on a larger set of images). The reported numbers also do not reflect the variability of the random design baseline - any design has equal probability under the uniform distribution and clearly some will perform better and other worse. * The benchmarks of Elata et al on CT and MRI data are difficult to interpret, as they do not report confidence estimates / standard errors and less fidelity in terms of the number of angles/rows selected. * Elata et al. do not report findings that indicate that expected gain from active learning are strongly dataset dependent. * Elata et al. do not report the same baselines as we do, for example the Laplace approximation, which has been proposed in the context of CT before: https://arxiv.org/abs/2207.05714 * Our work also compares different acquisition functions. * Our new accelerated MRI inference provides much faster inference than that of DDRM.

### MRI Setting and fastMRI data While we decided to focus on CT for this paper, our technique is general enough to work with linear and non-linear inverse problems. In particular, the technique directly extends to the MRI setting. **The updated paper includes preliminary evaluation on the fastMRI dataset**. Like in the CT setting, we found that diffusion active learning outperforms other generative models that we consider, and non-adaptive baselines. See Appendix B in the updated pdf. We will add a RL-based baselines for the final version of our paper. We also tested on the same datasets that we used for CT, using a synthetic MRI forward model to produce the observations. We show that DAL outperforms the other active learning strategies on our test datasets (these results are currently not included in the paper as evaluating the MRI model on CT data is not realistic). However, we also noticed the dependency on the structure of the dataset, where the gap between uniform sampling and DAL sampling narrows for non-heavily structured datasets, in this case the Lung dataset. In Appendix B of the updated pdf, we propose a new technique to accelerate MRI reconstructions using diffusion-based posterior sampling. This technique improves speed by an additional factor of 4x when compared with our baseline gradient descent strategy (see Figure 9), which is already faster than DPS and Hard Data Consistency (see Figure 6). ### Ablation: Is the advantage of diffusion active learning due to the reconstruction method or better angle selection? Reviewers CoA4 and 8GRq raised **the question if the reported gain is due to better angle selection or due to better reconstruction, or a combination of both**. To answer this question, we **re-evaluated the sequence of angles selected by Bootstrap approach, and performed the reconstruction using the diffusion model**. As expected, this improves the PSNR of the reconstruction, but still distinctly below the quality achieved by diffusion active learning. This showcases that the best reconstruction is achieved by the combination of both: The diffusion model captures the data distribution, and the angles selected by diffusion active learning exploit the data distribution in a way that a distribution-independent approach cannot; Bootstrap and Swag use only information obtained from the current sample, and therefore intuitively cannot “reason” about the posterior distribution as the diffusion model does. Note also that the gains are not purely from better angle selection, as there is a significant gap between uniform and active selection on the composite and chip data. ### Ablation of our Soft Consistency vs Hard Data Consistency To highlight one of our technical contributions better, we performed an **ablation of soft vs hard data consistency**. Refer to Figure 6 and the extended discussion in subsection "Soft Data Consistency and Early Stopping" inAppendix A. While both our method and Hard Data Consistency (Song et al 2023) achieve similar PSNR as a function of sparsity, our method is more efficient in terms of computation time and number of gradient steps, with up to 5x speed up when many measurements are present. ### Using Simulated Data vs Real Measurements * Fortunately, forward models of CT and MRI are very well understood and the quality of several reconstruction algorithms depends heavily on their accuracy. Thus, there is little "unfair" advantage in generating synthetic projections. In fact, for the FastMRI dataset, the difference between real measurements and synthetic projections produced by our forward model is always less than 1e-7. * One may also argue that the “inverse crime” problem is committed already in real data sets (like FastMRI), as the provided “ground-truth” is a reconstruction from the provided measurements assuming a forward model. Hence, it is not surprising that the measurements under the synthetic forward models match the real measurements extremely well. Therefore, this problem cannot be avoided completely even when working with real data. * This is a common practice in similar seminal works like Diffusion Posterior Sampling (Chung et al. 2022) and Hard Data Consistency (Song et al 2023). Synthetic data guarantees the existence of a ground truth that otherwise can introduce biases during testing. We believe that our techniques should be tested with real data, and we currently collaborate with researchers at a synchrotron facility to test this algorithm with real measurements. However, we are confident that the results on synthetic data showcase the benefits of our method.