Uncertainty modeling for fine-tuned implicit functions

Authors: Anna Susmelj, Mael Macuglia, Natasa Tagasovska, Reto Sutter, Sebastiano Caprara, Jean-Philippe Thiran, Ender Konukoglu
Primary area: applications to computer vision, audio, language, and other modalities
Keywords: uncertainty, implicit functions, 3D reconstruction, occupancy networks

Abstract:
Implicit functions such as Neural Radiance Fields (NeRFs), occupancy networks, and signed distance functions (SDFs) have become pivotal in computer vision for reconstructing detailed object shapes from sparse views. Achieving optimal performance with these models can be challenging due to the extreme sparsity of inputs and distribution shifts induced by data corruptions. To this end, large, noise-free synthetic datasets can serve as shape priors to help models fill in gaps, but the resulting reconstructions must be approached with caution. Uncertainty estimation is crucial for assessing the quality of these reconstructions, particularly in identifying areas where the model is uncertain about the parts it has inferred from the prior. In this paper, we introduce Dropsembles, a novel method for uncertainty estimation in tuned implicit functions. We demonstrate the efficacy of our approach through a series of experiments, starting with toy examples and progressing to a real-world scenario. Specifically, we train a Convolutional Occupancy Network on synthetic anatomical data and test it on low-resolution MRI segmentations of the lumbar spine. Our results show that Dropsembles achieve the accuracy and calibration levels of deep ensembles but with significantly less computational cost.

Decision: accept (poster)
Reviewer scores: 6, 6, 8, 3

Reviews and discussion:
<thread>
creation date:
1730108044612

thread_type:
official_review

<thread_object>
creator: reviewer_yvca

<content>
summary:
In this work, authors introduce Dropsembles, a novel method for uncertainty estimation in tuned implicit functions. They demonstrate the efficacy of this approach through a series of experiments, starting with toy examples and progressing to a real-world scenario. Our results show that Dropsembles achieve the accuracy and calibration levels of deep ensembles but with significantly less computational cost.

strengths:
1.  The topic of this work is really interesting. By utilizing an implicit shape model, points with sparse shapes can be densified with refined shape representations.
2.  This is the first work to model epistemic uncertainty in the implicit decoder during finetuning from synthetic data.

weaknesses:
1.  The core concept of this work is uncertainty modeling. There are two innovative designs proposed in this paper. For the Elastic weight consolidation strategy, it seems that no remarkable improvements are done. Authors simply employ EWC to adapt pretrained models from dataset A. The other component is Dropsembles, please provide a more detailed description on its contribution.
2.  Since Dropsembles is aimed for the tradeoff between computational costs and accuracy. However, no quantitative results are given in the experimental part. Authors should carefully list out the comparison between dropsembles and other methods on uncertainty estimation.
3.  The improvement on lumbar spine seems to be marginal, maybe a detailed analysis is required.
4.  The effectiveness of this method is highly required for the evaluation on more challenging datasets, such as Medshapenet, etc.

questions:
Please refer to the weakness part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730637715911

thread_type:
official_review

<thread_object>
creator: reviewer_bmvl

<content>
summary:
This paper explores the topic of uncertainty modeling for decoder fine-tuning in implicit neural representations for occupancy representations. The presented method employs a combination of widely used uncertainty modeling concepts (i.e., dropout, ensembling) and the elastic weight consolidation concept stemming from continual learning, with the notion of fine-tuning, allowing to obtain multiple (fine-tuned) models at less training cost, with partially improved or comparable performance in modeling the uncertainty. This is demonstrated in experiments with toy examples (Sinusoidal decision boundary, MNIST) to medical shapes with the perseverance of reconstruction accuracy.

strengths:
Originality: While the paper builds upon technical concepts from the uncertainty modeling community and does not introduce novel technical components, I believe the explored application and conclusions are novel and impactful.

Quality: The paper is very well structured and intuitive. The presentation and writing of the method are very clear, and in combination with the experiments, the method is sound.  

Clarity: The paper is well written in all sections, especially in related work and methods, making it a charm to read. It really flows. I also really enjoyed how well the method was put into the context of related work (i.e. very thoroughly research in different areas) and how the method was described in a simple and comprehensive way. Moreover, the authors detail every aspect of the experimental design in the main paper, and together with the supplementary this makes the evaluation of the presented method easy to follow. This, together with the accessibility of the source code, facilitates reproducability of the experiments. 

Significance: While implicit neural representations are often trained - individually - for (single) images and other signals, cohort-based training is quite common for shape completion in SDFs and occupancy representations. Given the prevalence of encoder-/decoder in these settings, this paper constitutes a meaningful exploratory study in the modeling of decoder uncertainty.

weaknesses:
1) Related Work:

As detailed above, the paper is very well written and conducts a very thorough related work section. However, I think it would be important to touch upon one more area - i.e. the more recent architectures used in cohort-based training that also use encoder frameworks. Specifically, it would be interesting to discuss the advancements in [1,2], given that these constitute more recent alternatives to the DeepSDF architecture. As both works [1,2] use encoder-/decoder settings in combination with weight modulation approaches, this would also add emphasis to the relevance of this work. 

2) Experimental Design: While the experimental design is sound, I believe the paper could have explored a higher variation of datasets and application scenarios.

(a)  Experiments not only for Re-LU-based networks but also for other activation functions:
While ReLU networks were routinely used in [3,4,5] medical applications, the INR community has shifted towards other activation functions [6] or embedding/projection functions [7,8] (particularly for images). While not all (novel) shape modeling applications use these (yet?), I believe it would have been interesting to explore if the same holds true for, e.g., SIREN [6]. This ablation would have been simple, as it merely requires a change in the MLP and DeepSDF [3], which also works well with sinusoid activation functions [6]. Moreover, it has been comparatively less explored if dropout and other anomaly detection frameworks are useful in this context. Thus, for decoder-based uncertainty modeling, for example, in the context of [2], such an ablation study would be ultimately interesting and relevant.

(b) Experiments on more common (shape) datasets, especially with SDFs: The authors state that their method is relevant for both SDF and occupancy representations but only conducts experiments for occupancy networks. Given that the training setup would not change much, I believe it would have been interesting to show experiments in an established shape dataset as well (e.g. (Med-) ShapeNet). 

References:

[1] Dupont E, Loya H, Alizadeh M, Goliński A, Teh YW, Doucet A. Coin++: Neural compression across modalities. arXiv preprint arXiv:2201.12904. 2022 Jan 30.

[2] Mehta I, Gharbi M, Barnes C, Shechtman E, Ramamoorthi R, Chandraker M. Modulated periodic activations for generalizable local functional representations. InProceedings of the IEEE/CVF International Conference on Computer Vision 2021 (pp. 14214-14223).

[3] Park JJ, Florence P, Straub J, Newcombe R, Lovegrove S. Deepsdf: Learning continuous signed distance functions for shape representation. InProceedings of the IEEE/CVF conference on computer vision and pattern recognition 2019 (pp. 165-174).

[4] Mescheder L, Oechsle M, Niemeyer M, Nowozin S, Geiger A. Occupancy networks: Learning 3d reconstruction in function space. InProceedings of the IEEE/CVF conference on computer vision and pattern recognition 2019 (pp. 4460-4470).

[5] Amiranashvili T, Lüdke D, Li HB, Menze B, Zachow S. Learning shape reconstruction from sparse measurements with neural implicit functions. InInternational Conference on Medical Imaging with Deep Learning 2022 Dec 4 (pp. 22-34). PMLR.

[6] Sitzmann V, Martel J, Bergman A, Lindell D, Wetzstein G. Implicit neural representations with periodic activation functions. Advances in neural information processing systems. 2020;33:7462-73.

[7] Tancik M, Srinivasan P, Mildenhall B, Fridovich-Keil S, Raghavan N, Singhal U, Ramamoorthi R, Barron J, Ng R. Fourier features let networks learn high frequency functions in low dimensional domains. Advances in neural information processing systems. 2020;33:7537-47.

[8] Müller T, Evans A, Schied C, Keller A. Instant neural graphics primitives with a multiresolution hash encoding. ACM transactions on graphics (TOG). 2022 Jul 22;41(4):1-5.

questions:
Importance of encoder design / Application in auto-decoder setting:

Q1: Given the suitability of INRs for non-grid data, it is sometimes challenging to design an encoder for non-grid data. How was the encoder used/designed for the MNIST case, where some pixels are occluded? Did the encoder see the entire image (e.g., with a CNN-based encoder)? Sorry if I missed this; I would really appreciate an answer, even if it does not change the comparability of the different methods/baselines.

Q2: In the INR community, auto-decoder settings are quite common as well (see DeepSDF or Mehta et al.). Could the same method also be used in this case, given that the encoding of, e.g., the shape or instance stems from the same MLP, or are there any theoretical limitations that prohibit a meaningful uncertainty in this case?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730652354537

thread_type:
official_review

<thread_object>
creator: reviewer_iepp

<content>
summary:
The paper presents Dropsembles, a novel uncertainty quantification technique for neural implicit functions such as Neural Radiance Fields (NeRFs) and Occupancy Networks. This method aims to achieve the high accuracy and calibration of deep ensembles while significantly reducing computational overhead. By integrating Monte Carlo dropout with ensemble strategies, and leveraging Elastic Weight Consolidation (EWC), Dropsembles effectively handles distribution shifts and sparse, noisy data.

The primary focus is on addressing the challenges of 3D shape reconstruction from sparse and corrupted inputs. The method leverages high-quality, densely sampled synthetic datasets as shape priors to fill gaps and correct corruptions in the target data, even when distribution shifts occur due to noise and corruption. The premise is that these priors can guide the reconstruction process, ensuring more reliable outputs.

Initially, a Convolutional Occupancy Network is trained on high-quality synthetic data. During testing, the network is fine-tuned on individual sparse and noisy samples using EWC, which regularizes the adaptation process to prevent forgetting the original training data and to transfer prior knowledge effectively. This method involves creating ensembles through dropout, where binary masks generate multiple network instances. Each instance is then fine-tuned independently on the testing sample, forming an ensemble of networks initialized with correlated weights.

The paper validates the method through extensive experiments on toy datasets and low-resolution MRI segmentations of the lumbar spine. Results demonstrate that Dropsembles maintains accuracy and reliable uncertainty estimation and efficiently handles computational constraints and distribution shifts.

strengths:
- The focus on uncertainty modeling in fine-tuned neural implicit functions addresses an underexplored area, contributing valuable insights into this challenging problem space.

- Dropsembles combines dropout and deep ensembles, offering a computationally efficient  training method for uncertainty estimation specifically tailored to neural implicit functions. This application of established uncertainty quantification and continual learning techniques to neural implicit functions is considered a novel contribution to neural implicit functions.

- The paper provides a rigorous experimental evaluation across synthetic and real-world spine datasets, effectively demonstrating that Dropsembles maintain prediction accuracy and uncertainty calibration under varying conditions.

- The method achieves significant reductions in computational costs during training compared to traditional deep ensembles, without compromising performance. However, computational demands during inference remain comparable to those of ensemble methods.

- The integration of Elastic Weight Consolidation (EWC) effectively mitigates the impact of distribution shifts, enhancing the model's robustness to noisy and sparse inputs.

- The approach is straightforward and versatile, capable of being integrated into various neural implicit networks and adapted to different task-specific training objectives, making it broadly applicable.

- The strategy of training on synthetic data while addressing distribution shifts between synthetic and real data is a pragmatic solution, particularly in scenarios where real-world data is scarce and noisy.

- The paper is clearly written, well-structured, and easy to follow.

weaknesses:
- While the method shows promise in specific tasks like lumbar spine reconstruction, broader application beyond medical imaging is not thoroughly explored or validated.
- Despite reduced costs compared to deep ensembles, Dropsembles still require significant resources during fine-tuning, especially on high-resolution data.
- Although Dropsembles handle sparse inputs well, the effectiveness on highly varied real-world scenarios with extreme data corruption could have been examined more comprehensively.
- While Dropsembles generally outperform baselines, the improvements in some scenarios, particularly when using EWC, are not consistently significant across all metrics.

questions:
- Could you provide more details on how the EWC regularization strength was selected? Is there a rationale beyond the empirical ablation studies, or could a more systematic method be developed?

- Have you tested Dropsembles on other types of neural implicit functions or datasets beyond those presented? How do you anticipate the method will perform in domains with significantly different characteristics?

- While the method reduces training costs, inference demands remain comparable to traditional deep ensembles. Are there potential strategies to also optimize inference efficiency?

- Given that synthetic data might not capture all real-world complexities, how do you ensure the generalization of the model to highly variable or unseen real-world data?

- Can you elaborate on the effectiveness of Dropsembles in handling extreme distribution shifts? Are there specific limitations or edge cases where the method struggles?

- How easily can Dropsembles be integrated into existing neural implicit networks? Are there any specific prerequisites or limitations that practitioners should be aware of?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730841062918

thread_type:
official_review

<thread_object>
creator: reviewer_6vsx

<content>
summary:
This paper introduces 'Dropsembles', an approach to uncertainty modeling in fine-tuned neural implicit functions. 'Dropsembles' combines dropout-based methods with ensemble learning to balance computational efficiency and robust uncertainty estimation.
The authors conduct experiments on synthetic anatomical datasets and MRI segmentation tasks, demonstrating that 'Dropsembles' maintains predictive accuracy comparable to deep ensembles while reducing computational demands.

strengths:
Key contributions include:

(1) A method for modeling epistemic uncertainty in fine-tuned implicit functions.  

(2) The application of Elastic Weight Consolidation (EWC) to address distribution shifts. 

(3) Experimental validation across toy and real-world datasets.

weaknesses:
(1) **Questionable motivation for EWC**:

**Lack of task distinction**: The dense and sparse datasets (Tasks A and B) are similar in content, which doesn’t align well with the typical use case for EWC in continual learning. EWC is usually applied to preserve knowledge across distinct tasks, so its application here appears unjustified or artificial.

(2) **Limited novelty in dropsembles**:

**Lack of specificity to NIR**:
Dropsembles combines dropout with ensembling in a straightforward way without tailoring it for neural implicit representations. 
This method may not provide significant innovation for uncertainty estimation in 3D reconstruction, particularly if it doesn’t leverage the unique characteristics of implicit functions.

(3) **Insufficient baseline comparisons**:

Dropsembles is not compared to other standard uncertainty estimation techniques such as Bayesian neural networks, or other methods tailored to implicit representations [1]. This limits the evidence for Dropsembles’ effectiveness.

(4) **Limited dataset diversity**:

**Narrow generalizability**: The experiments are focused on synthetic anatomical data and sparse MRI segmentation, which may not fully demonstrate Dropsembles' robustness across various domains or types of data shifts.

References:

[1] Stochastic Neural Radiance Fields: Quantifying Uncertainty in Implicit 3D Representations, 3DV 2021

questions:
**Q1**: Why is EWC necessary if Tasks A and B are similar?

Could you clarify why EWC is used in this context, given that Task A (dense dataset) and Task B (sparse/noisy dataset) do not appear to be significantly different? Is the intent to quantify uncertainty in a continual learning setting, or is EWC applied simply as a regularization tool?
What makes Dropsembles particularly suited to neural implicit representations?

**Q2**: Dropsembles seems like a standard dropout-ensemble approach. Are there specific adaptations that make it uniquely effective for implicit functions in 3D reconstructions, especially regarding spatial consistency or grid-based predictions?
Considering Dropsembles combines well-known methods, why weren’t additional baselines, like Bayesian neural networks or variational inference-based [1], tested to provide a clearer comparison?

**Q3**: How does Dropsembles handle spatial dependencies?

Neural implicit representations often require spatial coherence across grid or voxel predictions. 
Does Dropsembles include any mechanism to maintain this spatial consistency in its uncertainty estimation, like [2]?

Could you discuss the limitations of Dropsembles in more detail?

**Q4**: Given that Dropsembles simply ensembles thinned models from dropout, what are its limitations in terms of computational cost, reliability, or adaptability to different 3D tasks?

**Minor**: The term "thinned network" isn’t a standard term in the field and may lead to confusion. It could be misinterpreted, as it doesn't clearly indicate that dropout is being applied.

**References**:

[1] Stochastic Neural Radiance Fields: Quantifying Uncertainty in Implicit 3D Representations, 3DV 2021

[2] Stochastic Segmentation Networks: Modelling Spatially Correlated Aleatoric Uncertainty, NeurIPS 2021
</content>
</thread_object>

</thread>
<thread>
creation date:
1732564233227

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response to all reviewers

comment:
We thank the reviewers for their time and valuable feedback. We appreciate 
the thorough comments which after addressing have improved our manuscript. Here we will summarize our response addressing common/essential concerns and then follow with point by point responses. We additionally provide a revised version of the manuscript, where changes are highlighted in "red".

Here we summarize the main points raised by the reviewers:

# Generalizability and additional results

## 1. Generalizability beyond medical shapes and occupancy networks 

Upon reviewers' suggestion we provide additional experiments on **ShapeNet** dataset **with SDF** representations (details on this experiment could be found in the revised version on the manuscript):


| Method | mDICE [%] ↑ | mIoU [%] ↑ | ECE [%] ↓ |
|--------|-------------|-------------|--------|
| Dropout [trained only on A] | 78.4 ± 2.4 | 64.5 ± 3.2 | 31.66 ± 3.37 |
| Ensembles [trained only on A] | 76.4 ± 2.9 | 61.9 ± 3.7 | 36.33 ± 3.11 |
|--------|-------------|-------------|--------|
| Dropout | 82.8 ± 2.2 | 70.8 ± 3.3 | 19.31 ± 3.97 |
| Dropsembles | 83.5 ± 2.1 | 71.7 ± 3.2 | 19.51 ± 4.64 |
| Ensembles | 82.6 ± 2.2 | 70.5 ± 3.2 | 25.88 ± 3.27 |
| Dropout + EWC | 83.9 ± 2.4 | 72.3 ± 3.6 | **16.96** ± 3.81 |
| Dropsembles + EWC | **84.3** ± 2.1 | **72.9** ± 3.2 | 17.67 ± 4.45 |
| Ensembles + EWC | 83.4 ± 1.9 | 71.6 ± 2.8 | 24.59 ± 2.98 |


## 2.  Generalizability to other forms of INR

As pointed out by the reviewers, the INR community shifted to using activation functions beyond ReLU activations. Below the experimental results on ShapeNet dataset for **sine activation functions (SIREN)** [Sitzmann et al., 2020], demonstrating that Dropsembles can be efficiently applied to sine activation functions too.

| Method | mDICE [%] ↑ | mIoU [%] ↑ | ECE [%] ↓ |
|--------|-------------|-------------|--------|
| Dropout - SIREN [trained only on A] | 78.5 ± 3.3 | 64.8 ± 4.4 | 31.94 ± 3.89 |
| Ensembles - SIREN [trained only on A] | 75.9 ± 3.7 | 61.3 ± 4.8 | 36.44 ± 3.85 |
|--------|-------------|-------------|--------|
| Dropout | 80.2 ± 2.4 | 67.0 ± 3.3 | 27.99 ± 2.85 |
| Dropsembles | 83.8 ± 2.2 | 72.2 ± 3.2 | 21.04 ± 3.04 |
| Ensembles | 83.4 ± 2.5 | 71.6 ± 3.6 | 24.44 ± 3.03 |
| Dropout + EWC | 83.3 ± 2.2 | 71.5 ± 3.2 | 19.27 ± 2.74 |
| Dropsembles + EWC | **84.5** ± 2.3 | **73.2** ± 3.4 | **17.73** ± 3.05 |
| Ensembles + EWC | 80.9 ± 2.5 | 68.0 ± 3.5 | 28.95 ± 3.07 |



# Relevance of EWC
Addititonal resuls on ShapeNet from the tables above demonstrate consistent improvement across all methods with EWC regularization. In particular, in case of Dropsembles, EWC regularization brings additional +0.8% improvement over non-regularized version. 

# Specificity of Dropsembles to uncertainty modeling in INR

## Dropsembels are better than vanilla Dropout for modeling uncertainty in INR
We would like to highlight the effectiveness and particular benefit of using Dropsembles versus vanila Dropout for INR. We do agree that Dropout is used in INR modeling for the purpose of regularization [Park et al. 2019] and is a commonly used method for modeling uncertainty in computer vision tasks [Kendall et al. 2017]. However, the *application of Dropout for the purpose of uncertainty modeling in INR*, to our knowledge, has not been explored and could benefit from a closer inspection.
We illustrated this crucial point in Figure 2(c-d), where we demonstrated that Dropout produced "speckle-type artifacts" in both reconstructions and uncertainties. To additionally support this point, we add Figure 11 in Appendix E where we demonstrate exactly the same behavior on ShapeNet dataset modeled with DeepSDF approach. From these figures it is apparent that the Dropout impairs the ability of the model to capture the true surface of the shape, producing a shape speckled with wholes and overall coarse texture. These "speckle-type" artifacts in Dropout predictions and uncertainty are specific for INR training procedure: during training, dropout is applied randomly at each sample, but compared to normal segmentation/regression tasks in computer vision (e.g. segmentation or depth estimation) in INR a "sample of data" contains a random subsample of shape 3d-coordinates. This in turn leads that every random set of coordinates is optimized with a random network instance of dropout and dropout instances are not coherent over 3D coordinates anymore. Naturally, Dropout implemented in such a way, does not translate to trustworthy uncertainty estimates. But, complementing Dropout with ensembels as we do in Dropsembels, mitigates this issue (as presented in Figures 2 and 11), producing a smooth shape representation. Hence, with Dropsembels we introduce an uncertainty modeling method that is tailored to the complexity and requirements of 3D representations. Additionally, this is supported by quantitative results in Tables 1, 2 and Appendix Tables 4.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a method called Dropsembles for uncertainty estimation for implicit neural representations. The method combines dropout with model ensembles with a learned prior by fine-tuning a pre-trained model using elastic weight consolidation. The reviewers commented on both the simplicity of the approach and the fact that it constituted a general-purpose approach not necessarily tailored to implicit neural representations. However, they appreciated the importance of the application and (especially after revision) the breadth of evaluation in terms of datasets. There was some concern raised about the absence of comparison to Bayesian neural networks.

The paper received two borderline accepts and one strong accept, with these reviewers all suggesting that the paper will be of interest from the ICLR community.