Equivariant Splitting: Self-supervised learning from incomplete data

1. Problem Statement

The paper addresses inverse problems in imaging, formulated as $y = Ax + \epsilon$, where $x$ is the ground truth image, $y$ is the observed measurement, $A$ is a known forward operator (often rank-deficient), and $\epsilon$ is noise.

• The Challenge: In many applications (e.g., medical MRI, CT, astronomy), obtaining ground-truth images ($x$) for supervised training is expensive or impossible.
• The Specific Setting: The authors focus on the most challenging scenario: single incomplete forward operators. Unlike scenarios with multiple varying masks (where "splitting" methods work) or denoising (where $A=I$), here the operator $A$ is fixed and incomplete, meaning the nullspace of $A$ cannot be explored through operator diversity alone.
• Limitations of Existing Methods:
  • Measurement Splitting: Requires multiple operators or varying masks to access the full signal space. Fails with a single fixed operator.
  • Equivariant Imaging (EI): Assumes the image distribution is invariant to transformations (e.g., rotation, translation). While it works with single operators, it is computationally expensive (requiring 2–3 network passes per iteration) and does not guarantee convergence to the optimal Minimum Mean Squared Error (MMSE) estimator.

2. Methodology: Equivariant Splitting (ES)

The authors propose Equivariant Splitting (ES), a self-supervised strategy that combines the computational efficiency of splitting losses with the theoretical guarantees of equivariant imaging.

A. Core Concept: Virtual Multi-Operator Structure

The method relies on Assumption 1: The distribution of ground truth images $p(x)$ is invariant under a group of transformations $\mathcal{G}$ (i.e., $p(T_g x) = p(x)$).

• By applying a transformation $T_g$ to the image, the measurement equation becomes:
  $$y = Ax + \epsilon = A(T_g T_g^{-1} x) + \epsilon = (A T_g) (T_g^{-1} x) + \epsilon$$
• This reinterprets the single measurement $y$ as being generated by a virtual forward operator $A_g = A T_g$ acting on a virtual ground truth $x_g = T_g^{-1} x$.
• Since $x$ and $x_g$ are statistically indistinguishable (due to invariance), the problem effectively becomes a multi-operator problem where the operators are $\{A T_g\}_{g \in \mathcal{G}}$.

B. The Loss Function

The ES loss leverages this virtual multi-operator structure to apply a splitting loss:
$$L_{ES}(y, A, f) = \mathbb{E}_{g} \left[ L_{SPLIT}(y, A T_g, f) \right]$$
Where $L_{SPLIT}$ divides the measurement $y$ into an input part ($y_1$) and a target part ($y_2$) corresponding to the virtual operator $A T_g$. The network is trained to predict the full measurement $y$ from the partial input $y_1$.

Key Theoretical Result (Theorem 1):
Under the assumption of noiseless measurements and $G$-invariant $p(x)$, if the matrix $Q_{A_1} = \mathbb{E}_g[(A T_g)^\top (A T_g)]$ has full rank, the global minimizer of the ES loss is the MMSE estimator ($f^*(y) = \mathbb{E}[x|y]$). This is a significant improvement over standard EI, which does not guarantee MMSE optimality.

C. Equivariant Architectures and Efficiency

A major bottleneck in EI is the need to evaluate the network for every transformation $g$. The authors introduce a new definition of equivariant reconstructors for inverse problems:
$$f(y, A T_g) = T_g^{-1} f(y, A)$$
They prove (Theorem 2) that common architectures (Artifact Removal Networks, Unrolled Networks, MAP, MMSE) satisfy this property if their internal denoising blocks are equivariant.

• Theorem 3 (Synergy): If the network $f$ is an equivariant reconstructor, the ES loss simplifies mathematically to the standard splitting loss $L_{SPLIT}(y, A, f)$.
• Practical Impact: This allows the model to be trained using a single network pass per iteration (just like standard supervised learning), eliminating the computational overhead of EI while retaining its theoretical benefits.

3. Key Contributions

1. New Definition of Equivariance: Defined equivariance specifically for reconstruction functions $f(y, A)$, showing that standard architectures (like unrolled networks with equivariant denoisers) naturally satisfy this property.
2. Equivariant Splitting (ES) Loss: A novel self-supervised loss that unifies splitting and equivariant imaging. It provides an unbiased estimator of the supervised loss and guarantees convergence to the MMSE estimator under mild conditions.
3. Computational Efficiency: By combining ES with equivariant architectures, the method achieves the performance of EI without the 2–3x computational cost, requiring only a single forward pass per training step.
4. Theoretical Guarantees: Proved that the method yields the optimal MMSE estimator, a guarantee previously missing in single-operator self-supervised methods.

4. Experimental Results

The authors evaluated ES on four diverse inverse problems: Image Inpainting, Compressive Sensing, Accelerated MRI, and Sparse-View CT.

• Performance:
  • ES consistently outperforms Equivariant Imaging (EI) and other baselines (SURE, MC, IDFT) in terms of PSNR and SSIM.
  • In highly rank-deficient settings (e.g., Compressive Sensing with low compression rates, Sparse-View CT), ES approaches the performance of fully supervised methods, whereas EI often lags behind.
  • Example: In MRI (8x acceleration), ES achieved 28.54 dB PSNR vs. 27.88 dB for EI, and 28.74 dB for Supervised.
• Ablation Studies:
  • Architecture: Using equivariant architectures (e.g., UNets with shift-invariance or Reynolds averaging for rotation) significantly boosts performance when paired with the splitting loss.
  • Corollary 1 Validation: Experiments confirmed that if the forward operator $A$ is itself equivariant to the chosen transformations (e.g., a fixed grid mask with shift invariance), the method fails to learn (performance drops), validating the theoretical necessity of non-equivariant operators.
• Robustness: The method remains stable even when noise distributions are unknown or when the invariance assumption is only approximate (e.g., medical images).

5. Significance

• Bridging the Gap: This work successfully bridges the gap between "splitting" methods (efficient but require multiple operators) and "equivariant" methods (single operator capable but computationally heavy and theoretically weaker).
• Practical Impact: It enables state-of-the-art reconstruction in scenarios where ground truth is unavailable and data acquisition is constrained to a single, fixed, incomplete measurement model (common in medical imaging and remote sensing).
• Efficiency: By reducing the training cost to that of standard supervised learning while maintaining self-supervised capabilities, it makes advanced inverse problem solving more accessible and scalable.
• Theoretical Advancement: It establishes that architectural constraints based on equivariance are not just a heuristic for generalization but a fundamental tool for solving ill-posed inverse problems with provable optimality.