Source-linked AI summary
Invert to Learn to Invert
Patrick Putzky, Max Welling
TL;DR
Recovering an unobserved signal from measurements is nontrivial when the forward problem is ill-posed or nonlinear, and iterative inverse models face a memory cost from storing intermediate activations. The i-RIM modifies recurrent inference machines to be fully invertible, using an orthogonal 1×1 convolution layer designed for stable invertible learning of very deep models. Both i-RIM models consistently outperform baselines and sit atop the Single-Coil leaderboard; experiments train models on large 3D MRI volumes, while stable invertible learning reaches 400 layers.
Problem
Recovering an unobserved signal from measurements is nontrivial when the forward problem is ill-posed or nonlinear, and iterative inverse models face a memory cost from storing intermediate activations.
Source: §1 · p. 1, §1.2 · p. 2Method
The i-RIM modifies recurrent inference machines to be fully invertible, using an orthogonal 1×1 convolution layer designed for stable invertible learning of very deep models.
Source: §2 · p. 3, §2.2 · p. 5Results
Both i-RIM models consistently outperform baselines and sit atop the Single-Coil leaderboard; experiments train models on large 3D MRI volumes, while stable invertible learning reaches 400 layers.
Source: §1.5 · p. 3, §3.4 · p. 8Takeaways & Limitations
The approach supports deep iterative reconstruction on large 3D MRI volumes, beyond prior approaches limited to 2D slices.
Source: §1.5 · p. 3The accelerated MRI setup assumes single-coil measurements and uses the identity function for Ψ.
Source: §3.1 · p. 6
Abstract
from arXiv · showhide
Iterative learning to infer approaches have become popular solvers for inverse problems. However, their memory requirements during training grow linearly with model depth, limiting in practice model expressiveness. In this work, we propose an iterative inverse model with constant memory that relies on invertible networks to avoid storing intermediate activations. As a result, the proposed approach allows us to train models with 400 layers on 3D volumes in an MRI image reconstruction task. In experiments on a public data set, we demonstrate that these deeper, and thus more expressive, networks perform state-of-the-art image reconstruction.
1 Introduction
The introduction frames inverse problems as recovering an unobserved signal from measured data, a difficult task when the forward problem is ill-posed or nonlinear. It presents invertible learning as a way to reduce memory costs in iterative inverse models and enable deeper models.
1.1 Learn to Invert
- Iterative inverse models mimic iterative optimization through recursive updates, and are typically trained with simulated measurements paired with ground-truth signals.Their update function uses an estimate and an auxiliary memory variable at each iteration.Source: §1.1 · p. 1
1.2 Invert to Learn
- Invertible learning avoids storing intermediate activations, addressing the depth-versus-memory trade-off in training iterative inverse models.The proposed approach uses reversible networks to restore activations; its stated memory complexity is O(1), with computational complexity O(L) for network depth L.Source: §1.2 · p. 2
1.3 Invertible Neural Networks
- Invertible neural networks are used mainly in likelihood-based generative models; the paper assumes any such network can underpin invertible learning and notes potential future unsupervised training.The unsupervised-training direction is linked to Eq. (4).Source: §1.3 · pp. 2–3
1.4 Recurrent Inference Machines
- RIM updates an estimate and auxiliary memory from a data-consistency signal, but its original formulation is not invertible even if its component network is.The data-consistency term keeps estimates tied to measured data under the forward model.Source: §1.4 · p. 3
1.5 Contribution
- The contributions are a fully invertible iterative inverse model, a layer for stable training of 400-layer networks, and reconstruction of large 3D MRI volumes.The authors also identify potential relevance to other large-observation tasks, including radio-astronomy synthesis imaging.Source: §1.5 · p. 3
2 Method
The method makes Recurrent Inference Machines fully invertible, enabling training with O(1) memory, and introduces an orthogonal-convolution layer intended to stabilize deep invertible learning. It also specifies residual-block and related-work context for the proposed model.
2.1 Invertible Recurrent Inference Machines
- The proposed i-RIM combines a modified, fully invertible RIM update with layer-wise inversion to support training at O(1) memory complexity.The modification addresses the standard RIM’s non-invertible update structure; the method assumes its network function can be inverted layer by layer.Source: §2 · p. 3, §2.1 · p. 4
2 Method
- The method’s two components are a fully invertible RIM modification and a new invertible layer designed for stable learning in very deep i-RIMs.The authors identify numerical errors in very deep networks as a source of instability in common invertible layers.Source: §2 · pp. 3–4
2.3 Related Work
- The method builds on invertible-network designs and RIM applications in MRI and radio astronomy, while noting that prior invertible inverse-problem work assumes an unknown forward problem.The authors suggest their results may extend to other RIM applications.Source: §2.2 · p. 4, §2.3 · p. 6
2.2 An Invertible Layer with Orthogonal 1x1 convolutions
- The orthogonal 1×1 convolution embeds channels while avoiding vanishing or exploding gradients, and its transpose provides the inverse used during training.Projecting outputs back to the original basis keeps the convolution’s action local to its layer.Source: §2.2 · p. 5
2.2.1 Orthogonal 1x1 convolution
- The method parameterizes orthogonal convolution matrices with Householder reflections and uses D = 3 reflections in its experiments.Using a subset of reflections reduces parameter count and construction cost.Source: §2.2.1 · pp. 5–6
2.2.2 Residual Block with Spatial Downsampling
- Each residual block downsamples spatially, applies a 3×3 convolution, then reverses the downsampling; a gated linear unit supports stable training without special weight initialization.The block uses weight normalization for all convolutional weights and disables the final convolution’s bias.Source: §2.2.2 · p. 6
3 Experiments
Experiments evaluate accelerated single-coil MRI reconstruction on fastMRI, comparing U-Net, RIM, and i-RIM; both i-RIMs outperform the baselines, and i-RIM 3D performs almost as well as i-RIM 2D.
3 Experiments
- The experiments compare a U-Net baseline, RIM, and i-RIM on single 2D slices, and also train a 3D i-RIM to explore the memory benefits.The data come from the public fastMRI challenge.Source: §3 · p. 6
3.1 Accelerated MRI
- The accelerated MRI task models sub-sampled K-space measurements as a Fourier-transformed image selected by a sampling matrix, assuming single-coil data.The experiments set Ψ to the identity function.Source: §3.1 · p. 6
3.2 Data
- The fastMRI data comprise 973 training, 199 validation, and 108 test volumes; 3D training excludes volumes with fewer than 32 slices.Training and validation data are fully sampled, while the test set is sub-sampled; evaluation uses the central 320 × 320 portion of each magnitude slice.Source: §3.2 · pp. 6–7
3.4 Results
- i-RIM memory consumption is mostly independent of network depth, while its 3D machine state occupies more than 1.25 GB, or 7.5% of a 16 GB GPU.The authors say this makes network depth mostly a computational consideration.Source: Table 1 · p. 7, §3.4 · p. 8
- Both i-RIM models consistently outperform the baselines on validation and test data; i-RIM 3D performs almost as well as its 2D counterpart.The authors report that all three models were at the top of the fastMRI Single-Coil leaderboard at the time of writing.Source: Table 2 · p. 7, §3.4 · p. 8
3.3 Training
- The i-RIM models use 10 invertible layers per time step, without parameter sharing across steps, and differ between 2D and 3D in their convolution dimensions.Iterative models use eight inference steps; i-RIM loss is evaluated only at the last step, while RIM loss is averaged across steps. Training uses a random mask on 1% of pixels or voxels per back-propagation step, and the machine state has 64 channels.Source: §3.3 · pp. 7–8
4 Discussion
The discussion presents invertible networks as a way to address training-memory constraints for iterative inverse models and train deep models for large-scale imaging. It also highlights a stable invertible layer and potential extensions to other structured prediction tasks and unsupervised or semi-supervised training.
4 Discussion
- The approach uses invertible neural networks to address memory issues in training iterative inverse models, enabling very deep models on a large-scale imaging task.The authors describe the resulting reconstructions as state of the art.Source: §4 · p. 8
- A new invertible layer is introduced to support stable training of deep models and is expected to be useful for other structured prediction tasks.Source: §4 · p. 8
- The authors identify unsupervised or semi-supervised training as future directions for their approach.Source: §4 · p. 8