Memory-efficient optimization of implicit neural representations for CT reconstruction

TL;DR

This paper introduces a memory-efficient stochastic gradient method for optimizing implicit neural representations in CT reconstruction, enabling 3D imaging with reduced GPU memory without sacrificing accuracy.

Contribution

It proposes a novel gradient approximation technique that reduces memory usage in INR optimization for CT, especially beneficial for 3D imaging.

Findings

Gradient approximation significantly reduces GPU memory consumption.

Reconstruction quality remains comparable to standard methods.

Enables memory-efficient 3D cone beam CT reconstruction in sparse-view scenarios.

Abstract

Implicit neural representations (INRs) provide a parameter-efficient and fully differentiable image model for CT reconstruction. However, optimizing INRs for CT reconstruction using standard auto-differentiation techniques can be prohibitively GPU memory-intensive, especially in 3D imaging, due to the large number of INR evaluations needed to simulate ray projections. To address this issue, we propose a memory-efficient stochastic gradient approximation based on decomposing the gradient into a Jacobian-vector product that is amenable to stochastic subsampling. This approximation allows the user to trade-off between GPU memory usage and gradient approximation accuracy. Our experiments on synthetic 2D data demonstrate that gradient approximation uses far less GPU memory than standard INR training, while yielding reconstructions that are comparable in convergence behavior and mean squared error. Finally, we demonstrate that the proposed approach allows for memory-efficient 3D cone beam CT reconstruction in a sparse-view setting.