Scalable Infomin Learning
Abstract
The task of infomin learning aims to learn a representation with high utility while being uninformative about a specified target, with the latter achieved by minimising the mutual information between the representation and the target.
It has broad applications, ranging from training fair prediction models against protected attributes, to unsupervised learning with disentangled representations.
Recent works on infomin learning mainly use adversarial training, which involves training a neural network to estimate mutual information or its proxy and thus is slow and difficult to optimise.
Drawing on recent advances in slicing techniques, we propose a new infomin learning approach, which uses a novel proxy metric to mutual information.
We further derive an accurate and analytically computable approximation to this proxy metric, thereby removing the need of constructing neural network-based mutual information estimators.
Experiments on algorithmic fairness, disentangled representation learning and domain adaptation verify that our method can effectively remove unwanted information with limited time budget.
Introduction
Learning representations that are uninformative about some target but still useful for downstream applications is an important task in machine learning with many applications in areas including algorithmic fairness [1, 2, 3, 4], disentangled representation learning [5, 6, 7, 8], information bottleneck [9, 10] and invariant representation learning [11, 12, 13, 14].
A popular method for the above task is adversarial training [1, 3, 4, 7, 15, 2, 11], where two neural networks, namely the encoder and the adversary, are trained jointly to compete with each other.
The encoder’s goal is to learn a representation with high utility but contains no information about the target.
The adversary, on the contrary, tries to recover the information about the target from the learned representation as much as possible.
This leads to a minmax game similar to that in generative adversarial networks [16].
Adversarial training is effective with a strong adversary, however, it is often challenging to train the adversary thoroughly in practice, due to time constraints and/or optimisation difficulties [17, 18].
In fact, recent studies have revealed that adversarial approaches may not faithfully produce an infomin representation in some cases [19, 20, 18, 21, 22].
This motivates us to seek a good, adversarial training-free alternative for scalable infomin learning.
In this work, we propose a new method for infomin learning which is almost as powerful as using a strong adversary but is highly scalable.
Our method is inspired by recent advances in information theory which proposes to estimate mutual information in the sliced space [23].
We highlight the following contributions:
• We show that for infomin learning, a general, accurate estimate of mutual information or its bound is unnecessary: testing and minimising dependence randomly in some sliced space is sufficient; • We develop an analytical approximation to the aforementioned sliced independence test, along with a scalable training algorithm based on this approximation.
No adversarial training is needed.
Importantly, the proposed method can be applied to a wide range of infomin learning tasks without any constraint on the form of variables or any assumption about the distributions.
This contrasts our method to other adversarial training-free methods which are either tailored for discrete or univariate variables [24, 25, 25, 26, 22] or rely on variational approximation to distributions [10, 27, 19, 12].
Background
Infomin representation learning .
Let X ∈ R D be the data,
Y ∈ R D
′ be the target we want to predict from X .
The task we consider here is to learn some representation Z = f
( X ) that is useful for predicting Y but is uninformative about some target T ∈ R d .
Formally, this can be written as
[Equation]
where f is an encoder, L is some loss function quantifying the utility of Z for predicting Y and
I ( f ( X ); T ) quantifies the amount of information left in Z about T .
β controls the trade-off between utility and uninformativeness.
Many tasks in machine learning can be seen as special cases of this objective.
For example, by setting T to be (a set of) sensitive attributes e.g. race, gender or age, we arrive at fair representation learning [28, 1, 2, 21].
When using a stochastic encoder, by setting Y to be X and T to be some generative factors e.g., a class label, we arrive at disentangled representation learning [5, 6, 7, 8].
Similarly, the information bottleneck method [9, 10] corresponds to setting
T = X , which learns representations expressive for predicting Y while being compressive about X .
Adversarial training for infomin learning .
A key ingredient in objective (1) is to quantify
I ( f ( X ); T ) as the informativeness between f ( X ) and T .
One solution is to train a predictor for
T from f ( X ) and use the prediction error as a measure of I ( f ( X ); T ) [1, 3, 4, 29].
Another approach is to first train a classifier to distinguish between samples from p ( Z, T ) vs. p ( Z ) p ( T ) [7, 15] or to distinguish samples Z ∼ p
( Z | T ) with different T [2, 11], then use the classification error to quantify I ( f ( X ); T ) .
All these methods involve the training of a neural network t to provide a lower-bound estimate to I ( f ( X ); T ) , yielding a minmax optimisation problem
[Equation]
where ˆ I t ( f ( X ); T ) is an estimator constructed using t that lower-bounds I ( f ( X ); T ) .
The time complexity of optimising (2) is O
( L 1 L 2 ) where L 1 and L 2 are the number of gradient steps for the min and the max step respectively.
The strength of t is crucial for the quality of the learned representation [19, 20, 18, 21, 22].
For a strong adversary, a large L 2 is possibly needed, but this means a long training time.
Conversely, a weak adversary may not produce a truly infomin representation.
Methodology
We propose an alternative to adversarial training for optimising (1).
Our idea is to learn representation by the following objective, which replaces I ( f ( X ); T ) in objective (1) with its ‘sliced’ version:
[Equation]
where SI denotes the sliced mutual information, which was also considered in [23].
Informally, SI is a ‘facet’ of mutual information that is much easier to estimate (ideally has closed form) but can still to some extent reflect the dependence between Z and T .
Optimising (3) is then equivalent to testing and minimising the dependence between Z and T from one facet.
Importantly, while testing dependence through only a single facet may be insufficient, by testing and minimising dependence through various facets across a large number of mini-batches we eventually see I ( Z ; T ) → 0 .
We show one instance for realising SI whose empirical approximation ˆ SI has an analytic expression.
The core of our method is Theorem 1, which is inspired by [4, 23].
Theorem 1.
Let Z ∈ R D and T ∈ R
d be two random variables that have moments.
Z and T are statistically independent if and only if SI ( Z, T ) = 0 where SI ( Z, T ) is defined as follows
[Equation]
where ρ is the Pearson correlation, h, g : R → R are Borel-measurable non-constant functions, and θ ∈ S D − 1 , ϕ ∈ S d − 1 are vectors on the surfaces on D -dimensional and d -dimensional hyperspheres.
Proof .
See the Appendix.
We sketch here how this result relates to [4, 23].
[23] considers SI ( Z ; T ) , defined as the expected mutual information E
[ I ( θ ⊤ Z, ϕ ⊤ T )] of slices of Z and T , where the expectation is taken over respective Haar measures θ ∈ S D − 1 , ϕ ∈ S d − 1 .
Instead of considering the mutual information
I ( θ ⊤ Z, ϕ ⊤ T ) in the average case, we handle Pearson correlation over the supreme functions h, g as defined above, which links to R´enyi’s maximal correlation [4, 30, 31, 32] and has some interesting properties suitable for infomin representation learning (explained later).
We call θ and ϕ the slices for Z and T respectively, and θ ⊤ Z , ϕ ⊤ T the sliced Z and T respectively.
Intuitively, Theorem 1 says that in order to achieve I ( Z ; T ) → 0 , we need not estimate I ( Z ; T ) in the original space; rather we can test (and optimise) independence in the sliced space as realised by (4).
Note other realisations of the sliced mutual information SI can also be used.
The major merit of the realisation (4) is that (a) it is bounded in [0 , 1] , which eases hyperparameter tuning and makes learning stable; (b) it allows an analytic expression for its empirical approximation, as shown below.
Analytic approximation to SI .
An empirical approximation to (4) is
[Equation]
where
[Equation]
i.e., we approximate (4) by randomly sampling a number of slices θ, ϕ uniformly from the surface of two hyperspheres S D − 1 and S d − 1 and pick those slices where the sliced Z and the sliced T are maximally associated.
With a large number of slices, it is expected that (5) will approximate (4) well.
We refer to [23] for a theoretical analysis on the number of required slices in estimator-agnostic settings.
In Appendix B we also investigate empirically how this number will affect performance.
For each slicing direction, we further assume that the supreme functions h i , g j : R → R for that direction can be well approximated by K -order polynomials given sufficiently large K , i.e.
[Equation]
where σ ( · ) is a monotonic function which maps the input to the range of [ − 1 , 1] .
Its role is to ensure that σ ( a ) always has finite moments, so that the polynomial approximation is well-behaved.
Note no information is lost by applying σ ( · ) .
Here we take σ ( · ) as the tanh function.
Other approximations can also be used (e.g. those based on random features [33]), which can be explored in the future.
One may wonder if ignoring the higher-order terms σ ( a ) k ′ , k ′ >
K will cause inaccuracy in dependence modelling.
In Appendix A we show that if | ˆ h i ( a ) − h
i ( a ) | ≤ ϵ and | ˆ g j ( a ) − g j ( a ) | ≤ ϵ for ∀ i, j, a , then as ϵ →
we have | ρ ∗ ij − ˆ ρ ∗ ij |
=
o ( ϵ ) where ρ ∗ ij =
ρ ( h i ( θ ⊤ i Z ) , g j ( ϕ ⊤ j T )) is the true correlation and ˆ ρ ∗ ij is its K -order approximation.
In such case we have ˆ ρ ∗ ij ≈
whenever ρ ∗ ij ≈ 0 .
With this polynomial approximation, the solving of h i , g j reduces to finding their weights w i , v j :
[Equation]
This is known as canonical correlation analysis [34] and can be solved analytically by eigendecomposition.
Hence we can find the weights for all pairs of h i , g j by doing S 2 eigendecompositions.
In fact, the functions h i , g j for all i, j can be solved simultaneously by performing a larger eigendecomposition only once.
We do this by finding w, v that maximise the following quantity:
[Equation]
where
[Equation]
The benefits for solving f i , g j for all slices jointly are two-fold.
The first is better computational efficiency in practice, as it avoids invoking a for loop and has better affinity to modern deep learning infrastructure and libraries (e.g. Tensorflow [35] and PyTorch [36]) which are optimised for matrixbased operations.
We highlight that this joint-solving strategy will not violate our original objective (i.e. minimising (5)), as backed up by Theorem 2 below.
The intuition behind Theorem 2 is that if all Z ′ i and T ′ j together cannot achieve a high correlation, they alone can not either:
Theorem 2.
Provided that each h i , g j in (5) are K -order polynomials , given the sampled Θ
= { θ i } Si =1 , Φ = { ϕ j } Sj =1 , we have ˆ SI Θ , Φ ( Z ; T ) ≤ ϵ ⇒ sup
i,j sup h i ,g j ρ ( h i ( θ ⊤ i Z ) , g j ( ϕ ⊤ j T ))
≤ ϵ .
Proof .
See Appendix A.
Theorem 2 says that the solution of (6) yields an upper bound of (5).
This essentially means that we are safe to replace (5) with (6).
In addition to computational efficiency, another benefit for solving g j , h j jointly is stronger power in independence testing.
More specifically, it allows us to unify the power of different slices: while some sliced directions may be weak for detecting dependence, they together as a whole (by treating them as a ‘big slice’ [37]) can compensate for each other, yielding a more powerful test.
For these reasons, we use ˆ SI Θ , Φ ( Z ; T ) as the approximation to SI ( Z ; T ) .
Input: data D = { X ( n ) , Y ( n ) , T ( n ) } N n =
Output: Z = f
( X ) that optimises (1)
Hyperparams: β , N ′ , L 1 , L
Parameters: encoder f , MI estimator t for l 1 in 1 to L 1 do sample mini-batch B from D sample D ′ from D whose size N ′ <
N // Max-step for l 2 in 1 to L 2 do t ← t + η ∇ t ˆ I t ( f ( X ); T ) with data in D ′ end for //
Min-step f ← f − η ∇ f
[ L ( f ( X ); Y )
+
β
ˆ
I t ( f ( X ); T ) ] with data in B end for return Z = f ( X )
Algorithm
Input: data D = { X ( n ) , Y ( n ) , T ( n ) } N n =
Output: Z = f
( X ) that optimises (1)
Hyperparams: β , N ′ , L , S
Parameters: encoder f , weights w, v in ˆ SI for l in 1 to L do sample mini-batch B from D sample D ′ from D whose size N ′ <
N // Max-step sample S slices Θ
= { θ i } Si =1 , Φ = { ϕ j } Sj
= solve the parameters w, v in ˆ SI analytically with Θ , Φ , D ′ by eigendecomposition //
Min-step f ← f − η ∇ f [ L ( f ( X ); Y )
+ β ˆ SI ( f ( X ); T ) ] with data in B end for return Z = f
( X )
Mini-batch learning algorithm .
Given the above approximation (6) to SI , we can now elaborate the details of mini-batch learning.
For each mini-batch B , we execute the following steps:
• Max-step .
Sample S slices Θ = { θ i } S i =1 , Φ = { ϕ j } S j =1 and a subset of the data D ′ ⊂ D (here |D ′ | can be larger than |B| ).
Learn the weights w, v of (6) with Θ , Φ , D ′ by eigendecomposition; • Min-step .
Set ˆ SI ( Z, T ) =
ρ ( w ⊤ Z ′ , v ⊤ T ′ ) with w, v solved in the max-step and Z ′ , T ′ defined in (6).
Update f by SGD: f ← f − η ∇ f
[ L ( f ( X ) , Y ) +
β ˆ SI ( f ( X ) , T )] with the mini-batch data.
The whole learning procedure is shown in Algorithm 2.
Compared to that of adversarial methods [2, 11, 1, 3, 21, 4] as shown in Algorithm 1, we replace the optimisation of neural net in the max step with an analytical eigendecomposition step.
The time complexity of eigendecomposition is O ( S 3 ) .
As an optional strategy, one may actively seek informative slices for dependence detection.
This can be done by optimising ˆ SI w.r.t the sampled slices by a few gradient steps (e.g. 3) after the max step:
[Equation]
which is still very cheap to execute.
Such a strategy is particularly useful when most of the sampled slices are weak in detecting dependence, which typically happens in later iterations of learning where I ( Z ; T ) ≈ 0 .
We can activate such strategy whenever the estimated SI is very low (e.g.
ˆ SI ( Z ; T ) < 0 . 05 ).
Since the optimisation of slices only happens in late learning iterations, and we only apply a small number of gradient steps, the overall execution time will only increase negligibly.
Related works
Neural mutual information estimators .
A set of neural network-based methods [38, 39, 40, 41] have been proposed to estimate the mutual information (MI) between two random variables, most of which work by maximising a lower bound of MI [42].
These neural MI estimators are in general more powerful than non-parametric methods [43, 44, 45, 46] when trained thoroughly, yet the time spent on training may become the computational bottleneck when applied to infomin learning.
Upper bound for mutual information .
Another line of method for realising the goal of infomin learning without adversarial training is to find an upper bound for mutual information [10, 19, 12, 47, 48].
However, unlike lower bound estimate, upper bound often requires knowledge of either the conditional densities or the marginal densities [42] which are generally not available in practice.
As such, most of these methods introduce a variational approximation to these densities whose choice/estimate may be difficult.
Our method on the contrary needs not to approximate any densities.
Slicing techniques .
A series of successes have been witnessed for the use of slicing methods in machine learning and statistics [49], with applications in generative modelling [50, 51, 52, 53], statistical test [54] and mutual information estimate [23].
Among them, the work [23] who proposes the concept of sliced mutual information is very related to this work and directly inspires our method.
Our contribution is a novel realisation of sliced mutual information suitable for infomin learning.
Fair machine learning .
One application of our method is to encourage the fairness of a predictor.
Much efforts have been devoted for the same purpose, however most of the existing methods can either only work at the classifier level [25, 24, 4, 55], or only focus on the case where the sensitive attribute is discrete or univariate [22, 56, 26, 55, 28], or require adversarial training [2, 11, 1, 3, 21, 4].
Our method on the contrary has no restriction on the form of the sensitive attribute, can be used in both representation level and classifier level, and require no adversarial training of neural networks.
Disentangled representation learning .
Most of the methods in this field work by penalising the discrepancy between the joint distribution P = q ( Z ) and the product of marginals Q = (cid:81)
Dd q ( Z d ) [6, 57, 27, 58, 7].
However, such discrepancy is often non-trivial to estimate, so one has to resort to Monte Carlo estimate ( β -TCVAE [27]), to train a neural network estimator (FactorVAE [7]) or to assess the discrepancy between P and Q by only their moments (DIP-VAE [57]).
Our method avoids assessing distribution discrepancy directly and instead perform independence test in the sliced space.
Experiments
We evaluate our approach on four tasks: independence testing, algorithmic fairness, disentangled representation learning, domain adaptation.
Code is available at github.com/cyz-ai/infomin .
Evaluation metric .
To assess how much information is left in the learned representation Z ∈ R D about the target T ∈ R K , we calculate the R ´ enyi’s maximal correlation ρ ∗ ( Z, T ) between Z and T :
[Equation]
8) which has the properties ρ ∗ ( Z, T )
=
if and only if Z ⊥ T and ρ ∗ ( Z, T )
= 1 if h ( Z ) =
g ( T ) for some deterministic functions h, g [30].
One can also understand this metric as the easiness of predicting (the transformed) T from Z , or vice versa.
As there is no analytic solution for the supremum in (8), we approximate them by two neural networks h, g trained with SGD.
Early stopping and dropout are applied to avoid overfitting.
The reliability of this neural approximation has been verified by the literature [4] and is also confirmed by our experiments; see Appendix B.
This metric is closely related to existing metrics/losses used in fairness and disentangled representation learning such as demographic parity (DP) [24] and total correlation (TC) [7].
For example, if ρ ∗ ( Z, T ) → 0 then it is guaranteed that ˆ Y ⊥ T for any predictor ˆ Y = F ( Z ) , so ρ ∗ ( Z, T ) is an upper bound for DP .
Similarly, ρ ∗ ( Z, T ) coincides with TC which also assesses whether Z ⊥ T .
In additional to this metric, we will also use some task-specific metric; see each experiment below.
Baselines .
We compare the proposed method (denoted as “Slice”) with the following approaches
[Figure]
• Pearson , which quantifies I ( Z ; T ) by the Pearson correlation coefficient
[Figure]
[Equation]
[Figure] [Equation] [Figure] It was used in [5, 57] as an easy-to-compute proxy to MI to learn disentangled representations; • dCorr , i.e. distance correlation, a non-parametric method for the quantifying the independence between two vectors [46].
It was applied in [41] as a surrogate to MI for representation learning; • Neural R ´ enyi , an adversarial method for fair machine learning [25] which quantifies I ( Z ; T ) by the R ´ enyi correlation ρ ∗ ( Z, T ) =
sup h,g ρ ( h ( Z ) , g ( T )) with h, g approximated by neural networks.
It can be seen as training a predictor to predict (the transformed) T from Z and is closely related to many existing methods in algorithmic fairness and domain adaptation [1, 2, 3, 4, 29]; • Neural TC , an adversarial method for learning disentangled representation [7, 62] which quantifies I ( Z ; T ) by the total correlation
TC ( Z, T ) =
KL [ p ( Z, T )
∥
p ( Z ) p ( T )] .
To computes TC, a classifier is trained to classify samples from p ( Z, T ) and samples from p ( Z ) p ( T ) .
This method can also be seen as a variant of the popular MINE method [38] for mutual information estimate.
• v-CLUB , i.e. variational contrastive log upper bound, which introduces a (learnable) variational distribution q ( T | Z ) to form an upper bound of MI [47]: I ( Z ; T ) ≤ E p ( Z,T ) [log q ( T | Z )] − E
p ( Z ) p ( T ) [log q ( T | Z )] .
Like adversarial method, q ( T | Z ) can be learned by a few gradient steps.
For a fair comparison, for adversarial training-based approaches (i.e. Neural R ´ enyi, Neural TC), we ensure that the training time of the neural networks in these methods is at least the same as the execution time of our method or longer.
We do this by controlling the number of adversarial steps
L 2 in Algorithm 1.
The same setup is used for v-CLUB.
See each experiment for the detailed time.
Hyperparameter settings .
Throughout our experiments, we use 200 slices.
We find that this setting is robust across different tasks.
An ablation study on the number of slices is given in Appendix B.
The order of the polynomial used in (6) namely K is set as K = 3 and is fixed across different tasks.
Computational resource .
All experiments are done with a single NVIDIA GeForce Tesla T4 GPU.
5.1 Independence testing
We first verify the efficacy of our method as a light-weight but powerful independence test.
For this purpose, we investigate the test power of the proposed method on various synthetic dataset with different association patterns between two random variables X ∈ R 10 , Y ∈ R 10 and compared to that of the baselines.
The test power is defined as the ability to discern samples from the joint distribution p ( X, Y ) and samples from the product of marginal p ( X ) p ( Y ) and is expressed as a probability p ∈ [0 , 1] .
The data is generated as Y =
(1 − α ) ⟨ t ( A X ) ⟩ + αϵ, X d ∼ U [ − 3 , 3] , A dd = 1 , A dk = 0 .
2 , ϵ ∼ N ( ϵ ; 0 , I ) , α ∈ (0 , 1) and ⟨·⟩ is a scaling operation that scales the operand to the range of [0 , 1] according to the minimum and maximum values in the population.
The function t ( · ) determines the association pattern between X and Y and is chosen from one of the following: t ( a )
= a, a 2 , sin( a ) , tanh ( a ) .
The factor α controls the strength of dependence between X and Y .
All tests are done on 100 samples and are repeated 1,000 times.
We choose this sample number as it is a typical batch size in mini-batch learning.
For methods involving the learning of parameters (i.e. Slice, Neural R´enyi, Neural TC), we learn their parameters from 10,000 samples.
The time for learning the parameters of Slice, Neural R´enyi and Neural TC are 0.14s, 14.37s, 30.18s respectively.
For completeness, we also compare with the ‘optimal test’ which calculates the R´enyi correlation ρ ∗ ( X, Y )
=
ρ ( h ( X ) , g ( Y )) with the functions h, g exactly the same as the data generating process.
[Table] [Table] Figure 1 shows the power of different methods under various association patterns t and dependence levels α .
Overall, we see that the proposed method can effectively detect dependence in all cases, and has a test power comparable to neural network-based methods.
Non-parametric tests, by contrast, fail to detect dependence in quadratic and periodic cases, possibly due to insufficient power.
Neural
TC is the most powerful test among all the methods considered, yet it requires the longest time to train.
We also see that the proposed method is relatively less powerful when α ≥ 0 .
8 , but in such cases the statistical dependence between X and Y is indeed very weak (also see Appendix B).
The results suggest that our slice method can provide effective training signals for infomin learning tasks.
5.2 Algorithmic fairness
For this task, we aim to learn fair representations Z ∈ R 80 that are minimally informative about some sensitive attribute T .
We quantify how sensitive Z is w.r.t T by R ´ enyi correlation ρ ∗ ( Z, T ) calculated using two neural nets.
Smaller ρ ∗ ( Z, T ) is better.
The utility of the learned representation i.e., L ( Z ; Y ) is quantified by ρ ∗ ( Z, Y ) .
This formulation for utility, as aforementioned, is equivalent to measuring how well we can predict Y from Z .
In summary, the learning objective is:
[Equation]
where ˆ I ( Z ; T ) is estimated by the methods mentioned above.
For each dataset considered, we use 20,000 data for training and 5,000 data for testing respectively.
We carefully tune the hyperparameter β for each method so that the utility ρ ∗ ( Z ; Y ) of that method is close to that of the plain model (i.e. the model trained with β = 0 , denoted as N/A below; other experiments below have the same setup).
For all methods, we use 5 , 000 samples in the max step (so N ′
= 5 , 000 in Algorithm 1, 2).
US Census Demographic data .
This dataset is an extraction of the 2015 American Community
Survey, with 37 features about 74,000 census tracts.
The target Y to predict is the percentage of children below poverty line in a tract, and the sensitive attribute T is the ratio of women in that tract.
The result is shown in Table 1.
From the table we see that the proposed slice method produces highly fair representation with good utility.
The low ρ ∗ ( Z, T ) value indicates that it is unlikely to predict
T from Z in our method.
While adversarial methods can also to some extent achieve fairness, it is still not comparable to our method, possibly because the allocated training time is insufficient (in
Appendix B we study how the effect of the training time).
Non-parameteric methods can not produce truly fair representation, despite they are fast to execute.
v-CLUB, which estimates an upper bound of MI, achieves better fairness than adversarial methods on average, but has a higher variance [63].
UCI Adult data .
This dataset contains census data for 48,842 instances, with 14 attributes describing their education background, age, race, marital status, etc.
Here, the target Y to predict is whether the income of an instance is higher than 50,000 USD, and the sensitive attribute T is the race group.
The result is summarised in Table 2.
Again, we see that the proposed slice method outperforms other methods in terms of both fairness and utility.
For this dataset, Neural R´enyi also achieves good fairness, although the gap to our method is still large.
Neural TC, by contrast, can not achieve a [Figure] [Figure] comparable level of fairness under the time budget given — a phenomenon also observed in the US
Census dataset.
This is possibly because the networks in Neural TC require longer time to train.
The v-CLUB method does not work very satisfactorily on this task, possibly because the time allocated to learn the variational distribution
q ( Z | T ) is not enough, leading to a loose upper bound of I ( Z ; T ) .
5.3 Disentangled representation learning
We next apply our method to the task of disentangled representation learning, where we wish to discover some latent generative factors irrelevant to the class label T .
Here, we train a conditional autoencoder X
≈ G ( Z, T )
to learn representation Z = E ( X ) which encodes label-irrelevant information of X .
The target to recover is Y = X .
The utility of Z is therefore quantified as the reconstruction error:
L ( Z ; Y ) =
E [ ∥ G ( Z, T ) − X ∥ 22 ] , resulting in the following learning objective:
[Equation]
The conditional autoencoder uses a architecture similar to that of a convolutional GAN [64], with the difference being that we insert an adaption layer Z ′ =
MLP ( Z, T ) before feeding the features to the decoder.
See Appendix B for the details of its architecture.
All images are resized to 32 × 32.
For all methods, we use 10 , 000 samples in the max step (so N ′ = 10 , 000 in Algorithms 1 and 2).
Dsprite .
A 2D shape dataset [65] where each image is generated by four latent factors: shape, rotation, scale, locations.
Here the class label T is the shape, which ranges from (square, ellipse, heart).
For this dataset, we train the autoencoder 100 iterations with a batch size of 512.
The dimensionality of the representation for this task is 20 i.e. Z ∈ R 20 .
As in the previous experiments, we provide quantitative comparisons of the utility and disentanglement of different methods in Table 3.
In addition, we provide a qualitative comparison in Figure 2 which visualises the original image, the reconstructed image and the reconstructed image with a swapped label.
From Table 3, we see that the proposed method achieve very low ρ ∗ ( Z, T ) while maintaining good MSE, suggesting that we [Table] [Table] may have discovered the true label-irrelevant generative factor for this dataset.
This is confirmed visually by Figure 2(b), where by changing T in reconstruction we only change the style.
By contrast, the separation between T and Z is less evident in adversarial approach, as can be seen from Table 3 as well as from Figure 2(a) (see e.g. the reconstructed ellipses in the third column of the figure.
They are more like a interpolation between ellipses and squares).
CMU-PIE .
A colored face image dataset [66] where each face image has different pose, illumination and expression.
We use its cropped version [67].
Here the class label T is the expression, which ranges from (neutral, smile).
We train an autoencoder with 200 iteration and a batch size of 128.
The dimensionality of the representation for this task is 128 i.e. Z ∈ R 128 .
Figure 3 and Table 4 shows the qualitative and quantitative results respectively.
From Figure 3, we see that our method can well disentangle expression and non-expression representations: one can easily modify the expression of a reconstructed image by only changing T .
Other visual factors of the image including pose, illumination, and identity remain the same after changing T .
Adversarial approach can to some extent achieve disentanglement between Z and T , however such disentanglement is imperfect: not all of the instances can change the expression by only modifying T .
This is also confirmed quantitatively by Table 4, where one can see the relatively high ρ ∗ ( Z, T ) values in adversarial methods.
For this task, v-CLUB also achieves a low ρ ∗ ( Z, T ) value, though it is still outperformed by our method.
5.4 Domain adaptation
We finally consider the task of domain adaptation, where we want to learn some representation Z that can generalise across different datasets.
For this task, a common assumption is that we have assess to two dataset D s =
{ X ( i ) , Y ( i ) } ni =1 and D t
= { X ( j ) } mj =1 whose classes are the same but are collected differently.
Only the data in D s has known labels.
Following [47], we learn Z as follows:
[Equation]
where Z c , Z d are disjoint parts of Z that encode the content information and the domain information of X separately.
C is the content classifier that maps X to a ( K − 1) -simplex
( K is the number of classes) and D is the domain classifier that distinguishes the domain from which X comes. Since the
[Figure] [Table] classifier C only sees labels in D s , we call D s the source domain and D t the target domain.
For the two encoders f c and f d , we use Resnets [68] with 7 blocks trained with 100 iterations and a batch size of 128.
Here Z c , Z d ∈ R 256 .
We use N ′
= 5 , 000 samples in the max step for all methods.
MNIST → MNIST-M .
Two digit datasets with the same classes but different background colors.
Both datasets have 50,000 training samples.
Table 5 shows the result, indicating that our method can more effectively remove the information about the domain.
This is further confirmed by the T-SNE [69] plot in Figure 4, where one can hardly distinguish the samples of Z c from the two domains.
This naturally leads to a higher target domain accuracy Acc ( ˆ Y t ) than other methods.
CIFAR10 → STL10 .
Two datasets of natural images sharing 9 classes.
There are 50,000 and 5, training samples in the two datasets respectively.
Following existing works [70, 71, 47], we remove the non-overlapping classes from both datasets.
Table 5 and Figure 4 show the result.
Again, we see that our method can largely remove domain information. However, the target domain accuracy is still significantly below the source accuracy, suggesting additional methods might help.
Conclusion
This work proposes a new method for infomin learning without adversarial training.
A major challenge is how to estimate mutual information accurately and efficiently, as MI is generally intractable.
We sidestep this challenge by only testing and minimising dependence in a sliced space, which can be achieved analytically, and we showed this is sufficient for our goal.
Experiments on algorithmic fairness, disentangled representation learning and domain adaptation verify our method’s efficacy.
Through our controlled experiments, we also verify that adversarial approaches indeed may not produce infomin representation reliably – an observation consistent with recent studies.
This suggests that existing adversarial approaches may not converge to good solutions, or may need more time for convergence, with more gradient steps needed to train the adversary fully.
The result also hints at the potential of diverse randomisation methods as an alternative to adversarial training in some cases.
While we believe our method can be used in many applications for societal benefit (e.g. for promoting fairness), since it is a general technique, one must always be careful to prevent societal harms.
Acknowledgement
AW acknowledges support from a Turing AI Fellowship under grant EP/V025279/1, The Alan Tur-