Estimating Total Correlation with Mutual Information Bounds
Abstract
Total correlation (TC) is a fundamental concept in information theory to measure the statistical dependency of multiple random variables.
Recently, TC has shown effectiveness as a regularizer in many machine learning tasks when minimizing/maximizing the correlation among random variables is required.
However, to obtain precise TC values is challenging, especially when the closed-form distributions of variables are unknown.
In this paper, we introduced several sample-based variational TC estimators.
Speciﬁcally, we connect the TC with mutual information (MI) and constructed two calculation paths to decompose TC into MI terms.
In our experiments, we estimated the true TC values with the proposed estimators in different simulation scenarios and analyzed the properties of the TC estimators
Introduction
Statistical dependency measures the correlation of random variables or factors in models, which is often an important concern in various scientiﬁc domains including statistics [12, 15], robotics [16, 4], bioinformatics [18, 24], and machine learning [7, 1, 14].
In recent deep learning studies, statistical dependency has increasingly served as learning objectives or regularizers for neural network training, and has achieved improvement in terms of model robustness [25], generalizability [1], interpretablity [7, 9], etc .
Among statistical dependency measurements, mutual information (MI) is commonly used in machine learning.
Given two random variables x , y , the mutual information is deﬁned as:
[Equation]
Recently, mutual information has shown signiﬁcant improvement when applied as a training criterion on learning tasks, such as conditional generation [7], domain adaptation [11], representation learning [6], and fairness [23].
However, MI can only handle the statistical dependency between two variables.
When considering optimization of correlation among multiple variables, MI requires computation of each variable pair, which leads to a quadratic increase in computation cost.
To address this problem, total correlation (TC) has been proposed by extending MI to multi-variable cases:
[Equation]
TC has also proven effective to enhance machine learning models in many tasks, such as independent component analysis [3], and disentangled representation learning [5, 19, 17].
However, TC suffers from the same numerical problem as MI: the exact values of TC are difﬁcult to calculate without the closed-form distribution p ( x i ) and with only samples accessible.
Previous works on disentangled representation learning [5, 10] avoid the estimation problem by assuming that both the latent priors and the inference posteriors follow multi-variate Gaussian distributions.
Poole, et al.
[22] proposed an upper bound of TC by further introducing another variable y .
With a strong assumption that given y , all x i | y are independent, p ( X | y ) =
(cid:81) ni
=1 p ( x i | y ) , Poole, et al.
[22] concluded that T C ( X )
= (cid:80) ni
=
I ( x i ; y ) − I ( X ; y ) .
All the aforementioned methods require additional assumptions to the distributions, which limits their application scenarios.
In this paper, we propose two TC estimation strategies based on mutual information variational bounds.
More speciﬁcally, we decompose TC into the summation of MI terms along two different calculation paths: the tree-like path and the line-like path.
Then the TC values are approximated by applying MI estimation to each decomposed term.
In our experiments, we test the performance of the proposed TC estimators under multivariate Gaussian simulations.
Method
With the deﬁnition of total correlation (TC) and mutual information (MI) in (2) and (1), we ﬁnd a connection between TC and MI summarized in Theorem 2.1.
Theorem 2.1.
Suppose A
= { i 1 , i 2 , . .
, i m } ⊆ { 1 , 2 , . . .
, n } is an index subset.
¯ A = { j : j / ∈ A} is the complementary set of A .
Denote X A =
( x i 1 , x i 2 , . . .
, x i m ) as the selected variables from X with the indexes A .
Then we have T C ( X )
= T C
( X A )
+ T C
( X ¯ A )
+ I ( X A ; X ¯ A ) .
Corollary 2.1.1. Given a variable group X and another y , T C ( X ∪ { y } ) = T C ( X )
+ I ( X ; y ) .
Corollary 2.1.2.
Given X
=
( x 1 , x 2 , . . .
, x n ) , we have T C ( X )
= (cid:80) n − 1 i
{+=
+}
I ( X 1: i ; x i +1 ) .
The Theorem 2.1 provides insight that the TC of a group of variables X can be decomposed into the TC of two subgroups X A and X ¯ A and the MI between the two subgroups.
Therefore, we can recursively represent the TC with MI terms.
More speciﬁcally, we propose two schemes with different structures to calculate TC with different MI terms (as shown in Figure 1).
Let X i :
j = ( x i , x i +1 , . . .
, x j ) denote a subset of variables with indexes from i to j .
Based on Theorem 2.1, we propose two recursive TC calculation schemes: (1) Line-like : T C ( X 1: i +1 ) =
T C ( X 1: i )
+ I ( X 1: i ; x i +1 ) ;
(2) Tree-like : T C
( X i : j ) =
T C ( X i : (cid:98) ( i + j ) / 2 (cid:99) )
+ T C ( X (cid:98) ( i + j ) / 2 (cid:99) +1: j ) +
I ( X i : (cid:98) ( i + j ) / 2 (cid:99) ; X (cid:98) ( i + j ) / 2 (cid:99) +1: j ) , where (cid:98) t (cid:99) indicates the largest integer smaller than t .
The line-like dynamic calculates the MI between a subgroup and a single variable, which leads to the representation of TC as the summation in Corollary 2.1.2.
The tree-like dynamic divides the variables into balanced subgroups, so that the MI between two subgroups can be calculated with two variable parts in similar dimensions.
Since the tree-like estimation is hard to summarize in an equation, we describe it in
Algorithm 1.
[-=
 I ( X 1: i ; x i +1 ) .-] With the total correlation being decomposed into summation of MI terms, we can derive total correlation estimators based on the previous mutual information variational bounds. Algorithm 1 Tree-like TC estimation algorithm
Prerequisite: MI estimation method ˆ I , samples { X ( i ) } Mi =
= { ( x ( i ) 1 , x ( i ) 2 , . . .
, x ( i ) n ) } Mi =
Function TC Tree -estimate( X i : j ) : if j − i ≤ 0 then returnelse m = (cid:98) ( i + j ) / 2 (cid:99) return
TC Tree -estimate( X i : m ) + TC Tree -estimate
( X m +1: j )
+ ˆ I ( X i : m ; X m +1: j ) end if
[Figure] [Figure]
Experiments
We derive our TC estimators based on four MI bounds (MINE [2], NWJ [20], InfoNCE [21], and
CLUB [8]) as TC-MINE, TC-NWJ, TC-InfoNCE, and TC-CLUB.
The detailed description and implementation to the four MI estimators are shown in the Supplementary Material.
Then we test the TC estimators with both tree-like and line-like strategies on simulations.
The simulation data are drawn from four-dimensional Gaussian distributions ( x 1 , x 2 , x 3 , x 4 ) ∼ N ( 0 , Σ ) , where Σ is a covariance matrix with all diagonal elements equal to 1 .
With this Gaussian assumption, the true TC value can be calculated as T C ( x 1 , x 2 , x 3 , x 4 ) =
− 12 log Det ( Σ ) , where Det ( Σ ) is the determinant of Σ .
Therefore, we can adjust the correlation coefﬁcients in Σ to set the ground-truth TC values in the range { 2 . 0 , 4 . 0 , 6 . 0 , 8 . 0 , 10 . 0 } .
At each TC true value, we sample data batches 4000 times, with batch size equal to 64, for the training of variational TC estimators.
In Figure 2, we report the performance of our TC estimators with different MI bounds at each training steps.
In each ﬁgure, the true TC value is shown as a step function with black line.
The estimation values are displayed among different steps with shadow blue curves.
The dark blue curves shows the local averages of estimated TC, with a bandwidth equal to 200.
Under both a tree-like and line-like path calculation, the TC-MINE, TC-NWJ and TC-InfoNCE remains a lower bound of the truth TC values, based on the fact that MINE, NWJ, and InfoNCE are lower bound of mutual information.
CLUB is an MI upper bound, while the TC-CLUB also behaves as an upper bound of total correlation.
The upper bound method TC-CLUB achieves better performance with line-like calculation.
This is because that CLUB requires a variational approximation q θ ( v | u ) when estimating I ( v ; u ) .
When we use line-like calculation path, v = x
i +1 is always a single variable, and u = X
: i is the concatenation of ( x 1 , . . .
, x i ) .
The q θ ( v | u ) as a neural network can have better performance with output v in a ﬁxed low dimension.
In contrast, the lower bound methods show better estimation with tree-like calculation than line-like calculation.
Because for all listed lower bound methods, the estimation of I ( v ; u ) is based on v and u equally.
With the tree-like strategy, each time the MI estimators are provided with samples in similar dimensions, which facilitates the learning of lower bound MI estimators.
The bias and variance of the TC estimators are shown in the Supplementary Material.
Discussion
We have derived the line-like and tree-like calculation strategies to decompose the total correlation into summation of mutual information.
By estimating mutual information terms with MI bounds, we introduced several TC estimators.
The tree-like and line-like calculation strategies can bring advantages to TC estimation depending on different MI estimation processes.
The proposed TC estimators can be further applied as learning criterion on many deep learning tasks, such as disentangled representation learning, ensemble learning, and model distillation.
A Proofs
Proof of Theorem 2.1.
Note that X A :
=
( x i 1 , x i 2 , . . .
, x i m ) and
X ˆ A = X / X A .
Denote X ˆ A =
( x j 1 , x j 2 , . .
, x j l ) .
Then
[Equation]
Proof of Corollary 2.1.2.
We denote X i : j :=
( x i , x i +1 , . . .
, x j − 1 , x j ) .
Note that
T C ( X 1: n ) =
E p ( x 1 , x 2 ,..., x n )
[Equation]
[Equation]
I ( x 1 , x 2 , . . .
, x n − 1 ; x n ) + TC ( X 1: n − 1 )
=
I ( X 1: n − 1 ; x n ) + T C ( X 1: n − 1 )
Similarly
[Figure] T C ( X 1: n ) =
I ( X 1: n − 1 ; x n ) + I ( X 1: n − 2 ; x n − 1 ) + T C ( X 1: n − 2 ) = n − 1 (cid:88) i =
I ( X 1: i ; x i +1 ) (3)
B Experiment Results
[Figure] [Figure]
C MI Estimators
The Mutual Information Neural Estimator (MINE) [2] is deﬁned as
[Equation]
where f ( · , · ) is a value function (or, a critic) approximated by a neural network.
The NWJ [20] lower bound is based on the f -divergence representation of MI:
[Equation]
The InfoNCE [21] lower bound is based on Noise Contrastive Estimation (NCE) [13]:
[Equation]
where the expectation is over N samples { x i , y i } Ni =1 drawn from the joint distribution p ( x , y ) .
The MI contrastive log-ratio upper bound (CLUB) estimator [8] is based on a parameterized distribution q θ ( y | x ) :
[Equation]
All the MI lower bounds require learning of a value function f
( x , y ) ; the CLUB upper bound requires learning of a network approximation q θ ( y | x ) .
To make fair comparison, we set the value function and the neural approximation with one hidden layer and the same hidden units.
For the multivariate
Gaussian setup, the number of hidden units is 20 .
On the top of hidden layer outputs, we add the
ReLU activation function.
The learning rate for all estimators is set to 1 × 10 − 4 .