Equivariant Multiscale Learned Invertible Reconstruction for Cone Beam CT
Nikita Moriakov, Jan-Jakob Sonke, Jonas Teuwen
TL;DR
This work addresses the challenge of high-quality CBCT reconstruction under memory and speed constraints by introducing LIRE+, a rotationally-equivariant multiscale learned invertible primal-dual scheme. By combining patch-wise memory efficiency, invertible networks, and P4-equivariant primal blocks, LIRE+ delivers competitive or superior image quality with substantially fewer parameters and faster inference compared to prior methods. On thorax data, it outperforms baselines and remains robust to orientation; on head & neck data, it performs comparably without finetuning and shows notable gains after targeted finetuning, indicating strong generalization and practical potential for clinical CBCT workflows. Overall, LIRE+ advances fast, memory-efficient CBCT reconstruction with improved robustness and scalability for 3D imaging tasks.
Abstract
Cone Beam CT (CBCT) is an essential imaging modality nowadays, but the image quality of CBCT still lags behind the high quality standards established by the conventional Computed Tomography. We propose LIRE+, a learned iterative scheme for fast and memory-efficient CBCT reconstruction, which is a substantially faster and more parameter-efficient alternative to the recently proposed LIRE method. LIRE+ is a rotationally-equivariant multiscale learned invertible primal-dual iterative scheme for CBCT reconstruction. Memory usage is optimized by relying on simple reversible residual networks in primal/dual cells and patch-wise computations inside the cells during forward and backward passes, while increased inference speed is achieved by making the primal-dual scheme multiscale so that the reconstruction process starts at low resolution and with low resolution primal/dual latent vectors. A LIRE+ model was trained and validated on a set of 260 + 22 thorax CT scans and tested using a set of 142 thorax CT scans with additional evaluation with and without finetuning on an out-of-distribution set of 79 Head and Neck (HN) CT scans. Our method surpasses classical and deep learning baselines, including LIRE, on the thorax test set. For a similar inference time and with only 37 % of the parameter budget, LIRE+ achieves a +0.2 dB PSNR improvement over LIRE, while being able to match the performance of LIRE in 45 % less inference time and with 28 % of the parameter budget. Rotational equivariance ensures robustness of LIRE+ to patient orientation, while LIRE and other deep learning baselines suffer from substantial performance degradation when patient orientation is unusual. On the HN dataset in the absence of finetuning, LIRE+ is generally comparable to LIRE in performance apart from a few outlier cases, whereas after identical finetuning LIRE+ demonstates a +1.02 dB PSNR improvement over LIRE.