C ALIBRATED A DVERSARIAL R EFINEMENT FOR S TOCHASTIC S EMANTIC S EGMENTATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Ambiguities in images or unsystematic annotation can lead to multiple valid solutions in semantic segmentation.
To learn a distribution over predictions, recent work has explored the use of probabilistic networks.
However, these do not necessarily capture the empirical distribution accurately.
In this work, we aim to learn a calibrated multimodal predictive distribution, where the empirical frequency of the sampled predictions closely reﬂects that of the corresponding labels in the training set.
To this end, we propose a novel two-stage, cascaded strategy for calibrated adversarial reﬁnement.
In the ﬁrst stage, we explicitly model the data with a categorical likelihood.
In the second, we train an adversarial network to sample from it an arbitrary number of coherent predictions.
The model can be used independently or integrated into any black-box segmentation framework to facilitate learning of calibrated stochastic mappings.
We demonstrate the utility and versatility of the approach by attaining state-of-the-art results on the multigrader LIDC dataset and a modiﬁed Cityscapes dataset.
In addition, we use a toy regression dataset to show that our framework is not conﬁned to semantic segmentation, and the core design can be adapted to other tasks requiring learning a calibrated predictive distribution
I NTRODUCTION
Even though ambiguities are often unavoidable in many real-world datasets, the majority of the research encompassing semantic segmentation focuses on optimising models that assign a single prediction to each input image (Ronneberger et al., 2015;
Jégou et al., 2017; Takikawa et al., 2019;
Chen et al., 2017a;b; 2016a;b; 2015).
These are often incapable of capturing the entire empirical distribution of outputs as they optimise for a one-ﬁts-all solution.
Consequently, conﬂicts between labels corresponding to a given input image can compromise the quality of the ﬁnal predictions.
To resolve such ambiguities, ideally, one would sample multiple consistent hypotheses, and use uncertainty information to identify potential errors in each one.
Such a system would be particularly useful for safety-critical applications, such as in medical diagnostics, or for guiding user interactions for reﬁnement in semi-automated settings.
Several approaches have been proposed to capture label multimodality in image-to-image translation tasks (Huang et al., 2018;
Lee et al., 2018; Zhu et al., 2017a; Bao et al., 2017; Zhang, 2018), with only a few of them applied on stochastic semantic segmentation (Kohl et al., 2018; 2019; Baumgartner et al., 2019;
Hu et al., 2019; Kamnitsas et al., 2017; Rupprecht et al., 2017).
These methods have the capacity to learn a diverse set of labels for each input, however, they are typically limited to a ﬁxed number of samples (Kamnitsas et al., 2017; Rupprecht et al., 2017), return uncalibrated predictions, or do not account for uncertainty.
First, restricting to a ﬁxed number of predictions can hinder a model from learning all of the modes of the ground truth distribution.
Second, an uncalibrated predictive distribution is one that does not accurately reﬂect the occurrence frequencies of individual modes in the training set (Guo et al., 2017; Kull et al., 2019).
Finally, uncertainty information is highly desirable, as it can be used for distinguishing between predictions that can be trusted, and those that should be reassessed (Kendall and Gal, 2017; Ovadia et al., 2019).
In this work, we tackle all three challenges by introducing a two-stage cascaded strategy for adversarial reﬁnement that facilitates the learning of a calibrated multimodal predictive distribution, and allows the extraction of reliable data uncertainty estimates.
The key contributions are as follows 1 :
• We propose a novel cascaded architecture that constructively combines explicit likelihood modelling with adversarial reﬁnement to sample an arbitrary number of conﬁdent, and self-consistent predictions given an input image.
• We introduce a novel loss term that facilitates learning of calibrated stochastic mappings when using adversarial neural networks.
To our knowledge this is the ﬁrst work to do so.
• The proposed model can be trained independently or used to augment any pretrained blackbox semantic segmentation model, endowing it with a multimodal predictive distribution.
R ELATED WORK
Straightforward strategies towards learning multiple predictions include ensembling (Lakshminarayanan et al., 2017; Kamnitsas et al., 2017) or using multiple prediction heads (Rupprecht et al., 2017).
Even though these approaches can capture a diverse set of sampled predictions, they are limited to only a ﬁxed number of samples.
Alternatively, a probability distribution over the outputs can be induced by activating dropout during test time (Gal and Ghahramani, 2016a).
This method does offer useful uncertainty estimates over the pixel-space, however, Isola et al.
(2017) and Kohl et al.
(2018) have demonstrated that it introduces only minor stochasticity in the output and returns incoherent samples.
Bhattacharyya et al.
(2018) identify the maximum likelihood learning objective as the cause for this phenomenon in dropout Bayesian neural networks (Gal and Ghahramani, 2016b).
They postulate that under cross entropy optimisation, all sampled models are forced to explain all the data, and thereby converge to the mean solution.
Accordingly, they propose to liberate the optimisation from a single solution constraint, making it conducive to multimodality, by replacing the cross entropy loss term with an adversarial loss term parametrising a synthetic likelihood (Rosca et al., 2017). In contrast to this method, our approach is simpler to implement as it is not cast in the framework of variational Bayes, which requires the speciﬁcation of weight priors and variational distribution family. Kohl et al.
(2018) take an orthogonal approach in combining a U-Net (Ronneberger et al., 2015) with a conditional variational autoencoder (cVAE) (Kingma and Welling, 2013) to learn a distribution over semantic labels.
In Kohl et al.
(2019) and Baumgartner et al.
the authors build on Kohl et al.
(2018) to improve the diversity of the samples by modelling the data on several scales of the image resolution.
Nonetheless, these methods do not explicitly calibrate the predictive distribution in the pixel-space, and consequently do not provide reliable aleatoric uncertainty estimates (Kendall and Gal, 2017; Choi et al., 2018; Gustafsson et al., 2019).
Hu et al.
(2019) address this shortcoming by using the intergrader variability as additional supervision.
A major limitation of this approach is the requirement of a-priori knowledge of all the modalities of the data distribution.
For many real-world datasets, however, this information is not readily available.
In the more general domain of image-to-image translation, alternative methods employ hybrid models that use adversarially trained cVAEs
(Zhu et al., 2017a; Bao et al., 2017) to learn a distribution over a latent code, capturing multimodality, in order to sample diverse and coherent predictions.
A common hurdle in conditional generative adversarial network (cGAN) approaches is that simply incorporating a noise vector as an additional input often results in mode collapse.
This occurs due to the lack of regularisation between the noise input and generator output, allowing the generator to learn to ignore the noise vector (Isola et al., 2017).
This issue is commonly resolved by using supplementary cycle-consistency losses (Huang et al., 2018;
Lee et al., 2018; Zhu et al., 2017a; Bao et al., 2017), as proposed by Zhu et al.
(2017b).
However, again, none of these methods address the challenge of calibrating the predictive distribution.
P RELIMINARIES
Given an input image x ∈ R H × W × C , semantic segmentation refers to the task of predicting a pixelwise class label y ∈ { 1 , . .
, K } H × W .
For a dataset of N image and label pairs, D = { x i , y i } Ni =1 , the conditional distribution p
D ( y | x ) can be explicitly modelled through a likelihood q θ ( y | x ) , parameterised by a convolutional neural network F with weights θ , and activated by a softmax function (Ronneberger et al., 2015; Jégou et al., 2017).
One simple, yet effective way to learn the class probabilities is to express y ∈ { 0 , 1 } H × W × K as a one-hot encoded label, and set q θ as a pixelwise factorised categorical distribution, given by:
[Equation]
The parameters θ are then optimised by minimising the cross entropy between p D and q θ , deﬁned as:
[Equation]
When trained with Eq.
(2), F θ learns an approximation of E p D [ y | x ] (Bishop, 2006) that captures the pixelwise class probabilities over the label corresponding to a given input.
This accommodates the quantiﬁcation of the amount of noise inherent to the data, also referred to as aleatoric uncertainty, which can be obtained by computing the entropy of the output of F θ , H ( F θ ( x ))
(Kendall and Gal, 2017).
Further, the ﬁnal segmentation map is typically obtained by applying the argmax function along the class dimension of the predicted probabilities.
However, due to its deterministic nature, this approach is unable to produce multiple alternative predictions for the same input image.
On the other hand, direct sampling from q θ yields incoherent semantic maps, as for noisy datasets, maximising the factorised likelihood results in unconﬁdent predictions in regions of inter-label inconsistencies, e. g. fuzzy object boundaries, exempliﬁed later in the model overview diagram in Fig.
The aforementioned limitations can be partially addressed by adapting the framework of generative adversarial networks (GANs) (Goodfellow et al., 2014) to the context of conditional semantic segmentation, as proposed by Luc et al.
Formally, this involves training a binary discriminator network D to optimally distinguish between ground truth and predictions, while concurrently training a generative network G to maximise the probability that prediction samples G ( x ) are perceived as real by D .
Importantly, in contrast to explicit pixelwise likelihood maximisation, the adversarial setup learns an implicit sampler through G , capable of modelling the joint pixel conﬁguration of the synthesised labels, and capturing both local and global consistencies present in the ground truth.
In practice, the generator loss is often complemented with the pixelwise loss from Eq.
(2) to improve training stability and prediction quality (Luc et al., 2016; Ghafoorian et al., 2018; Samson et al., 2019).
However, we argue that the two objective functions are not well aligned in the presence of noisy data.
While the categorical cross entropy optimises for a single solution for each input x , thus encouraging high entropy in q θ ( y | x ) within noisy regions of the data and calibrating the predictive distribution, the adversarial term allows multiple solutions while optimising for low entropy, label-like output.
Therefore combining these losses in an additive manner, and enforcing them on the same set of parameters can be suboptimal.
This issue can be mitigated to some extent by a scheduled downscaling of
L ce , however, the residual conﬂict between the two losses adversely affects the optimisation process.
M ETHOD 4. C ALIBRATED ADVERSARIAL REFINEMENT
In this work, we propose to avert potential conﬂict between the cross entropy and adversarial losses by decoupling them in a two-stage cascaded architecture consisting of a calibration network F θ optimised with L ce from Eq.
(2), the output of which is fed to a reﬁnement network G φ .
G φ is then optimised with an adversarial loss term, parameterised by an auxiliary discriminator D ψ trained with a binary cross entropy loss.
To account for the multimodality in the labels, we additionally condition the reﬁnement network on an extraneous noise variable (cid:15) ∼ N (0 , 1) .
In practice, we also condition the reﬁnement network and the discriminator on input x , however, we do not show this explicitly for notational convenience.
More formally, using the non-saturated version of the adversarial loss (Goodfellow et al., 2014), the objectives for the reﬁnement and discriminator networks respectively are given by:
[Equation]
To calibrate the predictive distribution, we impose diversity regularisation on G φ by introducing a novel loss term that encourages the sample average G φ ( F θ ( x )) :
= E p (
[Equation]
) [ G φ
[Equation]
F θ
[Equation]
x )
[Equation] )] to match the [Figure] class probabilities predicted by F θ ( x ) .
Here, G φ ( F θ ( x )) serves as an approximation of the implicit predictive distribution of the reﬁnement network.
To this end, we deﬁne an auxiliary fully-factorised categorical likelihood q φ as:
[Equation]
where φ is optimised by minimising the Kullback-Leibler divergence KL( q φ || q θ ) .
Since both q φ and q θ are categorical distributions, the divergence can be computed exactly.
We coin this loss term as the calibration loss , deﬁned as:
[Equation]
Since L cal optimises through G φ ( F θ ( x )) rather than a single sampled prediction, the model is not restricted to learning a single mode-averaging solution for each input x , and is therefore more compatible with L adv .
The total loss for the reﬁnement network then becomes:
[Equation]
where λ ≥ 0 is an adjustable hyperparameter.
Fig.
1 shows the interplay of F θ , G φ and D ψ and the corresponding loss terms.
Intuitively, the calibration network F θ serves three main purposes.
It accommodates the extraction of sample-free aleatoric uncertainty maps.
It provides G φ with an augmented representation of x enclosing probabilistic information about y .
Finally, it sets a calibration target used by L cal to regularise the predictive distribution of G φ , in a cycle-consistent manner (Zhu et al., 2017b).
The reﬁnement network can then be interpreted as a stochastic sampler, modelling the inter-pixel dependencies to draw self-consistent samples from the explicit likelihood provided by the calibration network.
Thus both the pixelwise class probability and object coherency are preserved.
This approach leads to improved mode coverage and training stability, and increased convergence speed, as demonstrated in Section 5.
4. P RACTICAL CONSIDERATIONS
An important consequence of the loss decomposition is that the weights of F θ can be kept ﬁxed, while the adversarial pair G φ and D ψ are being trained.
This allows F θ to be pretrained in isolation, consequently lowering the overall peak computational burden and improving training stability (see Algorithms 1 and 2 in Appendix A.1 for an outline of the training and inference procedures).
Further, computing L cal requires a Monte Carlo estimation of G φ ( F θ ( x )) , where the quality of the loss feedback increases with the sample count.
However, modern deep learning frameworks allow for the samples to be subsumed in the batch dimension, and can therefore be efﬁciently computed on
GPUs.
We also note that on noisy data points, the trained model does not require to be provided with all solutions for a given input.
Instead, during training a random input-label pair is sampled, and the model automatically learns to identify cases where the input is ambiguous.
Finally, our method can augment any existing black-box model B for semantic segmentation, furnishing it with a calibrated multimodal predictive distribution.
This can be done by conditioning F θ on the output of B , which we demonstrate in Section 5.2.2.
E XPERIMENTS 5. 1D BIMODAL REGRESSION
We give intuitive insight into the mechanics of the proposed calibration loss by designing and experimenting on a simple one-dimensional regression task.
To create the dataset, an input x ∈ [0 , 1] is mapped to y ∈ R as follows:
[Equation]
where b ∼ Bernoulli( π ) and (cid:15) ∼ N (0 , σ ) .
We generate 9 different scenarios by varying the degree of mode selection probability π ∈ { 0 . 5 , 0 . 6 , 0 . 9 } and the mode noise σ ∈ { 0 . 01 , 0 . 02 , 0 . 03 } .
For every data conﬁguration, we use a 4-layer MLP for each of F θ , G φ and D ψ , and train with and without calibration loss by setting the coefﬁcient λ in Eq.
(7) to 1 or 0, respectively.
All runs are trained with a learning rate of 1e − 4 , and each experiment is repeated ﬁve times.
Note that unlike the categorical likelihood used in semantic segmentation tasks, we employ a Gaussian likelihood with ﬁxed scale of 1.
This changes the formulation of both Eqs.
(2) and (6) to mean squared error losses
[Figure]
The predictions average in dark blue matches the calibration target in red.
The discriminator output is shown between ground truth labels y and predictions
ˆ y for L ce , and between the output of the calibration network F θ ( x ) and the average of multiple predictions G φ
( F θ ( x )) for L cal (see Appendix A.2).
The results, depicted in Fig. 2, show that when using calibration loss, the optimisation process shows improved stability, converges faster and results in better calibrated predictions, in comparison to the non-regularised baseline.
The effect is more pronounced in data conﬁgurations with higher bias.
Further plots of the individual experiments are presented in Appendix B.1.
5. S TOCHASTIC SEMANTIC SEGMENTATION
In this section we examine the capacity of our model to learn shape and class multimodality in real-world segmentation datasets.
We begin by sharing essential implementation details below.
Network architectures For the calibration network F θ , we use the encoder-decoder architecture from SegNet (Badrinarayanan et al., 2017).
For G φ , we designed a U-Net-style (Ronneberger et al., 2015) architecture with 4 down- and upsampling blocks, each consisting of a convolutional layer, followed by a batch normalisation layer (Ioffe and Szegedy, 2015), a leaky ReLU activation, and a dropout layer (Srivastava et al., 2014) with 0.1 dropout probability.
We use a base number ofchannels, doubled or halved at every down- and upsampling transition.
To propagate the sampled noise vector to the output, we inject it into every upsampling block of the network in an afﬁne manner.
To do so, we project the noise vector using two fully connected layers into scale and residual matrices, with the same number of channels as the feature maps at the points of injection, and use these matrices to adjust the channel-wise mean and variance of the activations.
This is similar to the mechanism used for adaptive instance normalisation (Huang and Belongie, 2017).
We base the architecture for
D ψ on that used in DC-GAN
(Radford et al., 2015) except that we remove batch normalisation.
Any deviations this setup are described sections.
Training details We utilise the Adam optimiser (Kingma and Ba, 2014) with an initial learning rate of 2e − 4 for F θ and G φ , and 1e − 5 for D ψ .
The learning rates are linearly decayed over time and we perform scheduled updates to train the networks.
Additionally, the discriminator loss is regularised by using the R 1 zero-centered gradient penalty term (Mescheder et al., 2018).
For a detailed list of hyperparameter values, see Appendix A.1.
Metrics Following Kohl et al.
(2018; 2019), Huang et al.
(2018) and Baumgartner et al.
(2019), we use the Generalised Energy Distance (GED) (Székely and Rizzo, 2013) metric:
[Equation]
where d ( s, y )
= 1 − IoU( s, y ) .
As an additional metric, we follow Kohl et al.
(2019) in using the
Hungarian-matched IoU (HM-IoU)
(Kuhn, 2005).
All experiments are performed in triplicate and we report results as mean and standard deviation.
Further details regarding the exact implementation of the GED and HM-IoU metrics for each experiment can be found in Appendix A.3.
5.2. L EARNING SHAPE DIVERSITY ON THE LIDC DATASET
The Lung Image Database Consortium (LIDC) (Armato III et al., 2011) dataset consists ofthoracic CT scans from 1010 lung cancer patients, graded independently by four expert annotators.
We use the 180 × 180 crops from the preprocessed version of the LIDC dataset used and described in Kohl et al.
The dataset is split in 8882, 1996 and 1992 images in the training, validation and test sets respectively.
All models are trained on lesion-centered 128 × 128 crops where at least one of the four annotations indicates a lesion.
The ﬁnal evaluation is performed on the provided test set.
To assess the accuracy and diversity of samples generated by our model we examine how well the learnt conditional distribution captures the shape diversity of lesion segmentations in the dataset.
We pretrain F θ ( x ) with L ce in isolation, ﬁx its weights, and then train G φ with L G , where we estimate
L cal with 20 samples from G φ .
Further experiments with varying sample size are disclosed in Tablein Appendix B.2.
As a control experiment, we train using the same architecture but replace L cal in the reﬁnement network loss function with a cross entropy loss L ce , as used in Luc et al.
The L cal -regularised model performs on par with other state-of-the-art methods w. r. t. the GED score, and outperforms them on the HM-IoU score (only available for Kohl et al. Numerical [Table] x y 1gt y 2gt y 3gt
y 4gt y gt
F θ ( x ) y ref y 1ref y 2ref y 3ref y 4ref
y 5ref y 6ref results are summarised in Table 1.
The diversity and ﬁdelity of sampled predictions are illustrated in Fig.
In contrast, the L ce -regularised baseline collapses the predictive distribution, showing no perceptible diversity in the samples (see Fig.
9b in Appendix B.2.3), which results in a stark increase in the mean GED score, and decrease in the HM-IoU score, as shown in the bottom section of Table 1.
5.2. L EARNING A CALIBRATED DISTRIBUTION ON A MULTIMODAL C ITYSCAPES DATASET
The Cityscapes dataset contains 1024 × 2048 RGB images of urban scenes, and their corresponding segmentation maps.
It consists of 2975 training, 500 validation and 1525 test images.
Following Kohl et al.
(2018), we use the version of the Cityscapes dataset with 19 semantic classes and downsampled images and segmentation maps at a spatial resolution of 256 × 512 .
They establish controlled multimodality by augmenting the dataset with 5 new classes: sidewalk2 , person2 , car2 , vegetation2 and road2 , introduced by ﬂipping their original counterparts with probabilities 8 / 17 , 7 / 17 , 6 / 17 , 5 / 17 and/ 17 , respectively (see Fig. 4a).
Following Kohl et al.
(2018), we report results on the validation set.
Since we have full knowledge over the probability of each mode in the ground truth distribution, we use this dataset to assess the calibration of the tested models.
Further, to demonstrate that our approach can be easily integrated on top of any existing black-box segmentation model B , we employ the network from TensorFlow DeepLab Model Zoo (2020), ﬁnetuned on the ofﬁcial Cityscapes dataset, which achieves a mIoU of 0.79 on the test set.
We utilise its predictions as input to our calibration network, F θ , which consists of 5 convolutional blocks, each composed of a 3 × 3 convolutional layer, followed by a batch normalisation layer, a leaky ReLU activation, and a dropout layer with 0.
dropout rate.
We pretrain the calibration network in isolation, and subsequently apply it in inference unlabeledroad2vegetation2car2person2sidewalk2bicyclemotorcycletrainbustruckcarriderpersonskyterrainvegetationtrafﬁcsigntrafﬁclightpolefencewallbuildingsidewalkroad
[Figure]
E n t r o p y
( b i t s ) (c) mode while adversarially training the reﬁnement network.
We use a batch size of 16, and train with
L G , estimating L cal with 7 samples from G φ .
The same baseline as in the LIDC experiment is employed, where we replace
L cal with L ce .
As a second control experiment we completely omit the calibration network and instead condition the reﬁnement network on the known ground truth pixelwise categorical distribution over the label.
This allows us to directly evaluate the quality of sampling administered from the reﬁnement network.
When training the reﬁnement network with a cross entropy loss instead of the calibration loss L cal , the predictive distribution collapses, making the output deterministic.
Conversely, when we train our reﬁnement network with L cal , the learnt predictive distribution is well adjusted, with high diversity and reconstruction quality, signiﬁcantly outperforming the current state-of-the-art.
Fig.
4b shows representative sampled predictions from our model for three input images, and Fig.
4c illustrates the corresponding aleatoric uncertainty maps extracted from F θ ( x ) .
The learnt multimodality and noise in the dataset is reﬂected by regions of high uncertainty, where objects belonging to the different stochastic classes consistently display distinct shades of red, corresponding to their respective ﬂip probabilities.
Finally, we show that when using the ground truth pixelwise distribution as the input to the reﬁnement network, we attain an almost perfect GED score.
Thus, we demonstrate that G φ learns calibrated reﬁnement of the predictions from F θ , where the quality of the ﬁnal predictive distribution depends on the quality of F θ ( x ) .
The GED scores are reported in Table 2 and the calibration on the stochastic classes, evaluated over the entire dataset, is shown in Fig.
An important outstanding issue in our approach is that the calibration network does not perfectly capture the class-probabilities, e. g. for the car2 category in Fig.
The limitations of modern neural networks w. r. t. calibration are well documented (Guo et al., 2017; Kull et al., 2019; Zhang et al., 2020).
This has been attributed to several factors, such as long-tailed data distributions, out-of- [Figure] i x e l w i s e p r o [Figure]
a b ili t y [Figure] Method GED
[Table]
distribution inputs, speciﬁc network architectural elements or optimising procedures.
Therefore, an important direction for future work is improving the overall calibration of deep neural networks.
C ONCLUSION
In this work we developed a novel framework for semantic segmentation capable of learning a calibrated multimodal predictive distribution, closely matching the ground truth distribution of labels.
We attained improved results on a modiﬁed Cityscapes dataset and competitive scores on the LIDC dataset indicating the utility of our approach on real-world datasets.
We also showed that our approach can be easily integrated into an off-the-shelf, deterministic, black-box semantic segmentation model, enabling sampling an arbitrary number of plausible segmentation maps.
By providing multiple valid label proposals and highlighting regions of high data uncertainty, our approach can be used to identify and resolve ambiguities, diminishing risk in safety-critical systems.
Therefore, we expect our approach to be particularly beneﬁcial for applications such as map making for autonomous driving or computer-assisted medical diagnostics.
Finally, even though the primary focus of this work is semantic segmentation we demonstrated its versatility through an illustrative toy regression problem, alluding to a broader applicability beyond semantic image segmentation
Appendices A I MPLEMENTATION DETAILS
In this section we describe the overall training procedure and delve into the training and evaluation details for the stochastic segmentation experiments on the LIDC dataset and the modiﬁed Cityscapes dataset.
A. T RAINING PROCEDURE
Algorithm 1 outlines the practical procedure used to pretrain the calibration network, and the subsequent training of the reﬁnement network.
Even though the two networks can be trained end-to-end at once, in our experiments we use the two-step training procedure to stabilise the training and reduce the memory consumption on the GPU.
This way we are able to ﬁt larger batches and/or more samples for the estimate of L cal .
Algorithm 2 shows the inference procedure for obtaining M output samples.
Algorithm require: training data D , number of samples M , learning rate η , calibration loss scale λ ; 1: procedure T RAINING ( D , M, η, λ ) 2: while not converged do (cid:46)
Pretraining of F θ 3:
Sample batch { x, y } t ∈ D 4:
Update θ t +1 with − η ∇ θ t L ce ( { x, y } t , θ t ) 5: end while 6: while not converged do (cid:46)
Adversarial training of G φ and D ψ 7:
Sample batch { x, y } t ∈ D 8: for i = 1 , 2 , . . .
, M do 9:
Sample y i,t ref =
G φ t ( F θ ∗ ( x t ) , (cid:15) i ) where (cid:15) i ∼ N (0 , 1) 10: end for 11:
Compute G φ ( F θ ( x )) t = 1 M (cid:80) Mi =
y i,t ref 12:
Compute L cal ( x t , θ ∗ , φ t ) =
(cid:80) i,j,k (cid:0) G φ
( F θ ( x )) t (cid:0) log G φ ( F θ ( x ))
t − log F θ ∗ ( x t ) (cid:1)(cid:1) i,j,k 13:
Update φ t +1 with − η ∇ φ t ( L G ( θ ∗ , φ t , { x, y } t ) + λ L cal ( x t , θ ∗ , φ t )) 14:
Update ψ t +1 with − η ∇ ψ t L D ( θ ∗ , φ t , ψ t , { x, y } t ) 15: end while 16: end procedure
Algorithm require: test data point x , number of samples M ; 1: procedure I NFERENCE ( x, M ) (cid:46) Using θ ∗ and φ ∗ from Algorithm2:
for i = 1 , 2 , . . .
, M do 3:
Sample y i ref =
G φ ∗ ( F θ ∗ ( x ) , (cid:15) i ) where (cid:15) i ∼ N (0 , 1) 4: end for 5: end procedure
Notice that any off-the-shelf optimisation algorithm can be used to update the parameters θ , φ and ψ .
For the segmentation experiments, we utilise the Adam optimiser (Kingma and Ba, 2014) with β 1 = 0 .
5 , β 2 = 0 .and
weight decay of 5e − 4 .
F θ is trained with a learning rate of 2e − 4 which is then lowered to 1e − 4 after 30 epochs.
G φ and D ψ are updated according to a schedule, where G φ is updated at every iteration, and D ψ is trained in cycles of 50 iterations of weight updating, followed by 200 iterations with ﬁxed weights.
The reﬁnement network is trained with an initial learning rate of 2e − 4 , lowered to 1e − 4 after 30 epochs, whereas the discriminator has an initial learning rate of 1e − 5 , lowered to 5e − 6 after 30 epochs.
Additionally, we utilise the R 1 zero-centered gradient penalty term (Mescheder et al., 2018), to regularise the discriminator gradient on real data with a weight of 10.
Other hyperparameter speciﬁcs such as the batch-size and whether we inject stochasticity via random noise samples or latent code samples, depend on the experiment and are disclosed in the
A. 1D BIMODAL REGRESSION
In the following we derive the mean squared error form of the cross entropy and calibration losses used in experiment Section 5.1 under the assumption that the likelihood model for q θ and q φ is a univariate Gaussian distribution with a ﬁxed unit variance.
Using the setup from Eq.
(2) it then follows that:
[Equation]
Based on the deﬁnition of the calibration loss in Eq
[Equation]
(6) we show that:
[Equation]
A. LIDC
Architectures For the calibration network, F θ , we use the encoder-decoder architecture from SegNet (Badri- narayanan et al., 2017), with a softmax activation on the output layer.
Training During training, we draw random 180 × 180 image-annotation pairs, and we apply random horizontal ﬂips and crop the data to produce 128 × 128 lesion-centered image tiles.
All of our models were implemented in PyTorch and trained for 80k iterations on a single 32GB Tesla V100 GPU.
We train all our models for the LIDC experiments using 8-dimensional noise vectors in the cGAN experiments, or latent codes in the cVAE-GAN experiments.
This value was empirically found to perform well, sufﬁciently capturing the shape diversity in the dataset.
Additionally, in the reﬁnement networks loss, we set the weighting parameter λ in the total generator loss, deﬁned in Eq.
(7) in the main text, so as to establish a ratio of
L G : L cal = 1 : 0 . 5 , where L G is the adversarial component of the loss, and L cal is the calibration loss component.
In practice, the actual weights used are 10 for L G , and 5 for L cal .
Evaluation Following Kohl et al.
, Kohl et al.
(2019), Huang et al.
(2018), and Baumgartner et al.
we use the Generalised Energy Distance (GED) (Székely and Rizzo, 2013) metric, given as:
[Equation]
where d ( s, y )
= 1 − IoU( s, y ) .
Intuitively, the ﬁrst term of Eq.
(18) quantiﬁes the disparity between sampled predictions and the ground truth labels, the second term—the diversity between the predictions, and the third term—the diversity between the ground truth labels.
It is important to note that the GED is a sample-based metric, and therefore the quality of the score scales with the number of samples.
We approximate the expectations with all 4 ground truth labels ( y ∼ p D ) and 16, 50 or 100 samples from the model ( s ∼ q φ ) for each input image x .
As Kohl et al.
(2019) have pointed out, even though the GED metric is a good indicator for how well the learnt predictive distribution ﬁts a multimodal ground truth distribution, it can reward high diversity in sampled predictions even if the individual samples do not show high ﬁdelity to the ground truth distribution.
As an alternative metric that is less sensitive to such degenerate cases, Kohl et al.
(2019) propose to use the Hungarianmatched IoU (HM-IoU), which ﬁnds the optimal 1:1 IoU matching between ground truth samples and sampled predictions.
Following Kohl et al.
(2019), we duplicate the set of ground truth labels so that the number of both ground truth and predicted samples are 16, and we report the HM-IoU as the average of the best matched pairs for each input image.
In the main text we show the evaluated performance with both GED and HM-IoU metrics over the entire test set and compute the IoU on only the foreground of the sampled labels and predictions.
In the case where both the matched up label and prediction do not show a lesion, the IoU is set to 1, so that a correct prediction of the absence of a lesion is rewarded.
Note that the methods of Baumgartner et al.
(2019) and Hu et al.
(2019), whose GED scores we report in Table 1, use test sets that differ from the original splits deﬁned in Kohl et al.
(2018), which are used in Kohl et al.
(2018; 2019) and our work.
Baumgartner et al.
(2019) uses a random 60:20:20 split for the training, testing and validation sets, and 100 samples to compute the GED score, whereas Hu et al.
(2019) use a random 70:15:15 split and 50 samples to compute the GED score.
Due to the lack of reproducibility, we do not consider this the conventional way of benchmarking.
Therefore, in Table 1 we only report the scores for our models evaluated on the original splits.
Nevertheless, we also trained and tested our cGAN+ L cal model on the split methodology deﬁned by Baumgartner et al.
(2019) to enable a fairer comparison.
This improved our GED score from 0 .
243 ± 0 .
to 0 .
228 ± 0 .
009 , while the HM-IoU, evaluated using 16 samples, remained similar at 0 .
590 ± 0 .
007 (in the original splits we achieved an HM-IoU score of 0 .
592 ± 0 .
This shows that different random splits can signiﬁcantly affect the ﬁnal performance w.r.t.
GED score, while HM-IoU appears to be a more robust metric.
A. C ITYSCAPES
Architectures For the calibration network F θ , we design a small neural network with 5 convolutional blocks, each comprised of a 3 × 3 convolutional layer, followed by a batchnorm layer and a leaky ReLU activation.
The network is activated with a softmax function.
Training During training, we apply random horizontal ﬂips, scaling and crops of size 128 × 128 on the image-label pairs.
All of our models were implemented in PyTorch and trained for 120k training iterations on a single 16GB Tesla V100 GPU.
We train all our models for the modiﬁed Cityscapes experiments using 32-dimensional noise vectors.
Similarly to the LIDC experiments, this value was empirically found to perform well, however, it can be further tuned.
As commonly practiced, we use the ignore-masks provided by the Cityscapes dataset to ﬁlter out the cross entropy, calibration and adversarial losses during training on the unlabelled pixels.
Similarly to our LIDC experiment, we use a weight of 10 for L G , and 5 for L cal in the reﬁnement networks loss.
Evaluation The GED metric for Cityscapes is implemented as described in the appendix of Kohl et al.
and evaluated across the entire validation set.
In this dataset we have full knowledge of the ground truth class distribution and therefore we compute the GED metric by using the probabilities of each mode directly, as follows:
[Equation]
where w ( · ) is a function mapping the mode of a given label y to its corresponding probability mass.
The distance d ( s, y ) is computed using the average IoU of the 10 switchable classes only, as done in Kohl et al.
In the cases where none of the switchable classes are present in both the ground truth label and the prediction paired up in d ( s, y ) , the distance score is not considered in the expectation.
We use 16 samples to compute the GED score.
For the calibration results presented in Fig.
5, Section 5.2.2 in the main text, we compute the calibration network class-probabilities using the raw predictions of F θ ( x ) .
We obtain class masks by computing the overlap between the ground truth labels and the black-box predictions for each class
B A DDITIONAL EXPERIMENT RESULTS
To reinforce the results reported in Section 5 we present supplementary results for the bimodal regression experiment and the LIDC and Cityscapes segmentation experiments.
B. 1D BIMODAL REGRESSION
Fig.
6 shows the data log-likelihoods for the 9 data conﬁgurations for varying mode bias π ∈ { 0 . 5 , 0 . 6 , 0 . 9 } and mode noise σ ∈ { 0 . 01 , 0 .
03 } trained with and without the calibration loss L cal .
Each experiment is repeated 5 times and the individual likelihood curves are plotted in Fig.
6b and Fig.
6d respectively.
The results show that high bias is harder to learn, reﬂected by a slowed down convergence, however, the L cal -regularised model shows greater robustness to weight initialisation.
In contrast the non-regularised GAN exhibits mode oscillation expressed as a ﬂuctuation of higher likelihood (one mode is covered) and lower one (between modes).
[Figure]
B. LIDC B.2. Q UALITATIVE A NALYSIS
To further examine our L cal -regularised cGAN model trained on the LIDC dataset, we illustrate representative qualitative results in Fig.
7 and Fig.
For every input image x , we show the ground truth labels y 1 gt , . . .
, y 4 gt provided by the different expert annotators, overlaying the input image, in the ﬁrst four columns, and 6 randomly sampled predictions yref , . . .
, y 6 ref in the last six columns.
From left to right, the three columns with the dark blue background in the center of the ﬁgures show the average ground truth predictions ¯ y gt , the output of the calibration network F θ ( x ) and the average of 16 sampled predictions from the reﬁnement network ¯ y ref .
Our results show that even though there is a signiﬁcant variability between the reﬁnement network samples for a given input image, ¯ y ref is almost identical to the calibration target F θ ( x ) , due to the diversity regularisation enforced by the calibration loss L cal .
[Figure] [Figure]
( x , y 2gt )
( x , y 3gt ) ( x , y 4gt )
y gt
F θ ( x ) y ref y 1ref y 2ref y 3ref
y 4ref y 5ref
y 6ref
From the qualitative results in Fig.
7 and Fig.
8, it can be seen that the calibration target F θ ( x ) does not always capture well the average of the ground truth distribution ˆ y gt , affecting the ﬁdelity of the predictive distribution of the reﬁnement network G φ .
This further highlights the importance of future work on improving the calibration of F θ , e.g. implementing the approaches of Guo et al.
Kull et al.
B.2. T UNING THE NUMBER OF REFINEMENT NETWORK SAMPLES
[Table]
To investigate the effect of the number of samples used to compute L cal on the learnt predictive distribution
[Table]
we experimented on the L cal -regularised cGAN model using 5, 10, 15 or 20 samples from the reﬁnement network G φ during training.
As a control experiment
[Table]
we also train the same model using one sample.
Our results, reported in Table 3, show that increasing the number of samples improves the quality of the predictive distribution, whereas using only one sample collapses it.
This is expected because increasing the number of samples reduces the variance of the sample mean G φ and reﬁnes the approximation q φ of the implicit predictive distribution realised by G φ ( x, (cid:15) ) .
Since in our implementation we reuse the samples from G φ in the adversarial component L
G of the total reﬁnement network loss L G , the discriminator D ψ interacts with a larger set of diverse fake samples during each training iteration, thus also improving the quality of L G .
It is important to note that the beneﬁt of increasing the sample size on the quality of L cal highly depends on the intrinsic multimodality in the data.
In theory, if the number of samples used matches or exceeds the number of ground truth modes for a given input, it is sufﬁcient to induce a calibrated predictive distribution.
However, we usually do not have a priori access to this information.
Conversely, if the sample size is too small, the L cal loss may introduce bias in the predictive distribution.
This could lead to mode coupling or mode collapse, as exempliﬁed in our control experiment with one sample.
In the LIDC dataset, even though we have access to four labels per input image, we argue that the dataset exhibits distributed multimodality, where a given pattern in the input space, e.g. in a patch of pixels, can be associated to many different local labels throughout the dataset.
As a result, an input image may correspond to more solutions than the four annotations provided.
Therefore increasing the number of samples to more than four shows further improvement in performance.
This however can come at the cost of decreased training speed which can be regulated by tuning the sample count parameter while considering the system requirements.
B.2. I NDUCING MULTIMODALITY IN LATENT VARIABLE MODELS ON THE LIDC DATASET
To examine whether conditioning the source of stochasticity in our model on the input is beneﬁcial, we adapt our framework in order to learn a distribution over a latent code z using variational inference (Graves, 2011).
Following most of the existing work on stochastic semantic segmentation (Kohl et al., 2018; 2019; Hu et al., 2019; Baumgartner et al., 2019), we maximise a modiﬁed lower bound on the conditional marginal log-likelihood, through a variational distribution q ( z | x, y ) .
This is realised by minimising the loss function in Higgins et al.
(2017), given by:
[Equation]
where β controls the amount of regularisation from a prior distribution p ( z ) on the approximate posterior q ( z | x, y ) .
Both q ( z | x, y ) and p ( z ) are commonly taken as factorised Gaussian distributions.
To this end, we compare our cGAN model to two baselines where the reﬁnement network G φ is given as a cVAE-GAN (Larsen et al., 2015).
In the ﬁrst one, we train G φ by complementing the adversarial loss L G with Eq.
(20), using β ∈ { 0 . 1 , 1 , 10 } and a ﬁxed standard normal prior.
In the second, we introduce a calibration network and train G φ using Eq.
This does not necessitate specifying a prior.
For a fair comparison, we use the same core models for all of our experiments, introducing only minor modiﬁcations to the reﬁnement network to convert the deterministic encoder into a probabilistic one with
Gaussian output.
This is achieved by splitting the output head of the encoder so as to predict the mean and standard deviation of the encoded distribution (Kingma and Welling, 2013; Kendall and Gal, 2017).
Instead of using random noise sampled from a standard Gaussian as our source of stochasticity, the decoder of the reﬁnement network is now injected with latent codes sampled from the Gaussian distribution encoded for each input image.
To train the model, we pretrain the F θ in isolation, and subsequently apply it in inference mode while training G φ .
We use a batch size of 32, an 8-dimensional latent code, and use 20 samples to compute L cal .
[Table]
We show that the cVAE-GAN model trained with L cal instead of the traditional complexity loss term, KL( q ( z | x, y )
||
p ( z )) from Eq.
(20) is able to learn a distribution over segmentation maps, and performs similarly to our cGAN+ L cal model.
This is important because it abrogates the need for specifying a latent-space prior, which is often selected for computational convenience, rather than task relevance (Hafner et al., 2018;
Louizos et al., 2019).
On the other hand, our cVAE-GAN models trained using the KL-divergence complexity term showed limited diversity, even for large β values.
The results, shown quantitatively in the bottom part of Table 4, and qualitatively in Fig.
9a, demonstrate that the interaction between L G and L cal can sufﬁciently induce a multimodal predictive distribution in latent variable models, and indicate that for the purpose of stochastic semantic segmentation the use of a probabilistic encoder is not strictly required.
x y 1gt y 2gt y 3gt
y 4gt y gt
F θ ( x ) y ref y 1ref y 2ref y 3ref
y 4ref y 5ref
y 6ref (a) 0.
x y 1gt y 2gt y 3gt
y 4gt y gt
F θ ( x ) y ref y 1ref y 2ref y 3ref y 4ref
y 5ref
y 6ref (b) unlabeledroad2vegetation2car2person2sidewalk2bicyclemotorcycletrainbustruckcarriderpersonskyterrainvegetationtrafﬁcsigntrafﬁclightpolefencewallbuildingsidewalkroad 0.
E n t r o p y
( b i t s ) unlabeledroad2vegetation2car2person2sidewalk2bicyclemotorcycletrainbustruckcarriderpersonskyterrainvegetationtrafﬁcsigntrafﬁclightpolefencewallbuildingsidewalkroad 0.
E n t r o p y
( b i t s )
B.3. Q UALITATIVE A NALYSIS
In this section we provide additional qualitative results for the L cal -regularised cGAN model trained on the modiﬁed Cityscapes dataset (Kohl et al., 2018).
In Fig.
10, we show 16 randomly sampled predictions for representative input images x , and their corresponding aleatoric uncertainty maps, obtained by computing the entropy of the output of the calibration network, H ( F θ ( x )) , as done in Kendall and Gal (2017).
The predicted samples are of high quality, evident by object coherence and crisp outlines, and high diversity, where all classes are well represented.
Our model effectively learns the entropy of the ground truth distribution in the stochastic classes ( sidewalk , person , car , vegetation and road ), as their distinct entropy levels are captured as different shades of red in the entropy maps, corresponding to the different ﬂip probabilities ( 8 / 17 , 7 / 17 , 6 / 17 , 5 / 17 and 4 /respectively).
Additionally, it can be seen that edges or object boundaries are also highlighted in the aleatoric uncertainty maps, which reﬂects inconsistency during manual annotation, which often occurs on input pixels that are difﬁcult to segment.
Fig.
11d shows the entropy of the predictive distribution of the reﬁnement network G φ , H (cid:0) G φ ( F θ ( x )) (cid:1) , where
G φ ( F θ ( x )) is computed as the average of 16 samples from G φ .
Our results demonstrate that H (cid:0) G φ ( F θ ( x )) (cid:1) is similar to H ( F θ ( x )) , depicted in Fig. 11c, as encouraged by the L cal regularisation.
Notice that object boundaries are also highlighted in H (cid:0) G φ ( F θ ( x )) (cid:1) , indicating that our model captures shape ambiguity as well as class ambiguity.
However, some uncertainty information from H ( F θ ( x )) is not present in H (cid:0) G φ ( F θ ( x )) (cid:1) , e.g. the entropy of the different stochastic classes are not always consistent across images, as evident from the different shades of red seen for the road class in Fig. 11d.
We expect that increasing the number of samples from the reﬁnement network will improve the aleatoric uncertainty estimates.
Nevertheless, the sample-free estimate extracted from F θ ( x ) is cheaper to obtain and more reliable than the sample-based average from G φ , highlighting an important beneﬁt of our cascaded approach.
Finally, we illustrate in Fig.
11e the high conﬁdence of the predictions from the reﬁnement network G φ ( F θ ( x )) , reﬂected by their low entropy, H ( G φ ( F θ ( x ))) .
This is attributed to the adversarial component in the reﬁnement loss function, which encourages the predictions to assume a one-hot representation, matching the ground truth annotations.
Even though each prediction of the reﬁnement network is highly conﬁdent, the average of the predictions G φ ( F θ ( x )) is calibrated, as shown in Fig. 11d.
This is a clear illustration of the advantage of complementing the adversarial loss term L G with the calibration loss term L cal in the training objective for the reﬁnement network.
[Figure]
[Figure]
[Figure]
E n t r o p y (
[Figure]
i t [Figure] [Figure] [Figure]
[Figure]
[Figure]
E n t r o p y
( b i t s )
[Figure]
[Figure] [Figure] E n t r o p y ( b i t s )
[Figure]