Unsupervised Attribute Alignment for Characterizing Distribution Shift
Abstract
Detecting and addressing distribution shift is an important task in machine learning.
However, most of the machine learning solutions to deal with distribution shift lack the capability to identify the key characteristics of such a shift and present it to humans in an interpretable way.
In this work, we propose a novel framework to compare two datasets and identify distribution shifts between the datasets.
The key challenge is to identify generative factors of variation, which we refer to as attributes , that characterize the similarities and differences between the datasets.
Producing this characterization requires ﬁnding a set of attributes that can be aligned between the two datasets and sets that are unique.
We address this challenge through a novel approach that performs both attribute discovery and attribute alignment across the two distributions.
We evaluate our algorithm’s effectiveness at accurately identifying these attributes in two separate experiments, one involving two variants of MNIST and a second experiment involving two versions of dSprites.
Introduction
Machine learning (ML) algorithms traditionally assume that data is drawn from a single, stationary distribution.
In many real world scenarios, however, such an assumption likely will not hold, leading to uncontrollable or sometimes catastrophic failures [22, 8].
As a result, detecting, quantifying and compensating for the changes in distribution are of paramount importance.
Many recent methods propose making models more robust to arbitrary shifts in the data at test time [15, 18, 8, 20].
Most of these techniques, as well as the efforts to make models more robust, have largely focused on learning to make invariant predictions despite these shifts, treating them as noise or nuisance variables.
In many cases, characterizing precisely how the test distribution is different from the training distribution can be extremely useful.
Unfortunately, despite the large body of work in this space, surprisingly few studies focus on characterizing shifts and communicating this information in a form that is understandable to human users [19].
Characterizing distribution shift is especially important in scientiﬁc applications where ML models are ﬁrst trained in a simulator [4, 14] before being applied to make inferences on real-world data.
Simulators provide a cost-effective way to generate more samples in the experimental design space, but the ﬁdelity of the simulated data to real-world data is an obvious concern.
Characterizing the distribution shifts between simulated and real-world data can give us insights into the knowledge gap in our current understanding of scientiﬁc principles that need to be addressed.
One critical component to achieve such a system is the unsupervised discovery of the underlying generative factors, or attributes, in a data distribution.
With the development of high quality generative models [6, 13, 11], this goal can be achieved using unsupervised attribute discovery methods [7, 21, 23], which show latent spaces learn the underlying attribute structure of the data.
Here, concept vectors [12], i.e. linear directions in the latent space of a generative model, that correspond to speciﬁc data attributes, are identiﬁed; moving along these directions induces meaningful changes to the semantics of an output image.
However, to characterize the difference between two distributions, it is crucial to align them ﬁrst.
Generative methods have been used extensively to map between two distributions, implicitly achieving alignment between them, like unpaired image translation methods using CycleGAN [24] and one-to-many mapping by separating image content and style using MUNIT [10].
Most of these methods [2] are motivated by, and geared toward, generating high quality translated images rather than explaining shifts.
In other words, they help identify “parallel” or analogous attributes in both the distributions, implicitly making an assumption that a one-to-one mapping exists between the two datasets, which is highly restrictive.
The problem of interest here is identifying similar and dissimilar attributes given two pre-trained generative models, particularly when the assumption of a one-to-one mapping cannot be made.
In this work, we aim to ﬁll in this gap and introduce a novel approach to provide attribute level characterization of the distribution shift.
We achieve this through a joint framework that combines attribute discovery and attribute alignment across distributions.
Our method identiﬁes linear directions that encode both unique and shared attributes from the respective generative models derived from the original data distributions, which we empirically demonstrate on two simple tasks.
Our key contributions are as follows: 1) to propose a novel perspective to characterize the distribution shift through attribute alignment across datasets and 2) to introduce a framework that solves the joint attribute discovery and attribute alignment problem.
Methodology
Unsupervised attribute alignment between data distributions requires us to solve both the attribution discovery as well as the alignment tasks.
Ideally, we want to solve both tasks simultaneously so that attributes can be more reliably identiﬁed.
Many existing unsupervised GAN attribute discovery works [7, 21] rely on closed-form optimization to extract attribute directions, thereby requiring a post-hoc attribute alignment step.
In comparison, the recent work by Voynov et al.
[23] approaches the problem from a rather different perspective, in which a predictive model can be learned together with the concept directions for predicting the respective semantic shifts.
In this work, we leverage the ﬂexibility of a similar formulation to solve the joint optimization problem.
Formulation Let G be a pre-trained generator network that learns a mapping from a prior distribution, e.g., Gaussian P ( z ) =
N (0 , I d ) , to a data distribution G : P ( z ) (cid:55)→ P ( X ) .
Here, we deﬁne an attribute as a direction, w ∈ R d , in the latent space, that corresponds to a meaningful semantic change as G ( z ) → G ( z + (cid:15)w ) , where the extent of change is determined by (cid:15) .
As such, attribute discovery is the problem of identifying such directions to recover known semantic factors (labels such as digit thickness, shape, color etc.) in a data distribution.
To make this explicit, we associate the k th direction, w k , with a known semantic factor, α k ≈ w k .
Next, we describe the unsupervised attribute discovery problem for a single generative model formulated in [7], which we extend to multiple models subsequently.
Let W K × d be a learnable matrix, with hyper-parameter K corresponding to the total number of attributes to discover from the data distribution, applied to the d -dimensional latent space of the generative network.
The W matrix contains all the possible directions W =
[ w 1 , . .
. , w K ] , used to explore the latent space.
In order to identify meaningful changes in an unsupervised fashion, W is trained based on the fact that these directions are expected to be predictable – i.e., a discriminator network, D , is trained to take a pair of generated data as input, ( G ( z ) , G ( z + W ( (cid:15)k )) , to predict both (cid:15) and the speciﬁc direction k that induced the change.
The output of D is a probability distribution over the K attributes to learn P K and a prediction of the scalar ˆ (cid:15) .
Models D and W are jointly trained to minimize the cross entropy loss L CE of D ’s prediction of attribute k , k sampled k mean absolute error Loss L MAE between
[Equation]
Where λ = 0 .
25 is a hyperparameter to weight the L MAE
Attribute discovery for multiple generators We now extend this setup to the case of multiple generative models, in order to characterize distribution shifts using attribute level alignment.
Let G and G 2 be two pre-trained generators for both the data distributions of interest, and W K 1 ×
d 1and
W K 2 ×
d 2are the learnable matrices.
Unlike the case for a single model, the discriminator network D , must act on two sets of image pairs (one normal-perturbed pair from each generator).
Moreover, we modify its output to be M different attributes, such that they are composed of K 1 attributes from the ﬁrst data distribution and the K 2 attributes of the second data distribution.
M is composed of unique attributes for both data distributions, as well as shared directions, M > K 1 and M > K 2 .
If dataset one is a strict subset of dataset two, then we can deﬁne M = K 2 and K 1 < K 2 .
A more principled way of deciding values of K and M could be to incorporate progressive learning, where K is iteratively increased each time an accuracy threshold is reached.
Earlier learned directions would be semantically easier to discovery, and could help guide a human’s exploration of the learned attributes.
We leave this iterative learning as investigation for future work.
Since this process requires estimation of the attributes and alignment of them to be similar across a distribution shift, the problem is fairly ill-posed with a high likelihood of reducing to trivial solutions.
As such, we place different kinds of constraints (orthogonality, sparsity, disentanglement) to prevent these failure cases.
We illustrate the entire model in Figure 1.
Generative model choice Since the joint alignment and discovery problem is ill-posed, we ﬁnd that disentanglement of the latent space, using a β -VAE [9] helps improve the optimization to ﬁnd good solutions.
As a result, in our studies we use the beta-VAE’s decoder as our generative model G .
However, disentanglement on its own is insufﬁcient, since there is a risk of the model combining multiple attributes that may appear unique to D , resulting in a failure to ﬁnd the truly unique attributes.
As a result, we apply both a sparsity regularization and an additional spectral norm orthogonality loss [1] on the learned W 1 and W 2 matrices.
The sparsity term ensures that not too many independent attributes are being modiﬁed at the same time (e.g. preventing a direction that modiﬁes thickness, color, and digit as the same time).
Orthogonality loss ensures any two learned directions are not nearly identical.
For example, without this loss, many directions would learn the same transformation into an eights digit, but with slightly altered rotations.
The sparsity loss is enforced using a simple (cid:96) 1 -norm constraint: L sparse ( W ) =
(cid:80) i,j || W i,j || 1 , and the orthonormality constraint : L ortho ( W ) =
σ ( W T W − I ) , where σ is the spectral norm.
Finally, the overall cost function is denoted as follows:
[Equation]
for all experiments we set λ (cid:15) = 0 .
25 following [23].
And λ s = 0 . 1 , and λ o = 0 .
01 were set through visual inspection, where smaller values caused no changes, and higher values learned unmeaningful directions or degenerated into identical directions.
Unidirectional Attributes
In the original formulation of unsupervised attribute discovery, the scalar value (cid:15) is sampled uniformly from U ( − (cid:15), (cid:15) ) .
This seems reasonable as a learned attribute can be added or removed (e.g. thick digits in the positive direction versus thin digits in the negative direction).
However, having negative (cid:15) values does not work when aligning two data distributions.
While some semantic attributes, such as thickness, make sense to be linearly deﬁned in both the positive and negative directions, there is too much ﬂexibility in the learned latent space for other attributes such as classes.
For example, consider aligning the shared attributes between two datasets which both contain classes 0, 1 and 2.
The latent space for dataset one could learn class 0 in the negative direction and class 1 in the positive, whereas dataset two could learn class 0 in the negative direction but class 2 in the positive. It would therefore be impossible to align these attribute directions.
For this reason, we simply change the scalar to be sampled uniformly from U (0 .
1 ∗ (cid:15),
Experiments
Datasets As a ﬁrst foray into attribute alignment, we selected two experiments with attributes that can be semantically identiﬁed quickly.
To achieve this effort, we design two experiments where one data distribution is a strict subset of the other: one distribution will contain only aligned attributes, and the other will contain both aligned and unique attributes.
Therefore, we conduct two main experiments: MNIST [3] versus colored-MNIST [5], and dSprites versus square dSprites [17].
MNIST contains images with a single channel and 28x28 dimensions, where colored MNIST randomly applies color to a given image to get 3-channeled instances.
dSprites is a synthetic dataset of 5 independent attributes (location, scale, rotation, shape), and we create a modiﬁed version with only square shapes.
Our goal is to leverage the ﬁndings from these datasets to then apply them to more sophisticated distributions, such as faces versus anime faces, or experimental scientiﬁc data versus simulated scientiﬁc data.
Implementation details We train a β -VAEs on each of the 4 datasets, with β = 4 . 0 .
We use a latent size d = 10 for MNIST experiments and d = 5 for dSprites experiments.
The values of K 1 and K 2 are somewhat arbitrary, so we set them to the value of the latent size for the subsets, set d-sprites K 2 = 6 and set color-MNIST is set to K 2 = 15 . We use a LeNet [16] backbone as our discrimintive model D for all our experiments.
The latent points z are always sampled from a normal distribution N (0 , I ) matching the dimensionality of the original generator. The direction index i is sampled from the uniform distribution U { 1 , K 1 } or U { 1 , K 2 } . We ﬁnd using a shift scalar multiplier (cid:15) maximum of 3 for both all experiments to provide the most consistent training results, however this value is not always large enough at test time.
In order to ensure the l 2 -norms of our learned W matrices are not too high, we force all W matrix columns to have a unit vector length, as the cross entropy loss learns W with high norms to simplify classiﬁcation without the constraint.
Lastly, all model are jointly optimized using an Adam optimizer with learning rate of 0 .
01 , trained for 5000 iterations with a mini-batch size of 32. We found additional training for these relatively lower dimensional image datasets to be unnecessary. Results We show the results of attribute alignment for our two experiments in Figures 2 and 3. The former shows the learned, aligned, attributes between (left) dSprites squares and full dSprites and (right) MNIST and colored MNIST.
The dSprites experiment mostly learns positional alignment, moving a given shape between the cardinal positions in directions 0 to 3, and rotation in direction 4.
The MNIST experiment shows changes in the digit class, the orientation of a digit, and the intensity of the instance.
We present more examples for both experiments in the supplemental material.
Discussion
The proposed method provides the ﬁrst unsupervised attribute alignment method (to the extent of our knowledge) for characterizing distribution shift.
We solve the problem through a joint optimization of attribute discovery and attribute alignment.
It provides an intuition and an accessible way to characterize semantics shifts between distributions.
As the ﬁrst solution to this particularly challenging problem, we are aware and plan to improve upon certain limitations of the proposed approach.
Since we learn the alignment regarding the latent spaces of generative models, the quality and properties of these models would have a great impact on the ﬁnal quality of the result.
Moreover, the current approach makes the assumption that semantically meaningful attributes are more predictive for the attribute discovery [23], which may not hold for true in a more general setting.
To this end, a more effective attribute search scheme/criteria is likely required to make the overall system robust and applicable to wider range of generative models.
Acknowledgement
This work performed under the auspices of the U.S. Department of Energy by Lawrence Livermore
National Laboratory under Contract DE-AC52-07NA27344.
The manuscript is reviewed and released under LLNL-CONF-827845-DRAFT.
A Appendix
Additional examples Figure 4 shows all learned directions from the main text’s example in Figure 2.
Furthermore, we provide additional random samples from our experiments.
Figure 5 and 6 shows a side-by-side example of the shared directions for our MNIST and dSprites experiments respectively; ﬁgures 7 and 8 show the unique attributes learned for MNIST and dSprites.
F i gu r e 5 : m n i s t v s c m n i s t s h a r e d
d
i r ec ti on s .
F i gu r e 6 : s qu a r e s v
s f u ll d s p r it e s s h a r e d d
i r ec ti on s .