AIR: Amortized Image Reconstruction Framework for Self-Supervised Feed-Forward 2D Gaussian Splatting

Abstract: 2D Gaussian splatting provides an efficient explicit representation for image reconstruction, but existing methods still require costly per-image iterative optimization or rely on handcrafted priors for primitive allocation. We present AIR, a self-supervised feed-forward framework that amortizes iterative Gaussian fitting into a single network pass, eliminating per-image test-time optimization. AIR adopts a stage-wise residual architecture that progressively predicts additional Gaussian primitives from reconstruction residuals, together with an explicit Stage Control mechanism that activates new primitives only in under-reconstructed regions. A Predict--Optimize--Distill training strategy stabilizes multi-stage prediction by distilling short-horizon optimized Gaussian increments back into the predictor. The stabilized predictor is then jointly finetuned across stages and equipped with an image-adaptive quantizer for compact Gaussian storage. Experiments on Kodak and DIV2K show that AIR achieves better reconstruction quality than representative Gaussian-based baselines while reducing encoding time to 160--300\,ms. Code: https://github.com/whoiszzj/AIR.git

• The paper presents a novel framework that replaces per-image optimization with a self-supervised, stage-wise feed-forward network to predict 2D Gaussian representations.
• It employs a Predict–Optimize–Distill (POD) training cycle and explicit spatial control, achieving high fidelity with reduced encoding time and compact storage.
• The method establishes a scalable, efficient approach for real-time image reconstruction, demonstrating superior PSNR, MS-SSIM, and LPIPS on DIV2K and Kodak datasets.

Amortized Image Reconstruction via Self-Supervised Feed-Forward 2D Gaussian Splatting

Introduction and Motivation

Explicit image representations based on 2D Gaussian splatting enable efficient and faithful reconstruction, facilitating compact visual representation and fast decoding compared to dense pixel grids and coordinate-based neural implicit representations. While methods such as GaussianImage and Image-GS attain high-fidelity results, they depend on computationally expensive per-image iterative optimization. Recent attempts to overcome this with learned initialization (e.g., Instant-GI) still suffer from reliance on handcrafted priors for Gaussian allocation and residual test-time optimization, precluding end-to-end scalability and fully self-supervised approaches.

AIR (Amortized Image Reconstruction) introduces a new framework that replaces per-image optimization with a stage-wise, self-supervised feed-forward network for direct and adaptive prediction of 2D Gaussian representations. AIR eliminates handcrafted priors, and test-time optimization, amortizing Gaussian fitting across entire datasets for efficient and high-quality reconstruction.

Figure 1: Overview of AIR's architecture, highlighting stage-wise residual Gaussian prediction, explicit stage-control for spatial allocation, and the POD (Predict–Optimize–Distill) training cycle for stable learning.

Methodology

Stage-Wise Residual Gaussian Prediction

AIR reconstructs images progressively via multiple prediction stages. At each stage, the current network models the residual error left from previous reconstructions, predicting Gaussian increments ($\Delta G_i$) that correct under-reconstructed regions. The additivity of the 2D Gaussian renderer enables disjoint incremental improvements.

Each stage comprises:

• Residual input: $E_{i-1} = I_{gt} - I_{i-1}$.
• Network $f_{\theta_i}$ predicts $\Delta G_i$ from $E_{i-1}$.
• Additive rendering: $I_i = I_{i-1} + \mathcal{R}(\Delta G_i)$.
• Stage Control: Explicit mask selects spatial tokens for Gaussian activation where local PSNR/SSIM falls below prescribed thresholds, ensuring compact and quality-constrained allocation.

Figure 2: Network architecture for stage-wise Gaussian prediction based on hierarchical ViT encoding and local patch embedding.

This explicit spatial activation is structurally analogous to content-adaptive bit allocation in compression, but applies directly to the representation primitives.

Predict–Optimize–Distill (POD) Training

Supervising feed-forward Gaussian prediction directly with rendering loss leads to unstable, divergent training due to the under-determined inverse mapping between images and Gaussian parameters. POD addresses this with a three-phase block for each stage:

• Predict: Produce $\Delta G_i$ from the network.
• Optimize: Locally refine a detached copy $\widehat{\Delta G}_i$ of the prediction via gradient descent over the residual rendering loss for a short horizon.
• Distill: Use Gaussian-space regression to supervise $\Delta G_i$ with $\widehat{\Delta G}_i$ as target, stabilizing learning by mimicking local iterative refinement in a feed-forward manner.

Figure 3: Training curve comparison, showing the stability benefit of POD over direct image-space supervision.

After POD pretraining, finetuning is performed with image-space rendering loss over all accumulated predictions, coordinating joint optimization across stages.

Quantization and Compact Storage

AIR supports direct quantization for compact storage without post-prediction per-image range calibration. A global, learnable quantizer is combined with image-adaptive offsets predicted from the ViT [CLS] token, adjusting quantization ranges per image. This enables lightweight and effective quantized representation in a pure feed-forward mode.

Experimental Evaluation

AIR is pretrained on ImageNet and finetuned on DIV2K, with evaluation on DIV2K and Kodak. Metrics include PSNR, MS-SSIM, LPIPS, Gaussian count, and encoding time.

Reconstruction Quality and Efficiency

AIR achieves equal or better PSNR, MS-SSIM, and LPIPS than GaussianImage, Image-GS, and Instant-GI, while reducing encoding time by $E_{i-1} = I_{gt} - I_{i-1}$0--$E_{i-1} = I_{gt} - I_{i-1}$1 and eliminating per-image optimization. For a comparable number of Gaussians, AIR reconstructs images in 160–300 ms.

Figure 4: AIR attains superior PSNR and substantially lower reconstruction time than baselines across the DIV2K validation set.

Qualitatively, AIR preserves fine details and structural features with a single network inference.

Figure 5: AIR's qualitative reconstructions from Kodak and DIV2K show superior preservation of image structure and perceptual quality with pure feed-forward prediction.

AIR's adaptive allocation, driven solely by residual errors and stage-control, yields structure-aware Gaussian density distributions comparable to PPM-driven baselines, without reliance on handcrafted priors.

Figure 6: AIR produces spatially adaptive, context-aware Gaussian allocations, tightly matching underlying image structure without requiring pseudo-label–driven density maps.

Ablation Analysis

• Stage-wise vs. One-shot Prediction: Progressive prediction with explicit stage control yields similar or better compactness-reconstruction tradeoff, offering explicit rate-distortion operating points.
• Stage Control: Joint PSNR/SSIM thresholds yield the best utilization–quality tradeoff; stricter thresholds increase Gaussian count and fidelity at diminishing returns.

Figure 7: Gaussian utilization and LPIPS as a function of stage-control thresholds, illustrating tradeoffs in compactness and perceptual distortion.

• POD Training: Crucially prevents training collapse under image-space loss, stabilizing residual learning and maintaining feature diversity.

Figure 8: PCA visualization reveals POD's ability to maintain multi-dimensional, information-rich feature spaces across stages, preventing degenerate solutions.

• Quantization: AIR's adaptive quantization achieves near-original quality with low overhead and does not require per-image iterative search for parameter range.

Progressive Stage-wise Visualization

AIR's intermediate reconstructions validate the error-correcting nature of stage-wise inference: each subsequent stage reduces highly localized errors, focusing corrective effort on persistent under-reconstructed regions.

Figure 9: Visualization of AIR's progressive refinement—initial stages capture coarse content, while later stages focus on structurally complex or high-frequency residuals.

Discussion and Implications

AIR demonstrates that 2D Gaussian-based image representation can be effectively amortized through a residual, stage-wise, and strictly feed-forward architecture:

• Practical Implications: Eliminates runtime optimization, making Gaussian splatting viable for real-time applications and scalable large-scale deployment. End-to-end differentiability and lack of hand-crafted processing facilitate training on diverse datasets.
• Theoretical Implications: The success of residual stage-wise prediction with explicit spatial control offers a general strategy for parameter-efficient, adaptive explicit representations. POD unifies self-supervised regression and amortized optimization, suggesting dynamically refined learned targets can stabilize difficult prediction tasks over ambiguous inverse mappings.
• Future Directions: Further improvements may arise from parameter sharing across stages, incorporating more sophisticated entropy coding models, and extending the approach to natural video or higher-dimensional visual signals. Hybridization with implicit neural representations or meta-learning techniques could yield even more compact, generalizable image encodings.

Conclusion

AIR establishes a robust, efficient, and fully self-supervised pipeline for explicit 2D Gaussian image reconstruction. By reframing Gaussian fitting as progressive error-correcting prediction with explicit fidelity-driven spatial allocation and by leveraging the POD training paradigm, AIR advances the state of explicit visual representation. It achieves low-latency, high-fidelity reconstruction previously unattainable without hand-engineering or slow iterative optimization, offering a promising path for scalable image compression and advanced signal representation in vision systems.