Max-Sliced Mutual Information
Authors: Dor Tsur
Ziv Goldfeld
Kristjan Greenewald
Venue: NeurIPS 2023 poster (2023)
Keywords: CCA
dimensionality reduction
information theory
mutual information
neural estimation
slicing

Abstract: Quantifying dependence between high-dimensional random variables is central to statistical learning and inference. Two classical methods are canonical correlation analysis (CCA), which identifies maximally correlated projected versions of the original variables, and Shannon's mutual information, which is a universal dependence measure that also captures high-order dependencies. However, CCA only accounts for linear dependence, which may be insufficient for certain applications, while mutual information is often infeasible to compute/estimate in high dimensions. This work proposes a middle ground in the form of a scalable information-theoretic generalization of CCA, termed max-sliced mutual information (mSMI). mSMI equals the maximal mutual information between low-dimensional projections of the high-dimensional variables, which reduces back to CCA in the Gaussian case. It enjoys the best of both worlds: capturing intricate dependencies in the data while being amenable to fast computation and scalable estimation from samples. We show that mSMI retains favorable structural properties of Shannon's mutual information, like variational forms and identification of independence. We then study statistical estimation of mSMI, propose an efficiently computable neural estimator, and couple it with formal non-asymptotic error bounds. We present experiments that demonstrate the utility of mSMI for several tasks, encompassing independence testing, multi-view representation learning, algorithmic fairness, and generative modeling. We observe that mSMI consistently outperforms competing methods with little-to-no computational overhead.

Max-Sliced Mutual Information
Dor Tsur Ben-Gurion University
Ziv Goldfeld
Kristjan Greenewald MIT-IBM Watson Lab
Cornell University
Abstract
Quantifying dependence between high sige--imnnsional random variables is central to statistical learning and inference. Two classical methods are canonical correlation analysis (CCA), which identifies maximally correlated projected versions of the original variables, and Shannon's mutual information, which is universal depen- dence measure that also captures high-order dependencies. However, CCA only accounts for linear dependence, which may be insufficient for certain applications, while mutual information is often infeasible to compute/estimate in high dimen- sions. This work proposes middle ground in the form of scalable information- theoretic generalization of CCA, termed max-sliced mutual information (mSMI). mSMI equals the maximal mutual information between www-dimensional projec- tions of the high -igh-iinnssonal variables, which reduces back to CCA in the Gaussian case. It enjoys the best of both worlds: capturing intricate dependencies in the data while being amenable to fast computation and scalable estimation from samples. We show that mSMI retains favorable structural properties of Shannon's 's mutual information, like variational forms and identification of independence, We then study statistical estimation of mSMI, propose an efficiently computable neural estimator, and couple it with formal non-asymptotic error bounds. We present experiments that demonstrate the utility of mSMI for several tasks, encompassing independence testing, multi-view representation learning, algorithmic fairness, and generative modeling. We observe that mSMI consistently outperforms competing methods with little-to-no computational overhead.
Introduction
Dependence measures between random variables are fundamental in statistics and machine learning for tasks spanning independence testing [1-3], clustering [4, 5], representation learning [6, 7]. and self-supervised learning [8-10]. There myriad of measures quantifying different notions of dependence, with varying statistical and computational complexities. The simplest is the Pearson cor- relation coefficient [11], which only captures linear dependencies. At the other extreme is Shannon's mutual information [12], which is universal dependence measure that is able to identify arbitrar- ily intricate dependencies. Despite mutual information from data is infeasible in high-dimensional settings. First, mutual information estimation rates suffers from the curse of dimensionality, whereby convergence rates deteriorate exponentially with dimension [13]. Additionally, computing mutual information requires integrating a log-likelihood ratio over high-dimensional space, which is generally intractable.
universality and favorable properties, accurately estimating
Between these two extremes is the popular canonical correlation analysis (CCA) [14], which identifies maximally correlated linear projections of variables. Nevertheless, classical CCA still only captures linear dependence, which has inspired nonlinear extensions such as Hirschfeld- Sereeld-Geeeein--e -ebelein-Rényi (HGR) maximum correlation [15-17], kernel CCA [18, 19], deep CCA [20. 7], and various other generalizations [21-24]. However, HGR is computationally infeasible, while kernel and deep CCA can be burdensome in high dimensions, as they require optimization over reproducing kernel Hilbert spaces or deep neural networks, respectively. To overcome these shortcomings, this work proposes
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

max-sliced mutual information (mSMI) a scalable information theorrtti extension of CCA that captures the full dependence structure while only requiring optimization over linear projections.
1.1
Contributions
The mSMI is defined as the maximal mutual information between linear projections of the variables. Namely, the k-dimensional mSMI between X and Y with values in Rd and R respectively, is'
SI (X; =
sup (A.B))s
I(AT;
where St(k,, d) is the Stiefel manifold of d x l matrices with orthonormal columns. Unlike the nonlinear CCA variants that use nonlinear feature extractors in the high-dimensional ambient spaces, mSMI retains the linear projections of CCA and captures nonlinear structures in the low-dimensional feature space. This is done by using the mutual information between the projected variables, rather than correlation, the optimization objectíve. Beyond being considerably simpler from a computational standpoint, this crucial difference allows mSMI to identify the full dependence structure, akin to classical mutual information. mSMI can also be viewed as the maximized version of the average-sliced mutual information (aSMI) [25, 26], which averages I(ATX; X; BTY) with respect to (w.r.t) the Haar measure over Sk d')) x St(k: dy). However, we demonstrate that compared to aSMI, mSMI benefits from improved neural estimation error bounds and clearer interpretation.
We show that mSMI inherits important properties of mutual information, including identification of independence, tensorization, and variational forms. For jointly Gaussian (X, Y), the optimal mSMI projections coincide with those of k-dimensional CCA [27]. posing mSMI natural information- theoretic generalization. Beyond the Gaussian case, the solutions differ and mSMI may yield more effective representations for downstream tasks due the intricate dependencies captured by mutual information. We demonstrate this superiority empirically for multi-view representation learning.
For efficient computation, propose an mSMI neural estimator based on the Donsker- Varadhan (DV) variational form [28]. Neural estimators have seen a surge in interest due to their scalability and compatibility with gradient- optimization [29-36]. Our estimator employs single model that composes the projections with the neural network approximation of the DV critic, and then jointly optimizes them. This results in both the estimated mSMI value and the optimal projectionn matrices. Building on recent analysis of neural estimation of f-divergences [37, 38), we establish non-asymptotic error bounds that scale as (()2/((E-122 + kn -/2)), where E and n are the numbers of neurons and (X, Y) samples, respectively. Equating and n results in the (minimax optimal) parametric estimation rate, which highlights the scalability of mSMI and its compatibility to modern learning settings.
In our empirical investigation, we first demonstrate that our MSMI neural estimator converges orders of magnitude faster than that of aSMI [26]. This is because the latter requires (parallel) training of many neural estimators corresponding to different projection directions, while the mSMI estimator optimizes a single combined model. Notwithstanding the reduction in computational overhead, we show that mSMI outperforms average-slicing for independence testing. Next, we compare mSMI with deep CCA [20, 7] by examining downstream classification accuracy based on representations obtained from both methods in multi- learning setting. Remarkably, we observe that even the linear mSMI projections outperform nonlinear representations obtained from deep CCA. We also consider an application to algorithmic fairness under the infomin framework [39]. Replacing their generalized Pearson correlation objective with mSMI, we again observe superior performance in the form of more fair representations whose utility remains on par with the faimess-agnosti model. Lastly, we devise a max-sliced version of the InfoGAN by replacing the classic mutual information regularizer with its max-sliced analog. We show that despite the -ww-imensioonll projections, the max-sliced InfoGAN successfully learns to disentangle the latent space and generates quality samples.
2 Background and Preliminaries
Notation. For a, b E R. we use the notation a A = min(a, b} and V max(a, b}. For d 2 1, || || is the Euclidean norm in Rdd The Stiefel manifold of d x k matrices with orthonormal columns
'The parameter & is fixed and small compared the ambient dimensions dr, dy. often simply set ks L.

---

is denoted by St(k, d). For a d x k matrix A, we use p^ R + IR* for the orthogonal projection onto the row space of A. For A E Rdxk with rank(A) = S k A d, we write 01 (A)..... a,(A) for its non-zero singular values, and assume without loss of generality (w.l.o.g.) that they are arranged in descending order. Similarly, the eigenvalues of a square matrix y E Rdxd are denoted by (2)).... Aa(2). Let P(R4) denote the space of Borel probability measures on R. For u, P( we rl @ to denote product measure, while spt(u) designates the support of u.. All random variables throughout are assumed to be continuous w.r.t the Lebesgue measure. For measurable map f. the pushforward of j under is denoted by fap p f i.e., if X l then f(X) fap. For a jointly distributed pair Y) jxYY P(R BR,), write Ex and Exy for covariance matrix of X and cross-covariance matrix of (X, Y), respectively.
Canonical correlation analysis. CCA is classical method for devising maximally correlated linear projections of a pair of random variables (X, Y) ~ PXY E P(Rd x IR,,) via [14]
Exxx
(OCCA; (CCA) =
argmax (0.0)RR XR VOTEXX00 Eyro e Exx== d'E Eyø=l
argmax (0,0)ER xR
(1)
where the former objective is the correlation coefficient p(OT X, o Y) between the projected variables and the equality follows from invariance of to scaling, The global optimum has an analytic form as
1/=1/2b), vector pair associated with the largest singular value of Txy = xx' Ey E gd,xdd, This solution is efficiently computable in O((dz V dy)²) time, given that the population correlation matrices are known. CCA extends to k-dimensional projections via the optimization [27]
(OCCA, (CCA) (Ex* 1/2 !20,, E-" 1/2 d1).. where (01, ②h1 is the (unit-length) top left- and right-singular
(r(A ((ATExy)),
(A, X A A
with the optimal CCA matrices being ( (Acca, BECA) (Ex U2, Er1/'Va,, where U. and Va are the matrices of the first & left- and right-singular vectors of Txy. The optimal objective value then becomes the sum of the top k singular values of Txy (namely, its Ky Fan k-norm).
Divergences and information measures. Let E ((RR)) satisfy μ < v, i.e., p is absolutely con- tinuous w.r.t. v. The Kullback-Leibler (KL) divergence is defined as D(ullv) = S log(du/ du ddu, We have ((ul[ 2 O, with equality if and only if (iff) μ = v. Mutual information and differential entropy are defined from the KI. divergence as follows. Let (X, ~ Pxy E P(Rd= Rd.)) and denote the corresponding marginal distributions by μх and py. The mutual information between X and Y given by I(X; Y)= == D(μxy|llx @ py) and serves as a measure of dependence between those random variables. The differential entropy of X is defined as h(X) = h(ux) == -D(ux [[Leb). Mutual information between (jointly) continuous variables and differential entropy are related via Y(;:Y) h(X) + h(Y) h(X, Y); decompositions in terms of conditional entropies are also available [40].
3
Max-Sliced Mutual Information
We now define the k-dimensional mSMI, establish structural properties thereof, and explore the Gaussian setting and its connections to CCA. We focus here on the case of (linear) k-dimensional projections and discuss extensions nonlinear slicing in Section 3.3.
Definition 1 (Max-sliced mutual information). For 5 k S dn A dy the k-dimensional mSMI between (X,Y) μxx E P(R x RRd)) is
Slx(X;Y)= =
sup (A.B))SS(k.))xxStt.. I(ATX;BYY);
where Stikk, d) is the Stiefel manifold of d x l matrices with orthonormal columns.
The mSMI measures Shannon's mutual information between the most informative k-dimensional projections of X and Y. It can be viewed as maximized version of the aSMI (I; Y) from
3

---

[25, 26], defined as the integral of I(ATX; B w.r.t the Haar measure over St(k, dz) x St(k, dy). For d dæ dy. we have YY) = Y))
I(X;Y) due to invariance of mutual information to bijections. The supremum in mSMI is achieved since the Stiefel manifold is compact and the function (A, B3) + I(ATX; BTY) is Lipschitz and thus continuous (Lemma 2 of [26]).
Remark (Multivariate and conditional mSMI). The mSMI definition above extends to the mul- tivariate and conditional cases follows. Let (X, Z) PXYZ P x Rd x Rd)). The I t-dimensional multivariate and conditional mSMI functionals are, respectively, (I, = maxA,B.C C. ([ATX,BTY; CT Z) and (X; == maxA,B,C (ATX; between SI,(X; and its maltivariate and conditional versions are given in the proposition to follow. We also note that one may generalize the definition S1 (X: Y) to allow for projections into feature spaces of different dimensions, i.e., A E St(kg,dz) and B E St. dy) ) for kg # by We expect our theory to extend that case, but leave further exploration for future work.
Connections
In the spinit of mSMI, we define the max-sliced differential entropy.
Definition 2 (Max-sliced entropy). The k-dimensional max-sliced (differential) entropy of X μx E P((R is shg (X) = =hy (1) == SUPAESt(K,4) h(ATX).
An important property of classical differential entropy is the maximum entropy principle [40], which finds the highest entropy distribution within given class. In Appendix B, study the max-sliced entropy maximizing distribution in several common scenarios. For instance, we show that shş is maximized by the Gaussian distribution under a fixed (mean and) covariance constraint. Namely, letting P1 (m, Σ) = {u E P(Rd) spt(μ) = Rd [x m, E,, [ m) m)) E}.
we have argmax (n. sh& (14) = (m, E). An intimate connection between max-sliced entropy and PCA is established in the sequel, under the Gaussian setting.
Remark 2 (Sliced divergences). The slicing technique has originated as a means to address scalabil- ity issues concerning statistical divergences. Significant attention was devoted to sliced Wasserstein distances as discrepancy measures between probability distributions [41-47]. As such, the sliced Wasserstein distance differs from mutual information and its sliced variants, which quantify de- pendence between random variables, rather than discrepancy per se. Additionally, as Wasserstein distances are rooted in optimal transport theory, they heavily depend on the geometry of the underly- ing data space. Mutual information, on the other hand, is induced by the KL divergence, which only depends on the tog-likelihood of the considered distributions and overlooks geometry.
3.1 Structural Properties
The following proposition lists useful properties of the mSMI, which are similar to those of the average-sliced variant (cf. [26, Proposition 1]) as well as Shannon's 's mutual information itself.
Proposition 1 (Structural properties). The followirg properties hold:
1. Bounds:
For any integers k1 < k2 Sle, (X;Y)= S Slx, (X; Y) s She, (X;Y) & I(X;Y).
Identification of independence: SI (X;Y) N O with equality iff (X,Y) Y) are independent.
3. KL divergence representation: We have
SIx(X;Y)= (A.B)
sup
((p ,p",, ) (p", pv),
Sub-chain rule: For any random variables X Xn, Y, we have
(lx(X): + Σ (X,; X X1-1).
Tensorization: For mutually independent {(X, (X),, SIx ((X,) (Y. = (X;Y.)
The proof follows by similar arguments to those in the average-sliced case, but is given for com- pleteness Supplement A. Of particular importance are Properties and 3. The former renders MSMI sufficient for independence testing despite being significantly less complex than the classical mutual information between the high-dimensional variables. The latter, which represent mSMI as supremized KL divergence, is the basis for neural estimation techniques explored in Section 4.

---

Remark 3 (Relation to average-SMa. SMI). Beyond the inequality relationship in Property above, Propo- sition in [25] (paraphrased) shows that for matrices W a. W and vectors by. by of appropriate dimensions, we have su.... by 511 S + by: WTY TY + by) = SI (X;Y) Y), and the relation readily extends to projection dimension k 1. words, optimizing the uSMI over linear transforma- tions of the high-dimensional data vectors coincides with the max-sliced version. This further justifies the interpretation of (Xx((XYY) the information between the two most informative representations of X, in a k-dimensional feature space. also suggests that mSMI compatible for feature extraction tasks, as explored in Section 5.3 ahead.
3.2 Gaussian Max-SMI versus CCA
The mSMI is an information thhoreticc extension of the CCA coefficient PCCA Y), which is able to capture higher order dependencies. Interestingly, when (X, Y) are jointly Gaussian, the two notions coincide. We next state this relation and provide a closed-form expression for the Gaussian mSMI.
Proposition (Gaussian mSMI). Let X ~ N(mx, Ex) and Y ~ [(my, Σy) be dx- and dy- dimensional jointly Gaussian vectors with nonsingular covariance matrices and cross-covariance Σxy. For any k S dz A dy. we have
1 Σ Slx(X;Y) X B(Ca --
- log (1-o;(Txv)3)
En1/2 E Rd.xd, and where (AcCa, BECA) are the CCA solutions from (2), Txv = E= Exy on (Txx)) S ... S 0i (TxT) S 1 are the top singular values of Ty (ordered).
xv
This proposition is proven in Supplement A.2. We first show that the optimization domain of (lx(X;Y) can be switched from the product of Stiefel manifolds to the space of all matrices subject unit variance constraint (akin (2)), without changing the mSMI value. This implies that the CCA solutions (ACCA, BCCA) from (2) are feasible for mSMI and we establish their optimality using a generalization of the Poincaré separation theorem [48, Theorem 2.2]. Specializing Proposition to oie-dimensional projections, i.e., when h = 1, the mSMI is given in terms of the canonical correlation coefficient PecA (X, Y)= == ())().)) XR n(OT X, Φ Y). Namely,
0.5 log (1 (X, where (OCCA, dCcA)) are the global optimizers of pCca (Y,Y).
(X; 1
Remark 4 (Beyond Gaussian data). While the mSMI solution coincides with that of CCA in the Gaussian case, this is no longer expected to hold for non Gaussian distributions. CCA is designed maximize correlation, while mSMI has Shannon's 's mutual information between the projected variables as the optimization objective. Unlike correlation, mutual information captures higher order dependencies between the variables, and hence the optimal mSMI matrices will not generally coincide with (ACCA, BCCA). Furthermore, the intricate dependencies captured by mutual information suggest that the optimal mSMI projections may yield representations that are more effective for downstream tasks. We empirically verify this observation in Section 5 on several tasks, including classification, multi-veww representation learning, and algorithmic fairness.
Similarly to the above, the Gaussian max-sliced entropy is related to PCA [49, 14]. In Supplement A.3, we prove the following.
Proposition 3 (Gaussian max-sliced entropy). For a d-dimensional Gaussian variable X ~ (mm, E). we have sha (X) sUPAeStck.d) h(ATX) X) = 0.5 Dt log (2me);(E)) where APCA is optimal PCA matrix and X1 (I),.... (E) are the top k eigenvatues of Σ.
Note that the eigenvalues > (E)..... (E) are non-negative since Σ is a covariance matrix. Ex- trapolating beyond the Gaussian case, this poses max-sliced entropy as generalization of PCA for unsupervised dimensionality reduction. An analogous extension using the Rényi entropy of order 2 was previously considered in [50] for the purpose of binary classification. In that regard, (X) can be viewed as the a- aéRényi variant when + 1.
coommttion-thhoretin
5

---

3.3 Generalizations Beyond Linear Slicing
The notion of mSMI readily generalizes beyond linear slicing. Fix d dy 2 1, k S d, A dy and consider two (nonempty) function classes 9 S {g Rd== -> R and S {h: Rd - R.
Definition 3 (Generalized mSMI). The generalized mSMI between (X, Y) ~ pXY E P(R x IR W.E. the classes G and H. is SIg, (Xx Y) == ())(G, ((((X):(()))
The generalized variant reduces back to (X; by taking g = Sproj == {p* A E St dz)) and H. = H proj == {p": : B E St(k, dy)); but otherwise allows more flexibility in the way (X, are mapped into RR We also have that if G C g' and H C H', then HIg. Y) - Slg H (X:Y) S ((; which corresponds Property from Proposition 1. Further observations are follows.
Proposition 4 (Properties). For any classes 9, H. we have that SIG. h always satisfies Properties 3-5 from Proposition 1. further Sproj C 9 and Hproj S H. then SIg. also satisfies Property 2.
We omit the proof as it follows by the same argument as Proposition 1, up to replacing the linear projections with the functions (g.h) 9 x H. In practice, the classes G and H. chosen to be parametric, typically realized by artificial neural networks. As discussed in Remark ahead, this is well-suited to the neural estimation framework for mSMI (both standard and generalized), Lastly, note that (l (Y; corresponds to the objective of multi-view representation learning [5 1], which considers the maximization of the mutual information between NN-based representation of the considered variables. We further investigate this relation in Section 5.3.
Neural Estimation of Max-SMI
We study estimation of mSMI from data, seeking an efficiently computable and scalable approach subject to formal performance guarantees. Towards that end, observe that the mSMI is compatible with neural estimation [29, 38] due to its convenient variational form. In what follows we derive the neural estimator, describe the algorithm to compute it, and provide amm-syyptotiic error bounds.
4.1 Estimator and Algorithm
Fix dz, dy 2 1, k S de dy. and μxy E P(R x Rdx); suppress k. dg, dy from notation of the considered function classes. Neural estimation is based on the DV variational form:
I(X:Y) = sup Lov(f; pxy), Lov(f: pxx) = Eif(X, log (c fEF
where (X,Y) pxY, (X, px & Hy, and F is the class of all measurable functions f dd x Rd + R (often referred to as DV potentials) for which the expectations above are finite. As MSMI the maximal mutual information between projections of X, Y. we have
sup (); (.B))SS((.d,)xxS(((... feJ
sup Lov (J: (pA (.A,B)))xx) = sup (Dv(Ji FEJPRI
where F = {f (p^, : f F, (A, E St(k, da) x st((k,d)) The RHS above is given by optimizing the DV objective LDV over the composed class J.Proj, which first projects (X,Y)= - (ATX,BT) and then applies DV potential R* RR R to the projected variables.
Neural estimator. Neural estimators parametrize the DV potential by neural nets, approximate expectations by sample means, and optimize the resulting empirical objective over parameter space. Let Fun be class of feedforward networks with input space R* IR& and real-valued outputs." Given i...d. samples (X1, Yi),-... .,,,...,. (X,, Yn) from pxy, we first generate negative samples (i.e., from xx & py) by taking (X1, Yo(): ........................... (Xn, Yo))), where E S, a permutation such that
One may instead use the form that from convex duality: I(U; = sup, V[)
For now, we leave the architecture (number of layers/neurons, parameter bounds, nonlinearity) implicit to allow flexibility of implementation; will specialize to concrete class when providing theoretical guarantees.

---

oii) * i, for all i i= I....,n. The neural estimator of (X;Y) is now given by
51k (x", ,V")= = sup
Σ I(X, - log
(5)
JEFm
where F.roj == (p".p".pp)) : f Fan: (A, B) E St(kk dz) x St dy)] } is the composition of the neural network class with the projection maps. The projection maps can be absorbed into the network architecture as first linear layer that maps the (dz + dy) "ddmenssioo input to 2k-dimensional feature vector, which is then further processed by the original f E Fm network. Note that projection onto the Stiefel manifold can be efficiently and differentiably computed via QR decomposition. Hence, the Stiefel manifold constraint can be easily enforced by setting A, B to be the projections of unconstrained d x k matrices. Further details on the implementation are given in Supplement C.
Remark 5 (Nonlinear slicing). For learning tasks that may need more expressive representations (X, Y), one may employ the nonlinear mSMI variant from Section 3.3. In practice, the classes g= { and H { are taken to be parametric, realized by neural networks. These networks naturally compose with the DV critic fox, resulting in a single compound model fō o (go, ho).
4.2 Performance Guarantees
Neural estimation involves three sources of error: (i) function approximation of the DV potential; (ii) empirical estimation of the means; and (iii) optimization, which comes from employing suboptimal (e.g., gradient-based) routines. Our analysis provides sharp non-asymptotic bounds for errors of type (i) and (ii), leaving the account of the optimization error for future work. We focus on a class of f-neuron shallow ReLU networks, although the ideas extend to other nonlinearities and deep
architectures. Define F as the class of all f R I + R. f(2) Ei= 3,0 (ww,, 2) + b,) (wo. 2) + bo. whose parameters satisfy maxicice 11 11 V 16, S. 1. max sise [3.1 S 2E and
[bol, ||wll| I ag, where ((2) = V the ReLU activation and ag log log E V
Consider the neural mSMI estimator S1 == "i (X", y") (see (5)). We provide convergence rates for it over an appropriate distribution class, drawing upon the results of [37] for neural estimation of f-divergences. For compact X' c. Rd and y c Rdy, let Pac x J)) be the set of all Lebesgue absolutely continuous joint distribution μxx with (px() C X x Y. Denote the Lebesgue density of pxy by fxy. The distribution class of interest is*
Fm
Pe(M, = {IXV ( )))
Br C2 (2++ for some open set 41 3 X' xX)) )7
(6)
s.t. log fxy = rlxxy, I(X;Y) S M
which, in particular, contains distributions whose densities are bounded from above and below on X' x y' with smooth extension to an open set covering X' x y. This includes uniform distributions, truncated Gaussians, truncated Cauchy distributions, etc. The following theorem provides the convergence rate for the mSMI neural estimator, uniformly over b).
Theorem (Neural estimation error). For any M, b 2 O, we have
E ())x(x;Y)-Si?) < C (e +k +))
sup EP
where C' depends on M, b, k, and the radius of the ambient space ||'' = suP( [| v)ll-
The theorem is proven in Supplement A.4 by adapting the error bound from [38, Proposition 2] to hold for I(AT;; BTY). uniformly over (A, B) E St(k, d,) x St(k, d,y. To that end, we show that for any μXY E (b, M), the log-density of (ATX, (AY) ~ (p (A.pB) ))))) admits an extension (to an open set containing the support) with k: + continuous and uniformly bounded derivatives.
Remark 6 (Parametric rate and optimality). Taking i x n, the resulting rate in Theorem I is parametric, and hence minimax optimal. This result implicitly assumes that M is known when picking the neural net parameters. This assumption can be relaxed to mere existence of (an unknown) M, resulting in an extra polylog(E) factor multiplying the n 1/2 term.
*Here, C (24) = (f E C² (L) max.:|lollss fl S b), where D", the partial derivative operator of order Dd on. The restriction of f: Rd R to X C Rd flx.
oa)) E 2
7

---

Remark 7 (Comparison to average-SMII, Neural estimation of classic mutual information under the framework of [38] requires the density to have Hölder smoothness 2 [(d,- d,)/2] + For ( smoothness of + 3 is sufficient (even though the ambient dimension is the same), which means it can be estimated over larger distribution class. Similar gains in terms of smoothness levels were observed for aSMI in [26]. Nevertheless, we note that mSMI is more compatible with neural estimation than peeaagesslicing [25. 26]. The mSMI neural estimator integrates the max-slicing into the neural network architecture and optimizes a single objective. The aSMI neural estimator from [26] requires additional Monte Carlo integration step approximate the integral over the Steifel manifolds. This results in an extra k2/22m=-12 -1/2 term in the error bound, where m. is the number of Monte Carlo samples, introducing a burdensome computational overhead (see Section 5.1).
Remark 8 (Non-ReLU networks). Theorem employs the neural estimation bound from [38], which relies on [52] to control the approximation error. As noted [38], their bound extends to any other sigmoidal bounded activation with lim ( O and limg-+oc o(z) by appealing to the approximation bound from [53] instead. Doing so would allow relaxing the smoothness requirement on the extension to r E C in (6), but at the expense of scaling the hidden layer parameters as e1/2 log (as opposed to the ReLU-based bound, where the parameter scale is independent of E).
Experiments
-10
5.1 Neural Estimation
2.5
[nais]
2.3
We compare the performance of neural estimation methods for mSMI and aSMI synthetic dataset of correlated Gaussians. Let X, Z ~ (O, 1) be i.i.d. and set Y pX + VI p² for p E (0, 1). The goal estimate the k-dimensional mSMI and aSMI between (X.Y). Y). We train our mSMI neural estimator and the aSMI neural estimator from [26, Section 4.2] based on i.i.d. samples, and compare their perfor- mance as a function of n. Both average and max- sliced algorithms converge at similar rates; however, aSMI has significantly higher time complexity due to the need train multiple neural estimators (one for each projection direction). This is shown in Figure I, where compare the average epoch time for each algorithm against the dataset size. Implementation details are given in Supplement C.
2.1
0.35
mnts)
0.25
0.15
Time [sec.]
10
10¹
10e
10
Figure 1: Neural estimation performance with = 0.5. Convergence vs. n in upper figures and average epoch time n in lower figure.
5.2 Independence Testing
In this experiment, we compare mSMI and aSMI for independence testing. We follow the setting from [26, Section 5], generating d-dimensional samples correlated in latent d' d'-dimensional subspace and estimating the information measure to determine dependence. We estimate the aSMI with the method from [26]. using m 1000 Monte Carlo samples and the Kozachenko-Leonenee mutual information between the projected variables [54]. We then compute AUC-ROC over 100 trials,
estimator for the
k k= 2
A
k
k
k:
0.9
0.8
0.7
0.6
0.5
10
10
10a
10
10a
10"
10
102 n
10a
10'
10
10a
10'
10
10a
10'
10
10"
70
n
73.
72.
n
Figure 2: ROC-AUC comparison. Dashed and solid lines show results for aSMI [26] and mSMI (ours). respectively. Blue plots correspond to (d, d') = (10,4). while red correspond to (d, d') = (20, 6).
8

---

considering various ambient and projected dimensions. For mSMI, as we cannot differentiate through the Kozachenko- Leoneenko estimator, we resort to gradient-free methods. We employ the LIPO algorithm from [55] with a stopping criterion of 1000 samples. This choice is motivated by the w.r.t. the Frobenius norm on St dz) St dy) (cf. [26, Lemma 2D). Figure 2 shows that when k > 2. mSMI captures independence better than aSMI, particularly in the lower sample regime. We hypothesize that this is due to the fact that the shared signal lies in low-dimensional subspace, which mSMI can isolate and perhaps better exploit than aSMI, which averages over all subspaces. When k is much smaller than the shared signal dimension d', mSMI fails to capture all the information and aSMI, which takes all slices into account, may be preferable. Results are averaged over 10 seeds. Further implementation details are in Supplement C.
Lipschitzness of (A, B) + I(ATX;
5.3 Vulti-View Representation Learning
Linear CCA
Linear mSMI
ML.P DCCA
MIP MSMI
+0.03
0.274-00.2
0.284-0003
0.291+00.2
We next explore MSMI as an information- theoretic generalization of CCA by examining its utility in multi-view representation learning- a popular CCA application. Without using class labels, we obtain mSMI-based k-dimensional representations of the top and bottom halves of MNIST images (considered as two separate views of the digit image). This is done by com- puting the k-dimensional mSMI between the views and using the maximizing projected variables as the representations. We compare to similarly obtained CCA-based -eaaedded representations, following the method of [20]. Both linear and nonlinear (parameterized by an MLP neural network) slicing models are optimized with similar initialization and data but different loss functions. Performance is evaluated via downstream 10-class classification accuracy, utilizing the learned top-half representations. Results are averaged over 10 seeds. As shown in Table 1, mSMI outperforms CCA for learning meaningful representations. Interestingly, linear representations learned by mSMI outperform nonlinear represen- tations from the CCA methodology, demonstrating the potency of mSML. Full implementation details and additional results are given in Supplements C and D, respectively.
0.32.10.02
0.346.0.02
0.314.10:3
0.417.10.02
0.42:40.01
0.478-0002
0.441=0.04
0.546+0001
0.553-10.03
0.666:0.01
0.645:10.02
0.665
0.614:1.022
0.751:0.01
0.697:0001
0.75311001
16
0.673-40.02
0.775:0001
0.730=0002
0.779:10001
0.704
0.79-00.006
0.774.1001
0.798.10.01
Table I: Downstream classification accuracy from
MNIST representations by CCA and mSMI.
The aSMI is not considered for this experiment since it does not provide a concrete latent space representation (as it is an averaged quantity). Moreover, if one were to maximize aSMI as an objective to derive such representations, this would simply lead back to computing mSMI; cf. Remark 3.
5.4 Learning Fair Representations
Another common applica- tion of dependence mea- sures is learning fair rep- resentations of data. We seek a data transformation Z = fXX that is useful for predicting some outcome or label Y. while being statis- tically independent of some sensitive attribute A (e.g., gender, race, or religion of the subject). In other words, fair representation is one that is not affected by the subjects' protected attributes that downstream predictions are not biased against protected groups, even if the training data may have been biased. Following the setup of [39], we measure utility and fairness using the HGR maximal correlation PHGR( :) = suph.g p(h(),, (,g(()). seeking large PHGR Z, Y) and small PHGR (Z, A) where h and g are parameterized by neural networks. As solving this minimax problem directly is difficult in practice, following [39] we learn Z by optimizing the bottleneck equation PHGR Y) B51, (Z, A), where we use a neural estimator for the mSMI and B3, k are hyperparameters.
Table 2: Learning a fair representation of the US Census Demographic dataset, following the setup of [39]. Results are shown as the median over 10 runs with random data splits. The fairest result is k 6.
N/A
Slice [39]
mSMM (ours)
2
k
2
R
3
k
PHOR | FT1
0.949
0.967
0.955
0.958
0.557
ONIRR (Z, A)
0.793
0.116
0.220
0.099
0.067
0.048
0.029
0.036
0.047
Table 2 shows results on the US Census Demographic dataset extracted from the 2015 American Community Survey, which has 37 features collected over 74,000 census tracts. Here Y is the fraction of children below the poverty line in tract, and A is the fraction of women in the tract. Following the same experimental setup as [39], the learned Z is 80-dimensional. As [39] showed that their "Slice" approach significantly outperformed all other baselines on this experiments under computational

---

8
8
P
F
6
6
O
0
9
4
9
4
4
4
3
3
3
3
2
2
2
2
2
7
2
2
2
2
5 5
6
5
5
5
5
5
3
3
3
3
3
3
3
9
9
9
9
4
4 / S
7
4
4
4
9
7
7
7 7
7
7
7
7
7
/
1
6 6
6
6
6
4
6
5 7
5
5
6
6
5
5
5
8
8
8
8
7
5
4
9
9
O 0
0
D
0
2 2
7
2
2
2
(a) FNSH
(b) MNIST, A = 5
(c) MNIST, = 10
Figure 3: MNIST images generated via the max-sliced InfoGAN.
constraint', we apply the same computational constraint to our approach and compare only to Slice and to the "N/A" aairnes--agnostc model trained on the bottleneck objective with 3 = 0. Note that for k >>>> 1, mSMI learns more fair representation Z (lower PHGR (Z, A)) than Slice, while retaining a utility PHGR Z, Y) on par with the fairness agnostic N/A model. We emphasize that due to the reasons outlined in Section 5.3, aSMI is suitable for the considered task and is thus not included in the comparison. Results on the Adult dataset are shown in Supplement E.
5.5 Max-Sliced InfoGAN
We present application of max-slicing to generative modeling under the InfoGAN framework [56]. The InfoGAN learns a disentangled latent space by maximizing the mutual information between latent code variable and the generated data. We revisit this architecture but replace the classical mutual information regularizer the InfoGAN objective with mSMI. Our max-sliced InfoGAN is tested on the MNIST and FaSHion-MNIST datasets. Figure 3 presents the generated samples for several projection dimensions. We consider 3 latent codes (C₁, C2. Ca). which automatically learn to encode different features of the data. We vary the values of C₁ which is 10-state discrete variable, along the column (and consider random values of (C2, Ca) along the rows). Evidently, C successfully disentangles the 10 class labels and the quality of generated samples is on par with past implementations [56, 26]. We stress that since mSMI relies on low-dimensional projections, the resulting InfoGAN mutual information estimator uses reduced number of parameters (at the negligible cost of optimizing over linear projections). Additional details are given in Supplement C.
6 Conclusion
This paper proposed mSMI, an information theoretic generalization of CCA. mSMI captures the full dependence structure between high dimensional random variables, while only requiring an optimized linear projection of the data. We showed that mSMI inherits important properties of Shannon's 's mutual information and that when the random variables Gaussian, the mSMI optimal solutions coincide with classic k-dimensional CCA. Moving beyond Gaussian distributions, we present neural estimator of mSMI and establish non-asymptotic error bounds.
Through several experiments we demonstrate the utility of mSMI for tasks spanning independence test- ing, multi-view representation learning, algorithmic faimess and generative modeling, showing outperforms popular methodologies. Possible future directions include an investigation of an opera- tional meaning of mSMI, either information theoretic physical terms, extension of the proposed formal guarantees to the nonlinear setting, and the extension of the neural estimation convergence guarantees to deeper networks. Additionally, mSMI can provide mathematical foundation mutual information-bast representation learning, a popular area of self-supervised learning [10,55 57].
In addition to the above, we plan develop rigorous theory for the choice of k which is currently devised empirically and is treated as a hyperparameter. When the support the distributions lies in some d' < d dimensional subspace, the choice of & = d' is sufficient to recover the classical mutual information, and therefore it characterizes the full dependence structure. Extrapolating from this point, conjecture that the optimal value of & is related to the intrinsic dimension of the data distribution, even when it is not strictly supported on a wiw-dimensional subset.
Runtime per iteration not exceed the runtime of Slice per iteration. We used NVIDIA V V10 GPU.
10

---

References
[1] Sidney Siegel. Nonparametric statistics. The American Statistician, 11(3):13-19, 1957.
[2] Larry D Haugh. Checking the independence of two covariance-stationary time series; univariate residual cross-correlation approach. Journal of the American Statistical Association. 71(354):378-3855 1976.
[3] Thomas B Berrett and Richard J Samworth. Nonparametric independence testing via mutual information. Biometrika, 106(3):547-566, 2019.
[4] Atul Butte and Isaac S Kohane. Mutual information relevance networks: functional genomic clustering using pairwise entropy measurements. Biocomputing 2000, pages 418-429, World Scientific, 1999.
[5] Alexander Kraskov, Harald Stögbauer, Ralph Andrzejak, and Peter Grassberger. Hierarchical clustering using mutual information. Europhysics Letters, 70(2):278, 2005.
[6] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint (i0n:00788474 2018.
[7] Weiran Wang, Raman Arora, Karen Livescu, and Jeff Bilmes. On deep multi-view representation learning. In International conference on machine learning, pages 1083-1092. PMLR, 2015.
[8] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. simple framework for contrastive learning of visual representations. In International conference machine learning, pages 1597-1607. PMLR, 2020.
[9] Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self- supervised learning via redundancy reduction. In International Conference on Machine Learn- ing, pages 12310-12320, PMLR, 2021.
[10] Randall Balestriero, Mark Ibrahim, Vlad Sobal, Ari Morcos, Shashank Shekhar, Tom Goldstein. Florian Bordes, Adrien Bardes, Gregoire Mialon, Yuandong Tian, et al. A cookbook of self-supervised learning. arXiv preprint arXiv:2304.122111 2023.
[11] Karl Pearson. Vii. note on regression and inheritance in the case of two parents. proceedings of the royal society of London, 58(347-352::240-244 1895.
[12] Claude Elwood Shannon. A mathematical theory of communication. The Bell system technical journal, 27(3):379 423, 1948.
[13] Liam Paninski. Estimation of entropy and mutual information. Neural Computation, 15:1191- 1253, June 2003.
[14] Harold Hotelling. Analysis a complex of statistical variables into principal components. Journal of educational psychology, 24(6):417, 1933.
[15] Hermann O Hirschfeld. connection between correlation and contingency. Mathematical Proceedings of the Cambridge Philosophical Society, (1(::5220 524444 1935.
[16] Hans Gebelein. Das statistische problem der korrelation als variations-und eigenwertproblem und sein zusammenhang mit der ausgleichsrechoung. ZAMM-Jourdn of Applied Mathematics and Mechamics/Zeitschrff für Angewandte Mathematik und Mechanik, 21(6))364-377 1941.
[17] Alfréd Rényi. On measures of dependence. Acta Mathematica Academiae Scientiarum Hungar- ica, 10(3-4):441443 1959.
[18] Shotaro Akaho. A kernel method for canonical correlation analysis. arXiv preprint cs/0609071, 2006.
[19] David R Hardoon, Sandor Szedmak, and John Shawe-Taylor. Canonical correlation analysis: overview with application learning methods. Neural computation, 16(12):2639-266 2004.

---

[20] Galen Andrew, Raman Arora, Jeff Bilmes, and Karen Livescu. Deep canonical correlation analysis. In International conference on machine learning, pages 1247-1255. PMLR, 2013.
[21] Francis R Bach and Michael Jordan. A probabilistic interpretation of canonical correlation analysis. Technical Report, 2005.
[22] Tae-Kyun Kim, Josef Kittler, and Roberto Cipolla. Learning discriminative canonical correla- tions for object recognition with image sets. In Computer Vision-ECCV 2006: 9th European Conference on Computer Vision, Graz, Austria, May 7-13, 2006, Proceedings, Part III 9, pages 251-2622 Springer, 2006.
[23] Elena Parkhomenko, David Tritchler, and Joseph Beyene. Sparse canonical correlation analysis with application to genomic data integration. Statistical applications in genetics and molecular biology, 8(1), 2009.
[24] Amichai Painsky, Meir Feder, and Naftali Tishby. Nonlinear canonical correlation analysis: A compressed representation approach. Entropy, 22(2):208, 2020.
[25] Ziv Goldfeld and Kristjan Greenewald. Sliced mutual information: A scalable measure of statistical dependence. Advances in Neural Information Processing Systems, 34:17567-17578, 2021.
[26] Ziv Goldfeld, Kristjan Greenewald, Theshani Nuradha, and Galen Reeves. k-sliced mutual information: A quantitative study of scalability with dimension. Advances in Neural Information Processing Systems, 35:15982-15999 2022.
[27] Kantilal Varichand Mardia, John T Kent, and John M Bibby. Multivariate analysis. Probability and mathematical statistics, 1979.
(28) Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain Markov process expectations for large time. iv. Communications on Pure and Applied Mathematics, 36(2):183-212, 1983.
[29] Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeshwar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and Devon Hjelm. Mutual information neural estimation. In International Conference on Machine Learning, pages 531-540. PMLR, 2018.
[30] Ben Poole, Sherjil Ozair, Aäron van den Oord, Alexander Alemi, and George Tucker. On variational lower bounds of mutual information. In NeurIPS Workshop Bayesian Deep Learning, 2018
[31] Chung Chan, Ali Al-Bashabsheh, Hing Pang Huang, Michael Lim, Da Sun Handason Tam, and Chao Zhao. Neural entropic estimation: faster path to mutual information estimation. arXiv preprint arXiv: 1905. 12957, 2019.
[32] Jiaming Song and Stefano Ermon. Understanding the limitations of variational mutual infor- mation estimators. Proceedings of the International Conference on Learning Representations (ICLR), 2020.
[33] Jingjing Zhang, Osvaldo Simeone, Zoran Cvetkovie, Eugenio Abela, and Mark Richardson. Itene: Intrinsic transfer entropy neural estimator. arXiv preprint arXiv:191..07277, 2019.
[34] Sudipto Mukherjee, Himanshu Asnani, and Sreeram Kannan. Cemi: Classifier based conditional mutual information estimation. In Uncertainty in artificial intelligence, pages 1083-1093. PMLR, 2020.
[35] Dor Tsur, Ziv Aharoni, Ziv Goldfeld, and Haim Permuter. Neural estimation and optimization of directed information over continuous spaces. IEEE Transactions Information Theory, 2023.
[36] Qing Guo, Junya Chen, Dong Wang, Yuewei Yang, Xinwei Deng, Jing Huang, Larry Carin, Fan Li, and Chenyang Tao. Tight mutual information estimation with contrastive fenchel-legendre optimization. Advances in Neural Information Processing Systems, 15:28319-2833 2022.
12

---

[37] Sreejith Sreekumar, Zhengxin Zhang, and Ziv Goldfeld. Non-asympiotic performance guaran- tees for neural estimation of f-divergences. In International Conference on Artificial Intelligence and Statistics, pages 3322-3330. PMLR, 2021.
[38] Sreejith Sreekumar and Ziv Goldfeld. Neural estimation of statistical divergences. Journal of Machine Learning Research, 23(126):1-75, 2022.
[39] Yanzhi Chen, Weihao Sun, Yingzhen Li, and Adrian Weller Scalable infomin learning. Ad- vances in Neural Information Processing Systems, 35:2226-222 2239, 2022.
[40] Thomas M Cover and A Joy Thomas. Klements of Information Theory. Wiley, Yor-York, 2nd edition, 2006.
[41] Julien Rabin, Gabriel Peyré, Julie Delon, and Mare Bernot. Wasserstein barycenter and its application to texture mixing. In Scale Space and Variational Methods in Computer Vision: Third International Conference, SSVM 2011, Ein-Gedi, Israel, May 29-June 2, 2011, Revised Selected Papers 3, pages 435-446. Springer, 2012.
[42] Ishan Deshpande, Yuan-Timg Ting Hu, Ruoyu Sun, Ayis Pyrros, Nasir Siddiqui. Sanmi Koyejo, Zhizhen Zhao, David Forsyth, and Alexander G Schwing. Max-sliced wasserstein distance and its use for gans. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10648-10656, 2019.
[43] Soheil Kolouri, Kimia Nadjahi, Umut Simsekli, Roland Badeau, and Gustavo Rohde. Gener- alized sliced wasserstein distances. Advances in neural information processing systems, 32, 2019.
[44] Tianyi Lin, Chenyou Fan, Nhat Ho, Marco Cuturi, and Michael Jordan. Projection robust wasserstein distance and riemannian optimization. Advances neural information processing systems, 33:9383- 9397, 2020.
[45] Minhui Huang, Shiqian Ma, and Lifeng Lai. A riemannian block coordinate descent method for computing the projection robust wasserstein distance. In International Conference on Machine Learning, pages 4446-4455. PMLR, 2021.
[46] Tianyi Lin, Zeyu Zheng, Elynn Chen, Marco Cuturi, and Michael I Jordan. On projection robust optimal transport: Sample complexity and model misspecification. In International Conference on Artificial Intelligence and Statistics, pages 262-270. PMLR, 2021.
[47] Sloan Nietert, Ziv Goldfeld, Ritwik Sadhu, and Kengo Kato. Statistical, robustness, and computational guarantees for sliced wasserstein distances. Advances in Neural Information Processing Systems, 15:28179--2192, 2022.
[48] Radhakrishna Rao. Separation theorems for singular values of matrices and their applications in multivariate analysis. Journal of Multivariate Analysis, 9(3):362-377, 1979.
[49] Karl Pearson. Liii. on lines and planes of closest fit to systems of points in space. The London, Edinburgh, and Dublin philosophical magazine and journal of science,
1901.
[50] Wojciech Marian Czarnecki, Rafal Jozefowicz, and Jacek Tabor. Maximum entropy linear manifold for learning discriminative low- low-dimeniional representation. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2015, Porto, Portugal, September 7-11, 2015, Proceedings, Part 15, pages 52-67. Springer. 2015.
[51] Michael Tschannen, Josip Djolonga, Paul Rubenstein, Sylvain Gelly, and Mario Lucic. On mutual information maximization for representation learning. Proceedings of the International Conference on Learning Representations (ICLR), 2020.
[52] Jason M Klusowski and Andrew R Barron. Approximation by combinations of relu and squared relu ridge functions with é and e controls. IEEE Transactions on Information Theory, 6A(12):7649-7655 2018
[53] Andrew Barron. Universal approximation bounds for superpositions of sigmoidal function. IEEE Transactions Information theory, 39(3):930-945, 1993.
13

---

[54] Alexander Kraskov, Harald Stögbauer, and Peter Grassberger. Estimating mutual information. Physical review E, 69(6):066133) 2004.
[55] Cédric Malherbe and Nicolas Vayatis. Global optimization lipschitz functions. In Interna- lional Conference Machine Learning, pages 2314-2323, PMLR, 2017.
[56] Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Infogan: Interpretable representation learning by information maximizing generative adversarial nets. Advances in neural information processing systems, 29, 2016.
[57] Ravid Shwartz- Ziv and Yann LeCun. To compress or not scl-ersssssll-ssseriiisedddd learning and information theory: A review. arXiv preprint arXiv:2304.09355, 2023.
[58] Yury Polyanskiy and Yihong Wu. Information Theory: From Coding to Learning. Cambridge University Press, 2022.
[59] Jan Magnus and Heinz Neudecker. Matrix differential calculus with applications in statistics and econometrics. John Wiley & Sons, 2019.
[60] Djork-Amé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint 1.0022151..07289, 2015.
[61] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 2014.
[62] Aaron Defazio, Francis Bach, and Simon Lacoste-Julien. Saga: fast incremental gradi- ent method with support for non-strongly convex composite objectives. Advances in neural information processing systems, 27, 2014.
[63] Alex Krizhevsky, Geoffrey Hinton, et Learning multiple layers of features from tiny images.
2009.
14


Reviews:
[{"review_id": "3Tl3qioyhL", "reviewer": "Reviewer_oZ2r", "summary": "This paper defines a novel measure of independence called Max-Sliced Mutual Information (mSMI) which provides a non-linear generalization of Canonical Correlation Analysis. This measure can also be viewed as a variant of Mutual Information with a better tractability in high dimensions.\nIn fact, mSMI is closely related to average-sliced mutual information (aSMI) introduced in [26,27]. Whereas aSMI is computed by averaging an objective involving mutual information over a product of Stiefel manifold $St(k,d_x) \\times St(k,d_y)$, mSMI is actually obtained by maximizing this objective over the same domain.\nThis work also introduces a generalized mSMI where the optimization is performed (loosely speaking) on a class of non-linear functions. \n\nBy leveraging the Donsker-Varadhan variational form for the mutual information, the authors show that mSMI can be estimated thanks to a neural estimator for which Theorem 1 provides convergence rates. \n\nFrom the numerical perspective, three applications are studied: independence testing, multi-view representation learning and fair representation learning.\nFor independence testing, mSMI seems to better perform than aSMI when $k$ is not too small.\nThe time complexity of aSMI is also reported to be larger than the time complexity of mSMI.\nIn the comparisons with a few baselines in the case of multi-view representation learning or fair representation learning, mSMI is often the best method in terms of performance.\n\n\n", "questions": "- Here is a question about Theorem 1. Why is $\\|\\mathcal{X}\\times\\mathcal{Y}\\|$ defined at the end of this statement? This quantity does not appear explicitly in the theorem. This is a bit confusing.\n- In Table 2, the values of $k$ range from $1$ to $7$ whereas they range from $3$ to $7$ in Table 5 of the supplementary material (Adult data set). What happens for small $k$ in the case of the Adult data set ?", "limitations": "I do not foresee a negative societal impact of this work.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This paper is written clearly and its results are significant.\nThe theoretical contributions are interesting and the proofs given in supplementary material are easy to read.\nThe construction of mSMI nicely arises as a variant of aSMI [26,27] and the mSMI neural estimator has a clear interpretation.\nI find the numerical comparisons in section 5 convincing.", "weaknesses": "I do not see serious methodological weaknesses. Here are some minor remarks.\n- The rule for emphasizing entries with bold in Table 2 of section 5.4 is unclear to me. In particular, $\\rho_{HGR}$ equals $0.958$ for $k=2$ and is not displayed in bold whereas the entry corresponding to $k=6$ is emphasized although it takes a smaller value ($0.957$).\n- In the same section, it is unclear how $k$ has to be selected by the user since mSMI does not outperform SLICE for all the values of $k$.\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for feedback and comments, which we address below:\n\n**1. Bold display of both $\\rho_{\\mathsf{HGR}}$ values:** Thank you for this observation. The quantity that measures fairness is $\\rho_{\\mathsf{HGR}}(Z,A)$ and we aim to minimize it while maintaining a high value of $\\rho_{\\mathsf{HGR}}(Z,Y)$, so that the ability to predict $Y$ from $Z$ is not compromised. As seen in the table, the most fair result (i.e., when $\\rho_{\\mathsf{HGR}}(Z,A)$ is smallest) is achieved for $k=6$ and our intention was to highlight that. However, we acknowledge that marking both values of $\\rho_{\\mathsf{HGR}}(Z,A)$ and $\\rho_{\\mathsf{HGR}}(Z,Y)$ in boldface may be confusing, as the corresponding $\\rho_{\\mathsf{HGR}}(Z,Y)$ was not maximal. For that reason, in the revision we will remove the boldface from $\\rho_{\\mathsf{HGR}}(Z,Y)$ and mark only $\\rho_{\\mathsf{HGR}}(Z,A)$.\n\n**2. About the choice of $k$:** The choice of $k$ is indeed an interesting and important question—one that we plan to develop a theory for in the future. At the moment, we adopt an empirical approach that sweeps over a range of $k$ values and treats it as a hyperparameter of the task. \n\nWhile an in-depth investigation of this point is left for future work, we briefly discuss an interesting tradeoff concerning the choice of $k$: between sample complexity and capturing as much information as possible in the mSMI. On one hand, note that the $k$-dimensional mSMI increases with $k$, and is uniformly (in $k$) upper bounded by the classical mutual information between the variables in the ambient space. This fact encourages the choice of a larger $k$. On the other hand, recall from our Theorem 1 that the sample complexity of mSMI estimation grows rapidly with $k$. Hence, as $k$ increases there should be a tradeoff between the returns of increasing the population mSMI, and the growing sample complexity. A particularly interesting setting is when the supports of the distributions lie in $d'$-dimensional subspaces. In this case, $k=d'$ is sufficient to capture the full classical mutual information, and increasing $k$ further only serves to degrade sample complexity without further gain. Extrapolating from this point, we conjecture that the optimal value of $k$ is related to the intrinsic dimension of the data distribution, even when it is not strictly supported on a low-dimensional subset. We hope to prove this conjecture in the future and will discuss this point in the Conclusions section of the revision.\n\n**3. About the definition of $\\lVert\\mathcal{X}\\times\\mathcal{Y}\\lVert$ in Theorem 1:** We acknowledge that the original statement was somewhat confusing. Our intention was to list the parameters on which the constant $C$ depends, with $\\lVert\\mathcal{X}\\times\\mathcal{Y}\\rVert$ being one of them. To save space, we had defined this quantity at the same spot where it was listed. To avoid confusion, in the revision we will modify the wording after the bound to: \n“... where the constant $C$ depends on $M$, $b$, $k$, and the radius of the ambient space, which is given by $\\lVert\\mathcal{X}\\times\\mathcal{Y}\\rVert\\coloneqq \\sup_{(x,y)\\in \\mathcal{X}\\times\\mathcal{Y}}\\lVert(x,y)\\rVert$.”\n\nWe hope that this resolves the issue.\n\n**4. About omitting $k=$ 1,2 values in Table 5:** In Table 5 of the supplementary material we have omitted the results for k=1,2 since we thought it was sufficient to run the experiment around the clearly best performing values of $k=$ 4 (or 5). Nevertheless, we agree with the reviewer that it is better to include the $k=$ 1,2 values for consistency and will do that in the revision. The $\\rho_{\\mathsf{HGR}}(Z,A)$ for $k=1,2$ is $0.43$, $0.393$, respectively.\n"}, {"review_id": "p0FSJ3EjRv", "reviewer": "Reviewer_WmTq", "summary": "The paper introduces an adaptation of sliced mutual information that focuses on the maximal mutual information between linear projections of low dimensionality of random variables. This measure (mSMI) has desireable properties and is approachable by neural estimation. The authors show that mSMI can be approximated for a certain class of continuous random variables more efficiently than (sliced) mutual information. In some illuminating experiments potential use cases are illustrated and suggest good practical performance also in such use cases.", "questions": "How can the results for fairness aware methods in Table 2 be better than fairness agnostic?", "limitations": "The authors are honest in the limitations:\nThe experiments show that for some cases sliced mutual information might work better than mSMI. Further research on the best choice of the dimension k to use might be interesting but are clearly beyond the scope of this work and would only clutter the presentation.\nA proof for the non-linear method is not yet availabvle but most likely also something for future work and a separate manuscript.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is well written and presents a new measure of dependence from all essential perspectives: Useful theoretical properties, approximation is theoretically possible for a large class of random variables, and practical implementation is possible.\nThe paper is not overloaded in content and focuses on the topic at hand providing all necessary information without giving unnecessary details or too many experiments.", "weaknesses": "I don't see any constructive and actionable insights on how the work could improve towards its stated goals.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for feedback and comments, which we address below:\n\n**“How can the results for fairness aware methods in Table 2 be better than fairness agnostic?”**\n\n**Answer:** This is an excellent question. Note that the reported $\\rho_{\\mathsf{HGR}}$ coefficients are evaluated on test data that was not used in training. As a result, the reported $\\rho_{\\mathsf{HGR}}$ values depend on the ability of the learned representations to generalize to unseen examples. If it happens that the protected attributes $A$ in the problem are spurious features that are not causally predictive of the outcome $Y$, then a representation $Z$ that overlooks $A$ may generalize better for predicting $Y$. This can explain the small improvement in $\\rho_{\\mathsf{HGR}}(Z,Y)$ observed when using fairness-aware methods. That said, the above is a conjecture and it strongly relies on the intrinsic structure of the problem. It may be the case that this trend does not persist for other datasets/tasks."}, {"review_id": "Bm1PUL4oQA", "reviewer": "Reviewer_rZV3", "summary": "The paper proposes Max-Sliced Mutual Information (mSMI) which equals the maximal mutual information between low-dimensional projections of the high-dimensional variables. The mSMI can capture intricate dependencies in the data while being amenable to fast computation and scalable estimation from samples. In addition, the paper proposes multivariate conditional mSMI, and max sliced entropy which are extensions of mSMI. The paper discusses the structural properties of mSMI including bounding on different values of sub-space dimension and the original mutual information, Identification of independence, KL divergence representation, Sub-chain rule, and Tensorization. Moreover, the paper discusses how Gaussian mSMI is related to canonical correlation analysis (CCA) and how Max-sliced entropy is related to  PCA. In addition, generalized mSMI with two general classes of functions is also discussed. The neural estimation of mSMI and its error are presented in Section 4. Finally,  the paper presents experiments that demonstrate the utility of mSMI for several synthetic and real-world tasks, encompassing independence testing, multi-view representation learning, and algorithmic fairness.", "questions": "Can the paper include an application in the previous SMI papers e.g., InfoGAN?", "limitations": "Not applicable", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper is detailed and very well-written.\n* Max-sliced mutual information is a natural extension to SMI.\n* Connection between mSMI and conventional methods such as CCA and PCA.\n* Neural estimation of sMI has a better error than SMI since no Monte Carlo estimation is needed.\n* mSMI can be implemented efficiently with Neural Estimation by merging the linear projection to the neural network.\n* Experiments on independence testing show that mSMI is both better and faster than SMI.", "weaknesses": "* mSMI is quite incremental based on the existence of SMI, k-SMI, Max-sliced Wasserstein distance, and Max-K sliced Wasserstein distance.\n* There is no comparison between mSMI and SMI (K-SMI)  in Multi-View Representation and Learning Fair Representation. This is quite questionable since mSMI is an extension of SMI (K-MSI). I believe a comparison in both computation and performance is needed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for feedback and comments, which we address below:\n\n**1. max-SMI is incremental with respect to SMI and sliced Wasserstein measures:** Thank you for bringing this point up. We divide our response into two parts: (i) comparison with sliced Wasserstein distances, and (ii) comparison with other sliced information measures. \n\n(i) While the original SMI, proposed in [26], was inspired by slicing techniques for Wasserstein distances (as mentioned in the introduction of that work), we note that these objects quantify different things. Sliced Wasserstein distances measure discrepancy between probability distributions, while mutual information (MI) and its sliced variants quantify dependence between random variables. Additionally, as Wasserstein distances are rooted in optimal transport theory, the corresponding notion of discrepancy is geometric in nature and adapts to the structure of the metric space in which the data resides. Standard/sliced MI, on the other hand, is induced by the KL divergence, which is an entropy-based quantity that only depends on the log-likelihood of the considered distributions and overlooks the underlying geometry. For these reasons, we view SMI measures and sliced Wasserstein distances as not directly comparable/related, despite the said inspiration.\n\n(ii) Concerning the comparison to SMI or $k$-SMI, we believe that max-SMI (mSMI) addresses core limitations of the former which are crucial for machine learning applications. Specifically, SMI and $k$-SMI were proposed to circumvent the statistical and computational difficulties associated with classical MI in high-dimensional settings, while inheriting important structural properties from it. However, the average-sliced methods still suffer from a burdensome Monte Carlo (MC) step that was needed for their estimation/computation and lacked interpretability of the notion of dependence being quantified. The mSMI addresses both these issues as follows. First, by replacing the average over slices with a maximum, mSMI rids the MC step for estimation, replacing it with a simple optimization that is readily absorbed into the neural estimation paradigm with negligible cost. Thus, while estimation of average-SMI or $k$-SMI requires computing $m$ different estimators ($m$ being the number of MC samples), one for each slice, mSMI entails only a single estimation problem. As discussed in Section 5.1 and illustrated in Figure 1, this results in significant speedups in runtime with no loss of performance (see, e.g., Figure 2 for independence testing). \n\nSecond, we argue that mSMI also enjoys better interpretability than its average-sliced counterparts due to the equivalence, under the Gaussian setting, to the well-understood notion of CCA. Indeed, for Gaussian data, mSMI and CCA coincide, and both the optimal projection directions and the mSMI value adhere to simple closed-form solutions (see Proposition 2). Since CCA is a classical idea that has been thoroughly studied over the years, our understanding thereof is inherited by Gaussian mSMI, endowing it with a clear interpretation. By the same token, the relation between max-sliced entropy and PCA, which is discussed in Supplement A.3, serves a such an interpretability role. To the best of our knowledge, no such crisp connections are available for average-SMI variants. \n\nIn summary, while SMI and $k$-SMI were important steps towards a tractable and efficiently computable measure of information, mSMI provides another significant improvement over those approaches in several aspects. We therefore do not view the contribution of our work as incremental and believe that mSMI will serve as a useful tool for high-dimensional machine learning applications. To further clarify the above points, we will edit the Section 1.1 of the introduction in the final version to better motivate mSMI and contrast it with existing average-sliced methods. We will also add a remark before Section 3.1, where a discussion to the effect of the above will be provided.\n\n**2. Comparison between mSMI and SMI ($k$-SMI) in representation learning tasks:** This is a great question. The goal of these experiments was to use the considered measures (SLICE and mSMI for fairness; CCA and mSMI for multi-view representation) to extract the best latent features for a corresponding downstream task. While mSMI provides such a latent feature (namely, the optimal projection direction), average-sliced information measures do not. Indeed, SMI and $k$-SMI average multiple MI terms between the projected variables, and it is unclear how to extract a single feature to represent the data from them. For this reason, SMI and $k$-SMI were not used for this experiment. \n\nOne could alternatively attempt to maximize SMI or $k$-SMI as an objective for feature extraction. However, this would reduce back to mSMI, as discussed in Remark 3 of our paper and proven in Proposition 4 of [26]. This further justifies why average-sliced variants were not considered for these experiments. We will add text to Sections 5.3 and 5.4 to clarify the above.\n\n**3. Including an InfoGAN application:** Thank you for this suggestion. Our choice of experiments was driven by two main considerations: (i) present instances where the mSMI enables improving upon existing methods (e.g., neural estimation and independence testing, for which mSMI indeed performs better/faster), and (ii) diversify the experiment portfolio beyond the examples for which sliced information measures were previously used (e.g., multi-view representation and fairness, which were not considered before). That said, we recognize the potential of mSMI for constructing a 'max-sliced InfoGAN' and are happy to explore this application in order to include such results in the camera-ready version. We will also compare the performance to that of the average-sliced SMI InfoGAN from [27]."}, {"review_id": "K9LvhrLGbf", "reviewer": "Reviewer_YCWE", "summary": "The authors propose max-sliced mutual information (mSMI), which inherits important properties of mutual information. They show that the projection in mSMI is reduced to that in CCA for jointly Gaussian. They also provide a method to estimate mSMI using a neural network, which is computationally more efficient than an existing method for estimating average-SMI. They derive an error bound for the estimation of mSMI.", "questions": "In Remark 4, the authors insist that if the data is not Gaussian, the projections in mSMI do not coincide with those in CCA. Do you have any simple examples of how the projections in mSMI and those in CCA are different?\n\nMinor comments\n- The definition of I is missing in Eq. (1). The definition is in line 110, but it should be defined here.\n- In Definition 2, G and H should be subsets of $\\{g:\\mathbb{R}^{d_x}\\to\\mathbb{R}^k\\}$ and $\\{h:\\mathbb{R}^{d_y}\\to\\mathbb{R}^k\\}$ ?", "limitations": "Theorem 1 is for the shallow network with the ReLU activation function. Although the authors insist that the ideas can be extended to other nonlinearities and deep architectures, how we can extend the theoretical results like Theorem 1 to the deep case and other activation functions is not clearly explained. Are there any difficulties? If there are, do you have any ideas to overcome the difficulties?", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper includes fundamental properties of mSMI, connections with existing studies, and a practical method to estimate mSMI. The results provided in this paper are sufficient to understand mSMI.", "weaknesses": "Although the paper is well-written for the most part, some parts should be explained clearly, or need more detailed explanations. For example, what information can be extracted from mSMI if the data is not Gaussian (in Remark 4), and how we can extend the ideas in 4.2 to mode general neural networks. See the sections on questions and limitations for more details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for feedback and comments, which we address below:\n\n**1. Inequivalence between mSMI and CCA in the non-Gaussian case:** This is an interesting question. We first note that the equivalence between mSMI and CCA in the Gaussian case hinges on the sufficiency of the cross-covariance matrix to fully characterize the dependence between the projections of $X$ and $Y$. As non-Gaussian mutual information (MI) cannot be characterized merely by second order statistics, we do not expect this equivalence to hold in general. Capturing this inequivalence analytically seems challenging as it requires a closed-form formula for the MI between the projected variables, which can further be analytically optimized over the corresponding Stiefel manifolds. We are currently exploring such examples and hope to add one to the final version of the paper.\n\nIn the meantime, we provide numerical evidence for the above. Consider a 3-dimensional random variable X, whose components are independent and uniformly distribution on $[-1,1]$, i.e., $X=(X_1\\ X_2\\ X_3)^\\intercal$ with $X_i\\sim\\mathsf{Unif}[-1,1]$. Define $Y=(Y_1\\ Y_2\\ Y_3)^\\intercal$ by  $Y_i=X_i^2-\\mathbb{E}[X^2]+U_i$, where $U_i\\sim\\mathsf{Unif}[-1,1]$ for $i=1,2,3$. Using a simple CCA solver and numerically solving for mSMI, we obtain different maximizing projection directions in $\\mathsf{S}t_{1,3}$. Specifically, a python simulation yields\n\n$\\theta_{CCA}=(0.151, -0.449,  0.88)$, $\\phi_{CCA} = (0.129, 0.106, 0.985)$, \n\n$\\theta_{\\bar{\\mathsf{SI}}}=(-0.063, -0.99, -0.122)$, $\\phi_{\\bar{\\mathsf{SI}}} = (-0.062, 0.998, 0.0195)$\n\nwhich are significantly different. While we hope to have an analytical example to include in the revision, we are happy to include the numerical example in the text (perhaps with some supporting graphics) as a contingency plan.\n\n**2. Generalizing Theorem 1 to other nonlinearities and deep networks:** Thank you for this excellent question. We divide our response into two parts: addressing other activation functions first and then discussing deep neural estimation (NE). A list of references that are used throughout this answer that do not appear in the paper is provided below.\n\n(i) Non-ReLU activations: Our NE bound from Theorem 1 (as well as those in the average-sliced MI papers) are based on the theory developed in [39]. As highlighted therein, their bounds extend to nonlinearities beyond ReLU activation. Specifically, their theory accounts for any sigmoidal bounded activation with $\\lim_{z\\to-\\infty}\\sigma(z)=0$ and $\\lim_{z\\to\\infty}\\sigma(z)=1$.\nTo adapt our NE error bound to sigmoid activations, one should use the approximation error bounds from [A], instead of the currently used ones from [B]. This change will require a small modification of the class of distributions $\\mathcal{P}_{\\mathsf{KL}}$ that our theory accounts for (see definition in Section 4.2). In particular, one would have to consider an extension $r\\in\\mathcal{C}_b^{k+2}$, rather than $r\\in\\mathcal{C}_b^{k+3}$. However, as noted in [39], to achieve the $O(\\ell^{-1/2})$ approximation error with sigmoid activations one must scale the hidden layer parameters as $\\ell^{1/2}\\log \\ell$, where $\\ell$ is the number of neurons. With ReLU activations, on the other hand, the network has bounded parameters independent of $\\ell$. For that reason, we opted to use ReLU networks in our statement. We plan to include discussion to the effect of the above in a remark following Theorem 1.\n\n[A] Barron, Andrew R. “Universal approximation bounds for superpositions of a sigmoidal function.” IEEE Transactions on Information theory 39.3 (1993): 930-945.\n‏\n[B] Klusowski, Jason M., and Andrew R. Barron. “Approximation by combinations of ReLU and squared ReLU ridge functions with $\\ell^ 1$ and $\\ell^ 0$ controls.” IEEE Transactions on Information Theory 64.12 (2018): 7649-7656\n\n(ii) Concerning the extension to deep NEs, this is an interesting question that is on our research agenda going forward, as noted in the second paragraph of Section 6. The NE theory for $f$-divergences from [39], on which our Theorem 1 relies, decomposes the error analysis into two parts: function approximation and statistical estimation. The former is controlled using approximation error bounds for sparse shallow nets, while the latter requires control over the covering number of the corresponding neural network class. Sharp bounds for these aspects are available in the literature (cf. Theorems 8 and 11 from [39]), which results in the optimal $n^{-1/2}$ convergence rate for shallow NEs, as established in [39] and exploited in our work. Similar approximation error and covering number bounds are available for deep networks; cf., e.g., [C,D,E]. However, we believe that these bounds are not sharp since the best NE error rate that can be obtained from them is $O(n^{-1/4})$. The advantage of deeper architectures is that approximation is possible under relaxed smoothness assumptions on the distributions (compared to the shallow case), but we expect that the same $n^{-1/2}$ parametric rate would be achievable. This suggests that more work is needed on the fundamental components of the deep NE analysis to obtain a satisfactory theory. While this is a fascinating research avenue, it falls outside the scope of the current paper and is left for future work. Nevertheless, we will expand our discussion on this topic in Section 6 and provide more details as above.\n\n[C] Schmidt-Hieber, Johannes. “Nonparametric regression using deep neural networks with ReLU activation function.” Annals of Statistics 48.4 (2020): 1916-1921.\n\n[D] Bresler, Guy, and Dheeraj Nagaraj. “Sharp representation theorems for relu networks with precise dependence on depth.” NeurIPS 33 (2020): 10697-10706.\n\n[E] Shen, Zuowei. \"Deep Network Approximation Characterized by Number of Neurons.\" Communications in Computational Physics 28.5 (2020): 1768-1811.\n\n**3. Minor comments:** Thank you. These will be corrected in the revision."}]

Author rebuttal:

Meta-review: The focus here is on the estimation of the dependence of two random variables in finite-dimensional Euclidean spaces (R^{d_x} and R^{d_y}), in a scalable manner. By taking the most informative linear projection meant in mutual information (MI) sense, the authors propose the max-sliced variant of MI (mSMI), conditional MI and entropy (sh). After showing structural properties of these measures (Proposition 1), and their relation to CCA (canonical correlation analysis) and PCA (principal component analysis) in the Gaussian case, nonlinear extension/slicing is also proposed (Def. 2) with analogous properties as in Proposition 1. Leveraging the variational form of mSMI, the authors devise a neural estimator (5) for it, with error guarantees (Theorem 1). The practical efficiency of the proposed estimator is demonstrated in the context of independence testing, multi-view representation learning, and algorithmic fairness.

Estimating information theoretical measures efficiently is a fundamental building block in various problems of data science. The authors present a new, practically relevant, principled tool in this well-structured and clearly-written manuscript, which can be of definite interest to the NeurIPS community.
Decision: Accept (poster)