Source-linked AI summary

We Know Where We Don't Know: 3D Bayesian CNNs for Credible Geometric Uncertainty

Tyler LaBonte, Carianne Martinez, Scott A. Roberts

arXiv:1910.10793v2eess.IVcs.CVcs.LG

TL;DR

Reliable segmentation for high-consequence CT applications requires uncertainty estimates, but existing approaches had not produced statistically credible geometric uncertainty maps that scale to 3D. The paper proposes a 3D Bayesian CNN and reports that it outperforms an MCDN on uncertainty metrics while generating interpretable geometric maps. The authors note that available labels cannot properly calibrate uncertainty, limiting validation.

  • Problem

    High-consequence CT segmentation needs uncertainty quantification, but existing methods lacked statistically credible geometric uncertainty maps and faced concerns about Bayesian scalability.

    Source: §1 · p. 2
  • Method

    The authors propose a 3D Bayesian CNN that measures weight-space uncertainty to generate geometric uncertainty maps interpretable as segmentation confidence intervals.

    Source: §1 · pp. 2–3
  • Results

    The BCNN outperforms MCDNs on recent uncertainty metrics and achieves equal or better segmentation accuracy in most cases.

    Source: §1 · p. 3, §6 · p. 15
  • Takeaways & Limitations

    The BCNN produces continuous, visually graded geometric uncertainty maps that can be interpreted as confidence intervals on segmentations.

    Source: §1 · p. 3, §5.1 · p. 10

    CT artifacts and noise make human labels inconsistent, preventing proper uncertainty calibration and leaving PAvPU insufficient without ground-truth labels.

    Source: §5.2 · p. 12

Abstract

from arXiv · show

Deep learning has been successfully applied to the segmentation of 3D Computed Tomography (CT) scans. Establishing the credibility of these segmentations requires uncertainty quantification (UQ) to identify untrustworthy predictions. Recent UQ architectures include Monte Carlo dropout networks (MCDNs), which approximate deep Gaussian processes, and Bayesian neural networks (BNNs), which learn the distribution of the weight space. BNNs are advantageous over MCDNs for UQ but are thought to be computationally infeasible in high dimension, and neither architecture has produced interpretable geometric uncertainty maps. We propose a novel 3D Bayesian convolutional neural network (BCNN), the first deep learning method which generates statistically credible geometric uncertainty maps and scales for application to 3D data. We present experimental results on CT scans of graphite electrodes and laser-welded metals and show that our BCNN outperforms an MCDN in recent uncertainty metrics. The geometric uncertainty maps generated by our BCNN capture distributions of sigmoid values that are interpretable as confidence intervals, critical for applications that rely on deep learning for high-consequence decisions. Code available at https://github.com/sandialabs/bcnn.

1 Introduction

The paper presents a 3D Bayesian CNN for credible geometric uncertainty in CT segmentation, addressing the need for trustworthy UQ in high-consequence applications. It reports that the BCNN outperforms an MCDN on recent uncertainty metrics and produces interpretable maps.

1 Introduction

  • High-consequence part certification requires segmentation UQ for safety confidence intervals, while CT noise and human error make segmentations unreliable.CNNs provide fast, accurate automated segmentation, but require UQ to identify untrustworthy predictions.Source: §1 · p. 2
  • The BCNN outperforms an MCDN on recent uncertainty metrics for CT scans of graphite electrodes and laser-welded metals.The paper also reports that the BCNN captures geometric uncertainty maps that the MCDN cannot.Source: §1 · p. 3
  • The BCNN measures weight-space uncertainty, producing geometric maps interpretable as confidence intervals on segmentations.Its maps capture continuity and visual gradients, unlike the pixelated, uninterpretable MCDN maps.Source: Figure 1 · p. 2, §1 · p. 3
  • The authors present the first deep-learning method for statistically credible geometric uncertainty maps and the first Bayesian CNN that scales to 3D data.They report successful segmentation of billion-voxel CT scans and describe their approach as a synthesis of techniques rather than a straightforward extension from 2D.Source: §1 · pp. 2–3

2 Related Work

Prior work established volumetric segmentation and probabilistic UQ approaches, but left challenges around statistical validity, scale, and geometric uncertainty. The paper situates its 3D BCNN against those approaches and their limitations.

2.1 Volumetric Segmentation

  • Volumetric segmentation progressed from 2D networks to encoder-decoder 3D CNNs for medical imaging, while materials work addressed CT scans of fibers and composites.V-Net and 3D U-Net map volumes through a latent representation to voxel-wise segmentations.Source: §2.1 · pp. 3–4

2.2 Uncertainty Quantification

  • Neural-network predictive probabilities are not reliable uncertainty measures, motivating Bayesian models and variational inference over network weights.The section distinguishes probabilistic modeling from treating softmax outputs as uncertainty.Source: §2.2 · p. 4

Monte Carlo Dropout Networks (MCDNs)

  • MCDNs use stochastic dropout passes for uncertainty estimates and are easy to add to large networks [6], but their statistical soundness has been questioned.Prior work argues dropout may approximate model risk rather than posterior uncertainty, with its posterior depending on dropout rate and model size.Source: Monte Carlo Dropout Networks (MCDNs) · p. 4

Bayesian Neural Networks (BNNs)

  • Variational BNNs approximate the intractable weight posterior by optimizing an objective balancing a simplicity prior and dataset fit.Bayes by Backprop uses backpropagation gradients to update posterior parameters; local reparametrization and Flipout reduce variance from sampled weights. Earlier 2D BCNN work did not produce geometric maps, while complex datasets remained difficult.Source: Bayesian Neural Networks (BNNs) · p. 5

3 Methodology

The methodology combines an encoder-decoder segmentation network with Bayesian convolutional layers in the decoder. Design choices address information transfer, normalization, and reliable convergence.

3.1 Architecture

  • A four-stage encoder compresses the input, while a three-stage decoder upsamples, concatenates skip features, and produces a same-size binary segmentation probability map.The encoder omits stochastic layers to maximize information transfer into the latent space.Source: §3.1 · p. 6

3 Methodology

  • The BCNN places probabilistic weights in volumetric convolutional layers of the decoder and uses Flipout to approximate their distributions.The layers use a standard normal prior; the BCNN has 1,924,964 trainable parameters, compared with 1,403,059 for its MCDN counterpart.Source: §3.1 · pp. 6–7

3.2 Design Decisions

  • Group normalization with four groups gave the most reliable convergence under memory-constrained batch sizes; the filter count balanced segmentation accuracy against training difficulty.The authors found a standard normal prior best balanced initialization and exploration, skip connections slightly improved accuracy, and max pooling versus downward convolutions made negligible difference.Source: §3.2 · p. 7

4 Experiments

The experiments evaluate the BCNN on hand-segmented CT volumes of graphite electrodes and laser-welded joints, with dataset-specific training and held-out testing. The procedures cover minibatch loss scaling, KL annealing, and uncertainty comparison with an MCDN.

4.1 Datasets

  • The Graphite dataset contains eight hand-segmented electrode volumes; the BCNN was trained on GCA400 and tested on the other seven.The volumes range from 500 million to 1 billion voxels each.Source: §4.1 · p. 7
  • The Laser Weld dataset comprises ten hand-segmented joint volumes; training used S2, S24, and S25, with the other seven held out for testing.Each volume contains tens of millions of voxels, and light regions represent weld voids or defects.Source: §4.1 · p. 8
  • Both datasets were normalized to zero mean and unit variance, and scans were processed as subvolumes because they were too large to ingest whole on the GPU.Source: §4.1 · p. 8

4.2 Training

  • Dividing the KL penalty by the number of minibatches distributes it across batches; monotonic KL annealing lets the model learn 3D segmentation before applying that penalty.Without the scaling, the KL term dominates and can lead to a posterior with suboptimal accuracy; annealing was necessary for reliable convergence.Source: §4.2 · p. 8
  • Training used Bayes by Backprop, with dataset-specific KL-annealing hyperparameters, Adam learning rates, and epoch counts.Each epoch processed 1,331 chunks of size 88 × 176 × 176 in about 17 minutes 30 seconds on two NVIDIA Tesla V100 GPUs; training ran for 2 Graphite epochs and 7 Laser Weld epochs.Source: §4.2 · p. 8

4.3 Testing

  • For testing, voxel confidence intervals are the 33rd and 67th percentiles of sigmoid values, and uncertainty is their difference; the MCDN uses the same architecture with spatial dropout.The BCNN does not guarantee a normal distribution of sigmoid values, so these percentiles are used to compare outputs while mimicking the MCDN's standard-deviation measurement.Source: §4.3 · p. 9

5 Results

Across CT test sets, the BCNN produces interpretable geometric uncertainty at large scale and generally matches the MCDN’s segmentation accuracy. It outperforms the MCDN on conditional-probability validation measures, while its simulation-oriented maps capture uncertainty distributions that the MCDN maps do not.

5.1 Prediction Results

  • On the billion-voxel Graphite sample, the BCNN scales successfully and provides interpretable geometric uncertainty with equivalent or better segmentation than the MCDN.Across the test sets, segmentation accuracy is similar overall; the authors qualify this comparison because the hand-segmented labels contain inaccuracies.Source: Figure 3 · p. 9, §5.1 · pp. 10–11
  • On the Laser Weld samples, BCNN uncertainty captures variation around the central void and stays consistent across weld borders, including where segmentation underestimates void size.The S4 uncertainty band reflects the void’s size despite the smaller predicted segmentation.Source: Figure 4 · p. 10, §5.1 · p. 10, Figure 5 · p. 11

5.2 Validation

  • PAvPU3D extends PAvPU to sliding cubic patches; its results use 2 × 2 × 2 patches, an accuracy threshold of 7/8, and an uncertainty threshold set to the map mean.The authors regard the metric as insufficient for validating UQ without ground-truth labels because scan boundaries are often unclear and hand labels cannot properly calibrate uncertainty.Source: §5.2 · p. 12
  • The BCNN consistently outperforms the MCDN in conditional probabilities, but underperforms on PAvPU3D because that metric penalizes predictions that are both accurate and uncertain.The authors note that accuracy and uncertainty are not mutually exclusive in the Bayesian view.Source: §5.2 · pp. 11–13, Table 2 · p. 13

5.3 Advantages for Material Simulations

  • BCNN maps represent differences between sigmoid-percentile intervals, capturing geometric changes across intervals that barely change in MCDN maps.The authors connect these interpretable intervals to probing uncertainty in simulations and bounding physical properties.Source: §5.3 · pp. 13–14, Figure 6 · p. 14

6 Conclusion

The BCNN introduces credible geometric uncertainty maps for 3D segmentation and robustly scales to three-dimensional data. Across graphite and laser-welded-metal CT scans, it matches or exceeds MCDN segmentation accuracy in most cases, outperforms MCDNs on recent uncertainty metrics, and yields more credible confidence intervals when probed.

6 Conclusion

  • The BCNN achieves equal or better segmentation accuracy than MCDNs in most cases and outperforms them on recent uncertainty metrics.Probing BCNN uncertainty also leads to more credible confidence intervals on the segmentation than the MCDN.Source: §6 · p. 15
  • The BCNN is the first deep learning method to generate statistically credible geometric uncertainty maps and the first Bayesian CNN to robustly scale with 3D data.It measures uncertainty in weight space to provide interpretable uncertainty quantification for segmentation.Source: §6 · p. 15
  • The reported UQ results use graphite-electrode and laser-welded-metal CT scans in material simulations where accurate geometric UQ is critical.Source: §6 · p. 15
  • Future work includes semantic segmentation, separating aleatoric and epistemic uncertainty, and integration with techniques such as deep ensembles.Source: §6 · p. 15

Appendix A Chunking and Prediction

The appendix describes chunking large CT volumes, predicting on chunks, and stitching results into a prediction and uncertainty volume. It also outlines a boundary-trimming procedure for reducing chunking artifacts and provides Python implementations.

Appendix A Chunking and Prediction

  • The chunking algorithm divides large CT volumes into user-sized chunks and returns their coordinates and the original volume shape.Users specify overlap through the step size; the authors recommend a step size of 3.Source: Appendix A · pp. 18–19
  • For each chunk, stochastic predictions are averaged and combined by voxel counts, while the 33rd-to-67th-percentile range supplies an uncertainty volume.The mean predictions are thresholded at 0.5 to produce the binary prediction volume.Source: Appendix A · pp. 18, 20
  • To address chunk-boundary artifacts, the procedure discards predictions near chunk borders; the described implementation uses a 10% threshold.The counts volume accounts for voxels discarded from a chunk.Source: Appendix A · pp. 18, 21
  • Python implementations of the algorithms are available in the authors’ codebase.Source: Appendix A · p. 18
Loading 1910.10793v2…