Semantic uncertainty intervals for disentangled latent spaces
Abstract
Meaningful uncertainty quantification in computer vision requires reasoning about semantic information—say, the hair color of the person in a photo or the location of a car on the street.
To this end, recent breakthroughs in generative modeling allow us to represent semantic information in disentangled latent spaces, but providing uncertainties on the semantic latent variables has remained challenging.
In this work, we provide principled uncertainty intervals that are guaranteed to contain the true semantic factors for any underlying generative model.
The method does the following: (1) it uses quantile regression to output a heuristic uncertainty interval for each element in the latent space (2) calibrates these uncertainties such that they contain the true value of the latent for a new, unseen input.
The endpoints of these calibrated intervals can then be propagated through the generator to produce interpretable uncertainty visualizations for each semantic factor.
This technique reliably communicates semantically meaningful, principled, and instance-adaptive uncertainty in inverse problems like image super-resolution and image completion.
Introduction
When making decisions with visual data, such as automated vehicle navigation with blurry images uncertainty quantification is critical.
The relevant uncertainty pertains to a low-dimensional set of semantic properties, such as the locations of objects.
However, there is a wide class of imagevalued estimation problems—from super-resolution to inpainting—for which there does not currently exist a method of producing semantically meaningful uncertainties.
As an example, imagine doing uncertainty quantification for medical image reconstruction from, say, a fast but undersampled MRI scan.
In such a setting, pixelwise intervals [16, 40, 5] are not very useful.
We need uncertainty on the underlying semantics—e.g., whether there is a tumor, and if so, of what size and shape.
We make progress on this problem by bringing techniques from quantile regression and distributionfree uncertainty quantification together with a disentangled latent space learned by a generative adversarial network (GAN).
We call the coordinates of this latent space semantic factors , as each controls one meaningful aspect of the image, like age or hair color.
We do not require these semantic factors to be statistically independent.
Our method takes a corrupted image input and predicts each semantic factor along with an uncertainty interval that is guaranteed to contain the true semantic factor.
When the model is unsure, the intervals are large, and vice-versa.
By propagating these intervals through the GAN coordinate-wise, we can visualize uncertainty directly in image-space without resorting to per-pixel intervals—see Figure 1.
The result of our procedure is a rich form of uncertainty quantification directly on the estimates of semantic properties of the image.
[Figure] More concretely, we receive input images X and then predict an uncertainty interval for each of D semantic factors Z d , d = 1 , ..., D , which are the elements of a disentangled latent space.
The method involves training an encoder (a neural network that takes images as input and produces outputs in the latent space of the GAN) on ( X, Z ) pairs to give us three different outputs:
The point prediction, f ( X ) .
This is the encoder’s best guess at the semantic factors Z . 2.
The estimated lower conditional quantile, q α 2 ( X ) .
The encoder believes that q α 2 ( X ) is a lower bound on the value of Z given X . 3.
The estimated upper conditional quantile, q 1 − α 2 ( X ) .
The encoder believes that q 1 − α 2 ( X ) is an upper bound on the value of Z given X .
Once the above encoder is trained, as described in Section 2.2, we use it to form an uncertainty interval for each semantic factor.
However, for the d th element of the latent code, the naive interval ( q α 2 ( X ) d , q 1 − α 2 ( X ) d ) is not guaranteed to contain the ground truth value in finite samples.
We propose to perform a calibration procedure to fix this problem, yielding the sets
[Equation]
(1) where q cal is a calibrated version of q constructed using the tools in Section 2.3.
Once we have done so, the intervals will contain α fraction of the true latent codes with high probability.
In other words, for any user-chosen levels α and δ , we can output intervals that with probability 1 − δ satisfy
[Equation]
for a new test point ( X, Z ) , regardless of the distribution of the data, the encoder used, and the number of data points used in the calibration procedure.
This guarantee, described more carefully in Definition 2.1, says that the intervals cover 1 − α fraction of the semantic factors unless our calibration data is not representative of our test data (which only happens with a probability δ which goes to 0 as the number of calibration data points grows).
We visualize each of the d ∈ { 1 , ..., D } intervals in latent space by propagating the d th lower and upper endpoints through the generator with all other entries in the latent fixed to the point prediction (see Section 2.3 for a formal explanation).
1.1 Central Contribution
To our knowledge, this is the first algorithm for uncertainty intervals on a learned semantic representation with formal statistical guarantees.
By propagating these intervals through the generator, we are able to visualize uncertainty in an intuitive new way that directly encodes semantic meaning.
This is an important step towards interpretable uncertainties in general image-valued estimation problems.
However, the reader should note that our technique requires access to a disentangled latent space, such as that of a StyleGAN, and provides no guarantees about the degree of disentanglement.
Method 2.1 Notation and goal
Our data consist of pairs ( X, Z ) —the corrupted image X in X =
[0 , 1] H × W , and the latent code
Z ∈ Z , where Z = R D .
As mentioned in the introduction, we think of Z as a disentangled representation with d semantic factors — i.e. factors of variation corresponding to interpretable features in an image, such as hair color and expression.
For simplicity, assume each of the d dimensions controls a single semantic factor; in practice, we ignore those that do not.
In our sampling model, X is generated from Z by composing two functions.
The first function is a fixed generator G : Z → Y , where Y =
[0 , 1] H × W , which takes the latent vector Z and produces the ground truth image Y ∈ Y (for ease of notation, we assume X and Y have the same shape).
The second function is a corruption model, F : Y → X , which degrades the ground truth image
Y to produce the corrupted image X , for example by randomly masking out part of the image.
To summarize our data-generating process, we have
[Equation]
Goal #1.
Our first goal is to train an encoder E to recover Z from X —in other words, to invert the mapping F ◦ G —with a heuristic notion of uncertainty.
The encoder’s point prediction will be a function f : X → Z .
The uncertainty will be parameterized by two functions, q 1 − α 2 : X → Z and q α 2 : X → Z , denoting our estimates of the 1 − α 2 and α 2 conditional quantiles, respectively.
These conditional quantiles are potentially bad estimates; they do not natively possess the statistical guarantee we desire.
Goal #2.
Having trained the encoder and the conditional quantile estimates, we will output uncertainty intervals in the disentangled latent space.
Each dimension will get its own interval, which has the form in (2).
Ultimately, our uncertainty intervals will come with a statistical risk control guarantee, like in (2), that is distribution-free—i.e., valid irrespective of the model or data distribution.
The risk function controlled is the false negative rate in the latent space .
In other words, we bound the fraction of semantic factors not covered by their calibrated intervals above by α (say, α = 10% ).
Definition 2.1 (Risk-Controlling Prediction Set (RCPS)) .
A set-valued function T : X → 2 Z is an ( α, δ ) risk-controlling prediction set if
[Equation]
where
[Equation]
The reader should note here that the function T depends on the calibration data.
The outer probability in (4) is over the randomness in this calibration procedure; the inner expectation is over the new test point, ( X, Y ) .
The reader should note that in our setting, because we can generate infinite data from the GAN, it is always possible to take δ arbitrarily close to 0 , effectively making T nonrandom.
In the two following subsections, we address each of our goals separately.
2.2 Goal #1: Training the encoder for quantile regression
Our job in this subsection is to learn the three functions f , q α 2 , and q 1 − α 2 .
We do so by training a neural network with three different loss functions, one for each of three linear heads on top of the same feature extractor—see Figure 2 for the training protocol.
Loss function for point prediction.
We supervise the point prediction with three loss functions.
The first is an L 1 loss directly in the latent space, and encourages f ( X ) to be close to Z .
[Equation]
The second loss term imposes a domain specific prior on the generated image G ( f ( X )) . For our experiments on faces, we use an identity loss which encourages G ( f ( X )) to contain attributes that have the same “identity” — or semantic content — as Y . We calculate the identity loss using a X
Y E G corrupt generator sample encoder image space , , latent space losses predict gradients pretrained embedding function, ID :
Y → R d ′ for some d ′ , which projects Y to an embedding space where images with different identities land far away from one another.
Finally, we calculate our loss using cosine similarity,
[Equation]
The third is a perceptual loss on the generated image G ( f ( X )) that imposes similarity between
G ( f ( X )) and X in the embedding space of a large pretrained image classification model.
We impose perceptual similarity by computing the LPIPS loss [50] which has been shown to preserve image quality [17]. The perceptual loss is calculated using a pretrained feature extractor in a similar manner to Eq 7. Following standard practice, the feature extractor used in this case is a VGG network, pretrained on Imagenet. Finally, we combine the loss functions to form the loss function for the prediction f ,
[Equation]
(8) where c 1 and c 2 are hyper-parameters whose values are chosen based on a held out set.
Loss function for quantile regression.
Quantile regression [30, 13, 32, 27, 28, 29, 31] is a statistical method for estimating the conditional quantiles of a distribution.
The key idea of quantile regression is to supervise the regressor using a quantile loss ,
[Equation]
The minimizer of the quantile risk is the true β conditional quantile of Z | X .
We supervise our conditional quantile estimates q α 2 and q 1 − α 2 with two separate instances of the quantile loss, L α 2 q and
L 1 − α 2 q respectively.
This concludes our explanation of the model training procedure, summarized in Algorithm 1.
Experimental details, such as the particular model architecture we use, are available in Section 3.
Algorithm
GAN Input: training dataset D , risk level α , number of epochs E , fixed generator G
Output: trained functions f, q α 2 , and q 1 − α 2 .
f, q α 2 , and q 1 − α 2 ← random model initialization for e = 1 to E do loss ←for ( X train , Z train ) in D do L1 ← L pred (cid:0) f ( X train (cid:1) , Z train ) L2 ← L α 2 q (cid:0) q
α 2 ( X train (cid:1) , Z train )
L3 ← L 1 − α 2 q (cid:0) q 1 − α 2 ( X train (cid:1) , Z train ) loss += L1 + L2 + L end for loss.backward() end for
2.3 Goal #2: Calibration
Having trained the model, we now calibrate it to achieve the statistical guarantee in Definition 2.
using a set of calibration data (cid:8) ( X i , Z i ) (cid:9) n i =1 generated from the model and the upper-confidence bound procedure from [8].
The output of the procedure will be the function T from (2); specifically, we will learn the calibrated conditional quantiles q cal α 2 and q cal1 − α 2 .
Our procedure will calibrate the conditional quantiles by rescaling their size multiplicatively.
We will ultimately choose a multiplicative factor ˆ λ that gives us the desired guarantee.
Towards that end, we index a family of uncertainty intervals scaled by a free parameter λ for each semantic factor,
[Equation]
When λ grows, the interval T λ ( X )
d also grows, and thus, the loss function L ( T λ ( X ) , Y ) shrinks.
Therefore, by taking λ large enough, we can always ensure the loss is zero.
The challenge ahead is to pick ˆ λ to be the smallest value such that T ˆ λ is an RCPS as in (4).
The algorithm for selecting ˆ λ involves forming an upper confidence bound (UCB) for the risk, then picking the smallest value of λ such that the upper confidence bound falls below α .
We give
Hoeffding’s UCB below, although we use the stronger Hoeffding-Bentkus bound from [8] in practice:
[Equation]
Note that in our setting, we can always generate enough samples to drive the last term to 0 for any δ ; however, if we only had a finite sample from some population, this would not be the case.
We can then select ˆ λ by scanning from large to small values, ˆ λ = min (cid:110) λ : ˆ R + ( λ ′ ) ≤ α, ∀ α ′ ≥
α (cid:111) , or running binary search.
Proposition 2.2 ( T ˆ λ is an RCPS [8]) .
With ˆ λ selected as above, T ˆ λ satisfies Definition 2.1.
For the proof of this fact, along with a discussion the tighter confidence bounds used in our experiments and extensions to the underlying theory, see [8] and [3].
Having proven that T ˆ λ is an RCPS, we can simply set T ( X ) =
T ˆ λ
( X ) in (2); in other words, we set
[Equation]
Visualizing uncertainty intervals in image space
We briefly describe our method for visualizing latent-space uncertainty intervals.
In order to see the effect of a single semantic factor, we set it to either the lower or upper quantile and hold the other factors fixed to the point prediction.
More specifically, for a particular dimension d ∈ { 1 , ..., D } , define the following vector:
[Equation]
We visualize the lower and upper quantiles in image space as G ( ˆ Z d α 2 ) and G ( ˆ Z d 1 − α 2 ) respectively.
Since each semantic factor corresponds to an attribute, visualizing the lower and upper quantiles per-factor gives interpretable meaning to the latent-space intervals.
For example, in Figure 1, the images of the child smiling give a range of possible expressions the model believes are consistent with the underlying image.
Experiments 3.1 Dataset descriptions
FFHQ We use the StyleGAN framework pretrained using the Flickr-Faces-HQ (FFHQ) dataset [25].
FFHQ is a publicly available dataset consisting of 70,000 high-quality images at 1024 ×resolution.
The data used to train the quantile encoder and for the experiments in this section is sampled from the generator pretrained on FFHQ.
CelebA-HQ We use the CelebA-HQ dataset [23] to demonstrate the effectiveness of our approach on real world data.
The dataset contains 30,000 high-quality images of celebrity faces.
Since the real dataset is used only for evaluation, we use images from standard test split.
CLEVR.
In order to have a controlled setup where we can easily identify disentangled factors of variation, we generate synthetic images of objects based on the CLEVR dataset [22].
This dataset provides a programmatic way of generating synthetic data by explicitly varying specific semantic factors.
We create a synthetic dataset by varying { color, shape } and fixing the other factors such as lighting, material and camera jitter.
3.2 Experimental setup
Model architectures In all our experiments, we use the StyleGAN2 [24] framework for the generator architecture G .
For the experiments involving faces, we use the pretrained model available from the official repository.
For the CLEVR-2D experiments, we train a simpler variant of StyleGAN generator from scratch.
For the quantile regression, we use a standard architecture: the encoder network consists of a ResNet-50 backbone [19] with the final layer branching into the point prediction and conditional quantiles.
The branching module is a standard combination of convolution and activation blocks followed by a fully connected layer of the expected output dimension; we call these the heads of the model.
Specific details of the model architecture are provided in the supplementary material.
We start by pretraining the generative model or acquiring an off-the-shelf pretrained generative model for the task at hand.
In generative models such as StyleGAN, the style space that offers fine grained control over image attributes, is very high dimensional.
From this high dimensional space, we extract only the disentangled dimensions following previous work on style space analysis [42].
In order to better focus the encoder’s capacity only on the disentangled dimensions, we mask out the irrelevant dimensions for applying the quantile loss.
However, the pointwise loss in (6) is applied to the full style vector to ensure that the pointwise prediction is able to match the true latents accurately, while the quantile heads focus on learning variablity only in the disentangled dimensions.
During the encoder training 2.2, the generative model G is held frozen and only the parameters of the encoder E are updated.
The point prediction and conditional quantile heads are trained jointly with the Ranger optimizer [42] and a flat learning rate of 0.001 for all our experiments.
{+The hyperparameter weights (Eq 8) are set to c 
+}
= c 2 = 10 . 0 .
For the image super-resolution training, we augment the input dataset by using different levels of downsampled inputs, i.e. , we take the raw input and apply a random downsampling factor from { 1 , 4 , 8 , 16 , 32 } and resize it to the original dimensions.
For the image inpainting task, we vary the difficulty by choosing a random threshold to create the mask – lower thresholds implies fewer pixels are masked and vice-versa.
The mask is concatenated to the image resulting in a C + 1 -channel input to the encoder, where C is the number of image channels.
The detailed description of the mask generation procedure can be found in the supplementary material.
Calibration and Evaluation For both the synthetic object experiments with CLEVR and face experiments with FFHQ, we train the quantile encoder on data points sampled from the latent space of the pretrained generative model.
This ensures that we have access to the true latents that resulted in each image.
We generate 100k samples per model and generate a random 80-10-10 split for training, calibration and validation.
3.3 Findings
In the following experiments, we explore different properties of our intervals.
The problem types include image super-resolution and image inpainting.
The risk level α and the user-specified error threshold δ are fixed to 0.1, unless specified otherwise.
3.3.1 Producing semantic uncertainties
Goal.
We qualitatively verify that the proposed approach outlined in Section 2 produces visually meaningful uncertainties.
Description.
We train a quantile encoder with Algorithm 1.
Then we generate n = 5000 images by sampling latents and propagating them through the encoder-generator combination.
We use these n images as calibration data for the procedure in 2.3.
Finally, we randomly sample a new test point, pass it through the calibrated encoder, and form the uncertainty intervals in image space as in Section 2.3.
[Figure]
The corrupt image is provided as input to the encoder which outputs a pointwise prediction and quantile predictions plot the and as well as their visualizations in image-space.
[ viewed in color in detail.]
Results.
The results are illustrated in Figure 3 on images sampled from the generator trained on
FFHQ (left) and CLEVR (right).
In case of the FFHQ image, the person is wearing glasses in the lower quantile image and not in the upper; hence, the model is not certain that the person is wearing glasses.
The model also expresses some uncertainty about the amount of gray versus brown hair.
This outcome was predictable from the input image, where the fact that the person is wearing glasses is not obvious, and there is some hair color ambiguity.
The results on the CLEVR dataset are analogous.
The lower and upper quantile images yield similar colors, which is predictable from the blurry input.
The model predicts that both a cylinder and sphere would be consistent with this blurry input.
The calibrated quantiles cover the ground truth color value, while the uncalibrated ones do not.
3.3.2 Experiments with real data
[Figure] factor of Though prediction Observe that in some cases, such as smile and glasses , there are clear visible changes due to the calibration approach.
In other cases such as hair , the changes are not visually explicit.
More such examples can be found in the supplementary material
Goal. We calibrate our encoder with real data sampled from the Celeb-A face dataset for the image resolution task.
The objective is to test to what extent we can extract meaningful uncertainty intervals for real data without explicitly training with them.
Description The difficulty with calibrating using real data is that the targets i.e the true style vectors that generated the real images are not available to us.
To circumvent that difficulty, we back-project the real images into the latent space of the GAN using existing approaches to GAN inversion [24, 1].
Even though the mapping from the GAN latent space to the output is not one-one, the back-projection optimization procedure is able to visually match the appearance of the real images convincingly.
This procedure provides us with the estimated "true" style vectors for each image in the real calibration set.
We combine this with the generated data and perform the RCPS calibration procedure as described in
Section 2.3.
Note that, our encoder was only trained on generated data i.e. the same model as used in
Section 3.3.1, we use the real data only for the calibration procedure.
Qualitative results The predicted quantiles and the effect of the calibration procedure for a real test image drawn from the Celeb-A dataset is visualized in Figure 4.
The pointwise prediction of the encoder does not match the identity of the true image due to the high level of input corruption.
However, using our approach, we are able to visualize the predicted uncertainty intervals in a meaningful manner.
The visual effect of our calibration procedure is demonstrated on the right.
Notice that some attributes such as smile and glasses have a visible change but others do not.
The change is due to the fact that the calibration procedure adjusts the intervals in such a way to satisfy the required coverage guarantee.
3.3.3 Exploratory results with purposeful corruptions
[Figure]
Goal.
We probe our procedure to see if it will have the expected qualitative behavior.
Description of experiment.
We sampled images from the held out set of the FFHQ pretrained GAN and applied purposeful corruptions to check if the resulting quantile estimates had semantic meaning.
Both image resolution and image masking were used as corruption models.
We qualitatively analyze the results by visualizing the predictions.
We include a quantitative measurement of variability estimated using image based metrics in the supplementary material.
Qualitative results.
Figure 5 shows the results of this experiment for Image inpainting (left) and super-resolution (right).
In the inpainting case, when nothing is masked, the quantiles are roughly identical.
When the eyes are masked, the quantiles indicate the model does not know if the person was wearing glasses.
When the mouth is masked, the model expresses uncertainty as to whether the woman is showing her teeth in the smile.
Finally, when almost everything is masked, the quantile images are very different, representing individuals with entirely different identities.
Similar behavior can be observed in the super-resolution case.
The results are shown for two separate inputs; in both cases, the input in the top row is uncorrupted while in the bottom row, it is undersampled by 16x.
The model is able to predict almost perfectly in the absence of corruption.
The quantile predictions are extremely close as well.
In the presence of corruption, both the pointwise prediction is off (as expected) and the quantile edges display much higher variability including in attributes such as hair shape, glasses, facial hair and perceived gender.
The results from both these experiments point to an expected qualitative behavior of our proposed approach: the model exhibits more uncertainty with increased information loss at the input.
Related Work
GANs for Inverse problems.
The remarkable image generation properties of recent approaches such as BigGAN [11] and StyleGAN [24] has led to the increasing use of these models to solve inverse problems relating to image restoration such as image super-resolution and completion.
All prior methods that use GANs for inverse problems such as the ones that use the pretrained generative model as an image prior [46, 9, 38, 36] or the ones that train an encoder model to project the input into the generator’s latent space [51, 42, 45, 41] focus on the accuracy of the point estimate but not on the uncertainty level of the input.
GANs for Interpretability.
Despite providing no guarantees of image likelihood, unlike others such as Normalized Flow [26] or Score-based models [52], GANs have been used to develop interpretable approaches to image generation.
The widespread use of GANs as opposed to other generative models in the interpretability is done due to the availability of an disentangled latent space [18, 49], which is a property we utilize in our work.
Quantile Regression.
Quantile regression was first proposed in [30].
Since then, many papers have used the technique, applying it to machine learning [21, 37, 39], medical research [7], and more.
Most relevant to us is conformalized quantile regression [43], which gives quantile regression a marginal coverage guarantee using conformal prediction.
Our work instead uses risk-controlling prediction sets, a different distribution-free uncertainty quantification technique.
Conformal prediction and distribution-free uncertainty quantification.
At the core of our proposed method is the distribution-free, marginal risk-control technique studied in [8] and [3].
These ideas have their roots in the distribution-free marginal guarantees of conformal prediction, proposed in [48].
Conformal prediction is a flexible method for forming prediction sets that satisfy a marginal coverage guarantee, under no assumptions besides the exchangeability of the test point with the calibration data [48, 47, 34, 35, 33, 2].
Conformal prediction has been studied in computer vision [20, 12, 6, 44, 4], natural language [14], drug discovery [15], criminal justice [10], and more.
We are not aware of work applying the notions of conformal prediction and quantile regression to the latent space in generative models.
Conclusion
Experiments indicate that with an appropriately disentangled model, latent space uncertainty intervals express a useful semantic notion of uncertainty previously unavailable in computer vision.
Limitations of our method include that the calibration data must be reflective of the data distribution, and that we assume access to a disentangled latent space.
We see our work as part of a larger tapestry of results in generative models, and our technique will remain applicable as progress gets made on disentanglement and backprojection.
Furthermore, by introducing a notion of statistical rigor to GAN based approaches , our work creates a path for application of such models to safety critical settings.
Ethics
The ethics of generative modeling itself has been called into question given recent events, e g. , the development of deep fakes.
Nonetheless, we believe the downstream consequences of this work will likely be positive.
The techniques herein do not change the predictions of a generative model; they simply provide a calibrated notion of its uncertainty in a relevant semantic space.
Thus, the standard criticism of generative modeling—that it will enable widespread deep fakes—is not applicable.
Furthermore, we expect having a statistically valid and semantically rich notion of uncertainty will provide a sobering reliability assessment of these models, perhaps mitigating the chance of harmful failures.
Finally, although we use face datasets due to their ubiquity in this literature, we have attempted to ethically treat topics like gender and race where they arise.
Acknowledgements
S. S acknowledges the MIT SuperCloud and Lincoln Laboratory Supercomputing Center for providing compute resources that have contributed to the results reported in this work.
A. N. A. was supported by NSF GRFP.
S. B. was supported by the Foundations of Data Science Institute and the Simons
Institute.
Y. R. was supported by the Israel Science Foundation (grant No. 729/21).
Y. R. thanks the
Career Advancement Fellowship, Technion, for providing research support.