Source-linked AI summary
Single-Stage Diffusion NeRF: A Unified Approach to 3D Generation and Reconstruction
Hansheng Chen, Jiatao Gu, Anpei Chen, Wei Tian, Zhuowen Tu, Lingjie Liu, Hao Su
TL;DR
Existing 3D methods often specialize in generation or reconstruction, motivating a unified model. SSDNeRF jointly trains a NeRF auto-decoder and latent diffusion model end to end, and reports robust performance across unconditional generation and single- or sparse-view reconstruction.
Problem
Developing one comprehensive framework that bridges 3D generation and reconstruction remains challenging, especially when NeRFs must operate from sparse observations.
Method
SSDNeRF jointly trains an expressive NeRF auto-decoder and a 3D latent diffusion model in a single end-to-end stage from multi-view images.
Results
SSDNeRF demonstrates robust performance comparable to or better than leading task-specific methods in unconditional generation and single- or sparse-view 3D reconstruction.
Takeaways & Limitations
The learned diffusion prior can be sampled unconditionally or combined with arbitrary observations for test-time NeRF reconstruction.
Takeaways & Limitations
The method relies on ground-truth camera parameters during both training and testing, and prolonged training can make the diffusion prior discontinuous and affect generalization.
Abstract
from arXiv · showhide
3D-aware image synthesis encompasses a variety of tasks, such as scene generation and novel view synthesis from images. Despite numerous task-specific methods, developing a comprehensive model remains challenging. In this paper, we present SSDNeRF, a unified approach that employs an expressive diffusion model to learn a generalizable prior of neural radiance fields (NeRF) from multi-view images of diverse objects. Previous studies have used two-stage approaches that rely on pretrained NeRFs as real data to train diffusion models. In contrast, we propose a new single-stage training paradigm with an end-to-end objective that jointly optimizes a NeRF auto-decoder and a latent diffusion model, enabling simultaneous 3D reconstruction and prior learning, even from sparsely available views. At test time, we can directly sample the diffusion prior for unconditional generation, or combine it with arbitrary observations of unseen objects for NeRF reconstruction. SSDNeRF demonstrates robust results comparable to or better than leading task-specific methods in unconditional generation and single/sparse-view 3D reconstruction.
1. Introduction
SSDNeRF targets a unified framework for 3D generation and reconstruction by learning generalizable 3D priors from multi-view images. Its single-stage training jointly learns NeRF reconstruction and diffusion priors, including with sparse views.
- The model learns a generative prior of scene latent codes with a 3D latent diffusion model.
- Single-stage training jointly optimizes NeRF reconstruction and diffusion weights end to end, addressing noisy latent codes from two-stage training.The approach is designed to blend generative and rendering biases and support sparse-view training.
- The learned diffusion priors support flexible test-time scene sampling from arbitrary observations.
- Evaluation on multiple categorical single-object scene datasets demonstrates strong overall performance and progress toward a unified 3D framework.
- SSDNeRF is proposed as a unified approach to unconditional 3D generation and image-based reconstruction.
2. Related Work
Related work includes 3D GANs, view-conditioned regression and generation, and multi-scene NeRF auto-decoders. These approaches provide generation or reconstruction capabilities but have limitations in cross-view reasoning, diversity, or faithful detailed reconstruction.
- 3D GANs: 3D GANs are mainly used for unconditional generation, while GAN inversion can perform image-based completion without guaranteed faithfulness.
- View-Conditioned Regression and Generation: View-conditioned regression methods encode input images into volume features and render supervised target views, but they cannot reason about ambiguity or generate diverse contents.This often leads to blurry results.
- Auto-Decoders and Diffusion NeRF: Multi-scene NeRF auto-decoders jointly learn shared decoder parameters and per-scene latent codes, enabling 3D completion and generation with suitable architectures.Their latent codes are described as insufficiently expressive for faithful reconstruction of detailed objects.
3. Background
The background introduces NeRF as a scene-fitting and rendering model, extends it to multi-scene auto-decoders with per-scene latent codes, and describes latent diffusion and guided sampling. It also highlights why two-stage training is problematic for sparse-view 3D tasks.
- NeRF as an Auto-Decoder: Given scene images and camera parameters, NeRF fits a scene model to reconstruct the light field in 3D space.
- NeRF as an Auto-Decoder: NeRF represents a scene’s light field through density and color functions of 3D position and viewing direction, with differentiable volume rendering composing ray samples.
- NeRF as an Auto-Decoder: Multi-scene NeRFs jointly optimize per-scene codes and shared parameters under an L2 rendering objective, treating the codes as latent variables and the plenoptic function as a decoder.
- NeRF as an Auto-Decoder: Auto-decoder generation requires a continuous low-dimensional latent space and complex decoder, which makes faithful reconstruction from observed views difficult.
- Latent Diffusion Models: Latent diffusion learns a prior over expressive scene latents by adding Gaussian noise and training a denoising network to recover clean codes.
- Latent Diffusion Models: Diffusion sampling recursively denoises Gaussian noise for unconditional generation and can use rendering-loss gradients to reconstruct scenes from test-time images.
- Latent Diffusion Models: Two-stage NeRF diffusion training can produce noisy latent patterns and makes sparse-view reconstruction exceptionally difficult without a learned prior.
4. Proposed Method
SSDNeRF unifies a triplane NeRF auto-decoder with a triplane latent diffusion model through single-stage joint training. Its learned prior supports sparse-view reconstruction, image-guided sampling, and finetuning while addressing limitations of two-stage training.
- SSDNeRF conjoins an expressive triplane NeRF auto-decoder with a triplane latent diffusion model.
- The single-stage objective jointly optimizes scene codes, prior parameters, and decoder parameters using rendering and diffusion losses.The objective combines λrendLrend and λdiffLdiff, with scene codes and model parameters updated together.
- Single-stage training constrains scene codes with both rendering and prior terms, allowing the learned prior to complete regions unseen by rendering.This is particularly beneficial when expressive triplane codes are severely underdetermined by sparse views.
- Two-stage methods omit the diffusion prior during auto-decoder training, producing biased and noisy scene codes that can hinder learning a clean latent manifold.The issue is especially relevant when inverse rendering is uncertain under sparse-view training.
- At test time, SSDNeRF performs image-guided diffusion sampling and finetunes sampled scene codes using both diffusion-prior and rendering objectives.Finetuning freezes diffusion and decoder parameters; the test-time prior weight is lower because the training prior is less reliable on a different dataset.
- Prior gradient caching accelerates reconstruction by reusing back-propagated diffusion-prior gradients across multiple Adam steps while refreshing rendering gradients.The technique addresses the higher evaluation cost of the diffusion loss during training and test-time finetuning.
5. Experiments
Experiments evaluate SSDNeRF on unconditional generation and sparse-view reconstruction using SRN Cars, SRN Chairs, and ABO Tables. The results show strong generation quality, improved single-stage training, and competitive or leading reconstruction performance across view settings.
- Unconditional Generation: SSDNeRF significantly outperforms EG3D and DiffRF on unconditional generation for ABO Tables, while improving KID over EG3D and FID over Functa on SRN Cars.The SRN Cars KID comparison is described as a clear margin, and the FID comparison against Functa as drastic.
- Unconditional Generation: KID/10^-3 improves from 6.38 with two-stage training to 3.47 with single-stage training on SRN Cars.The comparison uses the same model architecture and tuned TV regularization for the two-stage variant.
- Sparse-View NeRF Reconstruction: SSDNeRF achieves the best LPIPS across single-view and two-view reconstruction tasks, and the best performance on all relevant metrics in two-view reconstruction.The evaluation measures PSNR, SSIM, LPIPS, and FID against unseen target views.
- Sparse-View NeRF Reconstruction: The learned diffusion prior significantly boosts LPIPS and FID over rendering-loss-only finetuning, while finetuning mainly benefits views with higher overlap to the input.Single-view finetuning provides only marginal improvements over guided sampling without the diffusion loss.
- Sparse-View NeRF Reconstruction: A guidance-finetuning model reaches LPIPS 0.106 for single-view reconstruction, outperforming most prior methods that use the full training set.The sparse-view training setup can use as few as three views, whereas TV-regularized auto-decoders often produce severe geometric artifacts.
6. Conclusion
SSDNeRF combines NeRF and diffusion modeling through single-stage end-to-end training, supporting reconstruction and prior learning without first requiring dense observations. The method shows strong benchmark performance while retaining practical considerations for batch-wise rendering and known limitations.
- SSDNeRF combines the diffusion model and NeRF representation through a novel single-stage training paradigm with an end-to-end loss.
- The approach avoids requiring implicit neural fields from dense observations before diffusion training.
- Qualitative evaluations include generated images, three-view scene-code comparisons, and DDIM interpolation results.
- The method relies on ground-truth camera parameters during training and testing, while prolonged training can make the diffusion prior discontinuous.
- Rendering losses are rescaled according to the sampled ray-batch size during training and test-time reconstruction.
- Test-time gradient guidance treats the sampled ray batch as the full observation set when calculating reconstruction gradients.
B.1. Implementation Details
The implementation uses PyTorch, MMGeneration, and torch-ngp, with step decay applied to learning rate and inner-loop iterations. Reconstruction-oriented models additionally use early stopping, while Langevin correction is dataset-dependent.
- The models are implemented with PyTorch and MMGeneration, while the NeRF renderer uses torch-ngp with density-based grid pruning.
- Step decay is applied to both the learning rate and the number of inner-loop iterations during training.
- Reconstruction-purposed training stops early at 80K iterations, unlike the unconditional-generation schedule.
- Langevin correction is more effective for Chairs than Cars and is therefore omitted for the SRN Cars dataset to reduce inference time.
B.3. Training and Inference Time
Training and inference costs vary substantially by task and dataset. The implementation also includes latent bounding and denoising details, with Tanh mapping reported as potentially unnecessary under the final parameterization.
- Training uses two RTX 3090 GPUs, with 80K iterations taking around 11 hours and 1M iterations around 6 days.
- Unconditional generation takes 4.63 seconds for eight scenes using 50 DDIM steps on one RTX 3090 GPU.
- Single-view reconstruction takes around 23 seconds for eight Cars and 102 seconds for eight Chairs, including additional Langevin steps for Chairs.
- The denoising U-Net predicts noisy and denoised triplane features using v-parameterization.
- Latent codes are bounded through a Tanh layer, with raw unbounded parameters optimized during training and test-time finetuning.
- The final v-parameterization makes Tanh mapping potentially nonessential, according to the reported implementation analysis.
C.2. Additional L2 Regularization
SSDNeRF retains L2 latent regularization from Gaussian-prior formulations, but the reported analysis finds negligible impact under single-stage training. The section also references qualitative failure cases and sparse-to-dense reconstruction details.
- L2 latent regularization controls latent-code norms and discourages outlying values relative to sampling-time clipping.
- Table 6 reports sparse-to-dense reconstruction settings, including input views, finetuning iterations, learning rate, and novel-view synthesis results.
- Figure 10 shows single-view reconstruction failure cases and a sample resolving the issue by reducing guidance scale λgd.
- The regularization weight λreg is set to 0.003, while the latent dimensionality determines the normalization term.
- L2 regularization has negligible impact under the single-stage training framework.
D. Experiment Details and Additional Results
Additional experiments examine sparse-to-dense reconstruction, real-image generalization, sampling stability, and qualitative generation and reconstruction results.
- Experiment Settings: More input views require more finetuning iterations and higher learning rates, but stability requires keeping the learning rate at or below 0.08.For dense-view settings, 200 outer-loop iterations, totaling 1600 inner-loop iterations, are sufficient.
- Real-Image Reconstruction: SSDNeRF is evaluated for single-view reconstruction from real images to test generalization across a substantial synthetic-to-real domain gap.The model is trained on synthetic SRN Cars and evaluated using real vehicle images from KITTI.
- Qualitative Results and Failure Cases: Reducing the guidance scale resolves missing geometry, whereas excessive guidance can destabilize sampling and corrupt reconstructed geometries.The failure case involves a large portion of missing geometry; qualitative results include novel views and extracted meshes.
- Qualitative Results: The paper presents unconditional samples on SRN Cars, ABO Tables, and three-view SRN Cars, plus single-view comparisons on unseen Cars and Chairs objects.Single-view novel views are compared with CodeNeRF and VisionNeRF.