Semantic uncertainty intervals for disentangled latent spaces
Anonymous Author(s)
AffiliationAddress email
Abstract
Meaningful uncertainty quantification in computer vision requires reasoning aboutsemantic information—say, the hair color of the person in a photo or the location ofa car on the street.
To this end, recent breakthroughs in generative modeling allowus to represent semantic information in disentangled latent spaces, but providinguncertainties on the semantic latent variables has remained challenging.
In thiswork, we provide principled uncertainty intervals that are guaranteed to contain thetrue semantic factors for any underlying generative model.
The method does thefollowing: (1) it uses quantile regression to output a heuristic uncertainty intervalfor each element in the latent space (2) calibrates these uncertainties such thatthey contain the true value of the latent for a new, unseen input.
The endpoints ofthese calibrated intervals can then be propagated through the generator to produceinterpretable uncertainty visualizations for each semantic factor.
This techniquereliably communicates semantically meaningful, principled, and instance-adaptiveuncertainty 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 ofsemantic properties, such as the locations of objects.
However, there is a wide class of image-valuedestimation problems—from super-resolution to inpainting—for which there does not currently exist amethod of producing semantically meaningful uncertainties.
Many methods exist for getting per-pixelintervals [15, 37, 4] but directly giving uncertainty on semantically meaningful image properties hasremained challenging.
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 generativeadversarial network (GAN).
We call the coordinates of this latent space semantic factors , as eachcontrols one meaningful aspect of the image, like age or hair color.
Our method takes a corruptedimage input and predicts each semantic factor along with an uncertainty interval that is guaranteed tocontain the true semantic factor.
When the model is unsure, the intervals are large, and vice-versa.
Bypropagating these intervals through the GAN coordinate-wise, we can visualize uncertainty directlyin image-space without resorting to per-pixel intervals—see Figure 1.
The result of our procedure is arich form of uncertainty quantification directly on the estimates of semantic properties of the image.
More concretely, we receive input images X and then predict an uncertainty interval for each of Dsemantic factors Z d , d = 1 , ..., D , which are the elements of a disentangled latent space.
The methodinvolves training an encoder (a neural network that takes images as input and produces outputs in thelatent space of the GAN) on ( X, Z ) pairs to give us three different outputs:
[Figure]
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 alower 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 uncertaintyinterval 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.
Wepropose 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 doneso, 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 thenumber of data points used in the calibration procedure.
This guarantee, described more carefullyin Definition 2.1, says that the intervals cover 1 − α fraction of the semantic factors unless ourcalibration data is not representative of our test data (which only happens with a probability δ whichgoes 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 withall 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 areable to visualize uncertainty in an intuitive new way that directly encodes semantic meaning.
This isan important step towards interpretable uncertainties in general image-valued estimation problems.
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 disentangledrepresentation with d semantic factors — i.e. factors of variation corresponding to interpretable featuresin an image, such as hair color and expression.
For simplicity, assume each of the d dimensionscontrols a single semantic factor; in practice, we ignore those that do not. X Y E
G corrupt generator sample encoder image space , , latent space losses predict gradients
In our sampling model, X is generated from Z by composing two functions.
The first function is afixed generator G : Z → Y , where
Y =
[0 , 1] H × W , which takes the latent vector Z and producesthe 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.
Tosummarize 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 invertthe mapping F ◦ G —with a heuristic notion of uncertainty.
The encoder’s point prediction will bea function f : X → Z .
The uncertainty will be parameterized by two functions, q 1 − α 2 : X → Zand 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 statisticalguarantee we desire.
Goal #2.
Having trained the encoder and the conditional quantile estimates, we will output uncertaintyintervals in the disentangled latent space.
Each dimension will get its own interval, which has theform in (2).
Ultimately, our uncertainty intervals will come with the following statistical guarantee:
Definition 2.1 (Risk-Controlling Prediction Set (RCPS)) .
A set-valued function T : X → 2 Z is an( α, δ ) risk-controlling prediction set if ( α , δ )-Risk-Controlling Prediction Set if P [Equation]
(cid:2)
[Equation]
where
[Equation]
The reader should note here that the function T depends on the calibration data.
The outer probabilityin (4) is over the randomness in this calibration procedure; the inner expectation is over the new testpoint, ( X, Y ) .
The reader should note that in our setting, because we can generate infinite data from the GAN, itis always possible to drive δ arbitrarily close to 0 , effectively making T nonrandom.
In the twofollowing 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 aneural network with three different loss functions, one for each of three linear heads on top of thesame feature extractor—see Figure 2 for the training protocol.
Loss function for point prediction.
We supervise the point prediction with two loss functions.
Thefirst is an L 1 loss directly in the latent space, and encourages f ( X ) to be close to Z .
[Equation]
The second is an identity loss on the generated image G ( f ( X )) , which encourages G ( f ( X )) tocontain objects that have the same “identity” — or semantic meaning — as those in Y .
We calculatethe identity loss using a pretrained embedding function, ID :
Y → R d ′ for some d ′ , which projects Yto an embedding space where images with different identities land far away from one another.
Finally,we calculate our loss using cosine similarity,
[Equation]
Finally, we combine these two loss functions to form the loss function for the prediction f ,
[Equation]
where c = 0 .
7 was chosen to balance the two losses.
Loss function for quantile regression.
Quantile regression [27, 12, 29, 24, 25, 26, 28] is a statisticalmethod for estimating the conditional quantiles of a distribution.
The key idea of quantile regressionis 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 ourconditional 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.
2.3 Goal #2: Calibration
Having trained the model, we now calibrate it to achieve the statistical guarantee in Definition 2.1using a set of calibration data (cid:8)
( X i , Z i ) (cid:9) n i =1 generated from the model and the upper-confidencebound procedure from [7].
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 willultimately choose a multiplicative factor ˆ λ that gives us the desired guarantee.
Towards that end, weindex a family of uncertainty intervals scaled by a free parameter λ for each semantic factor,
[Equation]
[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 topick ˆ λ 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 [7] in practice:
[Equation]
Note that in our setting, we can always generate enough samples to drive δ → 0 ; however, if weonly had a finite sample from some population, this would not be the case.
We can then select ˆ λ byscanning from large to small values, ˆ λ = min (cid:110) λ : ˆ R + ( λ ′ ) ≤ α, ∀ α ′ ≥
α (cid:111) .
Proposition 2.2 ( T ˆ λ is an RCPS [7]) .
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 [7] and [2].
Having proven that T ˆ λ is an RCPS, we can simply set T ( X ) =
T ˆ λ
( X ) in (2); in other words, we setq cal α 2 ( X ) d
= f ( X ) d − ˆ λ (cid:0)
[Equation]
( X ) d − q α 2 ( X ) d (cid:1) + and q cal1 − α 2 ( X ) d =
[Equation]
( X
[Equation]
d + ˆ λ (cid:0) q 1 − α 2 ( X ) d − f
( X ) d (cid:1) + .
Visualizing uncertainty intervals in image space
We briefly describe our method for visualizing latent-space uncertainty intervals.
In order to see theeffect of a single semantic factor, we set it to either the lower or upper quantile and hold the otherfactors 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 quantilesper-factor gives interpretable meaning to the latent-space intervals.
For example, in Figure 1, theimages of the child smiling give a range of possible expressions the model believes are consistentwith the underlying image.
Experiments
3.1 Dataset descriptions
FFHQ We use the StyleGAN framework pretrained using the Flickr-Faces-HQ (FFHQ) dataset [22].
FFHQ is a publicly available dataset consisting of 70,000 high-quality images at 1024 × 1024resolution.
FFHQ has significant variation among age, ethnicity and image backgrounds.
The dataused train the quantile encoder and for the experiments this section is sampled from the generatorpretrained on FFHQ.
CLEVR.
In order to have a controlled setup where we can easily identify disentangled factors ofvariation, we generate synthetic images of objects based on the CLEVR dataset [20].
This datasetprovides a programmatic way of generating synthetic data by explicitly varying specific semanticfactors.
We create a synthetic dataset by varying { color, shape } and fixing the other factors such aslighting, material and camera jitter.
3.2 Experimental setup
Model architectures In all our experiments, we use the StyleGAN2 [21] framework for the generatorarchitecture G .
For the experiments involving faces, we use the pretrained model available fromthe official repository.
For the CLEVR-2D experiments, we train a simpler variant of StyleGAN2generator from scratch.
For the quantile regression, we use a standard architecture: the encodernetwork consists of a ResNet-50 backbone [17] with the final layer branching into the point predictionand conditional quantiles.
The branching module is a standard combination of convolution andactivation blocks followed by a fully connected layer of the expected output dimension; we call thesethe heads of the model.
Specific details of the model architecture are provided in the supplementarymaterial.
[Figure]
The corrupt image is provided as input to the encoder which outputs a pointwise prediction and quantile predictions for each style dimension.
We plot the calibrated and uncalibrated intervals as well as their visualizations in image-space.
[Best viewed in color, zoom in for detail.]
Model training We start by pretraining the generative model or acquiring an off-the-shelf pretrainedgenerative model for the task at hand.
In generative models such as StyleGAN, the style space thatoffers fine grained control over image attributes, is very high dimensional.
From this high dimensionalspace, we extract only the disentangled dimensions following previous work on style space analysis[39].
In order to better focus the encoder’s capacity only on the disentangled dimensions, we maskout the irrelevant dimensions for applying the quantile loss.
However, the pointwise loss in (6) isapplied to the full style vector to ensure that the pointwise prediction is able to match the true latentsaccurately, 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 ofthe encoder E are updated.
The point prediction and conditional quantile heads are trained jointlywith the Ranger optimizer [39] and a flat learning rate of 0.001 for all our experiments.
For the image super-resolution training, we augment the input dataset by using different levels ofdownsampled 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 thedifficulty by choosing a random threshold to create the mask – lower thresholds implies fewer pixelsare masked and vice-versa.
The mask is concatenated to the image resulting in a C + 1 -channelinput to the encoder, where C is the number of image channels.
The detailed description of the maskgeneration procedure can be found in the supplementary material.
Calibration and Evaluation For both the synthetic object experiments with CLEVR and faceexperiments with FFHQ, we train the quantile encoder on data points sampled from the latent spaceof the pretrained generative model.
This ensures that we have access to the true latents that resulted ineach 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 typesinclude image super-resolution and image inpainting.
The risk level α and the user-specified errorthreshold δ 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 visuallymeaningful uncertainties.
Description.
We train a quantile encoder with Algorithm 1.
Then we generate n = 5000 images bysampling latents and propagating them through the encoder-generator combination.
We use these nimages as calibration data for the procedure in 2.3.
Finally, we randomly sample a new test point, passit through the calibrated encoder, and form the uncertainty intervals in image space as in Section 2.3. [Figure]
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 the person is wearing glasses in thelower quantile image and not in the upper; hence, the model is not certain that the person is wearingglasses.
The model also expresses some uncertainty about the amount of gray versus brown hair.
Thisoutcome was predictable from the input image, where the fact that the person is wearing glasses isnot 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.
Thecalibrated quantiles cover the ground truth color value while the uncalibrated ones do not.3.3.2 Exploratory results with purposeful corruptions
Goal.
We probe the uncertainty quantification procedure to see if it will have the expected qualitativebehavior.
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 semanticmeaning.
Both image resolution and image masking were used as corruption models.
We qualitativelyanalyze the results by visualizing the predictions and perform a quantitative measurement as well bycomputing image based metrics. For the qualitative analysis, we use 500 images at each difficultylevel as inputs to our encoder.
Qualitative results.
Figure 4 shows the results of this experiment for Image inpainting (left) andsuper-resolution (right).
In the inpainting case, when nothing is masked, the quantiles are roughlyidentical.
When the eyes are masked, the quantiles indicate the model does not know if the personwas wearing glasses.
When the mouth is masked, the model expresses uncertainty as to whether thewoman is showing her teeth in the smile.
Finally, when almost everything is masked, the quantileimages are very different, representing individuals with entirely different identities.
Similar behaviorcan be observed in the super-resolution case.
The results are shown for two separate inputs; in bothcases, 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 predictionsare 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 hairshape, glasses, facial hair and perceived gender.
The results from both these experiments point to anexpected qualitative behavior of our proposed approach: the model exhibits more uncertainty withincreased information loss at the input.
Quantitative results.
For each input, we compute the calibrated uncertainty interval using ourapproach and compute the identity loss specified in Equation 7, and perceptual loss between the upperand lower edges.
We repeat the procedure for each image by varying the input difficulty, similarto 3.3.3. It can be observed from Table 1 that both perceptual and ID losses increase with increasingperceived input difficulty. This substantiates our claim that the calibrated quantiles display morevariability as we increase the difficulty of the task. Note that most of the style dimensions onlyaffect attributes like hair color/glasses/facial hair that do not necessarily change the identity of theindividual.
Given this observation, the change in ID loss is very much indicative of the variability ofthe quantile predictions.
[Table]
3.3.3 Interval sizes as a function of problem difficulty
Goal.
We seek to construct intervals that adapt to the uncertainty of the input, i.e. , result in lower
[Figure]
results in supplementary material.values for easier inputs and higher values for harder inputs.
This experiment tells us how informativeour intervals are about when the model makes an error.
Description of experiment.
For the image super-resolution case, we create the difficulty levels { easy , medium , hard } that correspond to { 1 x , 8 x , 32 x } downsampled versions.
All downsampled versions are resized to the same dimensions before presenting to the encoder.
For imageinpainting, we vary the masking threshold to create { easy , medium , hard } difficulty levels whichindicate the fraction of the image being masked, with { easy , medium , hard } corresponding to{ 10-15% , 40-50% , 90-95% } .
The results are computed on held out set sampled from the StyleGAN generative model pretrained on the FFHQ dataset.
In order to obtain the set size for each input,we scale the quantile width using the threshold obtained after the RCPS calibration procedure.
Results.
Figure 5 shows the set sizes for the super-resolution corruption model as a function ofproblem difficulty on two datasets, FFHQ and CLEVR.
As expected, the set sizes increase withincreasing problem difficulty indicating increasing uncertainty as corruption level increases.
Related Work
GANs for Inverse problems.
The remarkable image generation properties of recent approachessuch as BigGAN [10] and StyleGAN [21] has led to the increasing use of these models to solve inverse problems relating to image restoration such as image super-resolution and completion.
Allprior methods that use GANs for inverse problems such as the ones that use the pretrained generativemodel as an image prior [43, 8, 35, 33] or the ones that train an encoder model to project the inputinto the generator’s latent space [47, 39, 42, 38] focus on the accuracy of the point estimate but noton the uncertainty level of the input.
GANs for Interpretability.
Despite providing no guarantees of image likelihood, unlike others suchas Normalized Flow [23] or Score-based models [48], GANs have been used to develop interpretableapproaches to image generation.
The widespread use of GANs as opposed to other generative modelsin the interpretability is done due to the availability of an disentangled latent space [16, 46], which isa property we utilize in our work.
Quantile Regression.
Quantile regression was first proposed in [27].
Since then, many papers haveused the technique, applying it to machine learning [19, 34, 36], medical research [6], and more.
Most relevant to us is conformalized quantile regression [40], which gives quantile regression amarginal coverage guarantee using conformal prediction.
Our work instead uses risk-controllingprediction sets, a different distribution-free uncertainty quantification technique.
Conformal prediction and distribution-free uncertainty quantification.
At the core of ourproposed method is the distribution-free, marginal risk-control technique studied in [7] and [2].
These ideas have their roots in the distribution-free marginal guarantees of conformal prediction,proposed in [45].
Conformal prediction is a flexible method for forming prediction sets that satisfya marginal coverage guarantee, under no assumptions besides the exchangeability of the test pointwith the calibration data [45, 44, 31, 32, 30, 1].
Conformal prediction has been studied in computervision [18, 11, 5, 41, 3], natural language [13], drug discovery [14], criminal justice [9], and more.
We are not aware of work applying the notions of conformal prediction and quantile regression to thelatent space in generative models.
Conclusion
Experiments indicate the latent space uncertainty intervals express a useful semantic notion ofuncertainty previously unavailable in computer vision.
The intervals contain most of the true semanticfactors. These intervals are useful for semantically explaining a model’s uncertainties in a wayhumans can understand—for example, on high-level features in the image. Depending on theavailable disentangled latent space, this notion of uncertainty can be quite rich. Limitations ofour method include a) that we have only applied the technique to GAN generated data, b) thatthe calibration data must be reflective of the data distribution, and c) that we assume access toa disentangled latent space.
We see our work as part of a larger tapestry of results in generativemodels, and our technique will remain applicable as progress gets made on real-data backprojection,calibration under distribution shifts, and disentanglement.
Ethics
The ethics of generative modeling itself has been called into question given recent events, e.g. , thedevelopment of deep fakes.
Nonetheless, we believe the downstream consequences of this work willlikely be positive.
The techniques herein do not change the predictions of a generative model; theysimply provide a calibrated notion of its uncertainty in a relevant semantic space.
Thus, the standardcriticism 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 willprovide a sobering reliability assessment of these models, perhaps mitigating the chance of harmfulfailures.
Finally, although we use face datasets due to their ubiquity in this literature, we haveattempted to ethically treat topics like gender and race where they arise.