Efficient Distribution Matching of Representations via Noise-Injected Deep InfoMax
Authors: Ivan Butakov
Alexander Semenenko
Alexander Tolmachev
Andrey Gladkov
Marina Munkhoeva
Alexey Frolov
Venue: ICLR 2025 Poster (2025)
Keywords: information theory
deep infomax
self-supervised learning
representation learning
distribution matching
Gaussian embeddings
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: This paper enhances Deep InfoMax approach to self-supervised representation learning via the injection of additive independent noise, allowing for the representations conforming to specific distributions to be learned.

Abstract: Deep InfoMax (DIM) is a well-established method for self-supervised representation learning (SSRL) based on maximization of the mutual information between the input and the output of a deep neural network encoder. Despite the DIM and contrastive SSRL in general being well-explored, the task of learning representations conforming to a specific distribution (i.e., distribution matching, DM) is still under-addressed. Motivated by the importance of DM to several downstream tasks (including generative modeling, disentanglement, outliers detection and other), we enhance DIM to enable automatic matching of learned representations to a selected prior distribution. To achieve this, we propose injecting an independent noise into the normalized outputs of the encoder, while keeping the same InfoMax training objective. We show that such modification allows for learning uniformly and normally distributed representations, as well as representations of other absolutely continuous distributions. Our approach is tested on various downstream tasks. The results indicate a moderate trade-off between the performance on the downstream tasks and quality of DM.

Published as a conference paper at ICLR 2025
EFFICIENT DISTRIBUTION MATCHING OF REPRESEN- TATIONS VIA INJSE-INJECTED DEEP INFOMAX
Ivan Butakov*.1., Alexander Semenenko', Alexander Tolmachev!.2, Andrey Gladkov
Marina Munkhoeva', Alexey Frolov
Skolkovo Institute of Science and Technology; "Moscow Institute of Physics and Technology:
3 AAttificia Intelligence Research Institute;
(butakov.id, semenenko.. av, tolmachev.ad, munkhoeva@airl net, al. frolov@skoltech. too.
Mladkov.ao@@phyyeee www.saeteesetttcc.uuuu
ABSTRACT
Deep InfoMax (DIM) is a well-established method for self-supervised representa- tion learning (SSRL) based on maximization of the mutual information between the imput and the output of a deep neural network encoder. Despite the DIM and contrastive SSRL in general being well-explored. the task of learning representa- tions conforming to a specific distribution (i.e., distribution matching. DM) is still under-addressed. Motivated by the importance DM several downstream tasks (including generative modeling, disentanglement, outliers detection and other), we enhance DIM to enable automatic matching of learned representations to a selected prior distribution. To achieve this, we propose injecting independent noise into the normalized outputs of the encoder, while keeping the same InfoMax training objective. We show that such modification allows for learning uniformly and normally distributed representations, as well as representations of other ab- solutely continuous distributions. Our approach is tested various downstream tasks. The results indicate a moderate trade-off between the performance on the downstream tasks and quality of DM.
INTRODUCTION
Learning viable low-dimensional representations of complex data plays an important role in many modern applications of artificial intelligence. This task arises in various domains, including im- age (Haralick et al., 1973; Chen et al., 20206] Rombach et al. 2022),, audio (van den Oord et al.
2019),
and natural language processing (Mikolov al.] 2013 Radford et al., 2018 Devlin al. High-quality embeddings are particularly useful for multi-modal methods (Vinyals et al.
2019
2015
Radford et al. [021 [o & Salimans] 202T), statistical and topological analysis (Moor al.
2020
Duong & Nguyen, 2022) Butakov et al., (024b). data visualization (van der Maaten & Hin-
ton
2008
McInnes et 2018) and testing fundamental hypotheses (Brown al.
2023
Gurnee
&
Tegmark, 2024 Huh et al. 2024).
Existing approaches representation learning can be divided into three categories (Ericsson al., 2022]: supervised (requires labeled data), self-supervised and unsupervised (no labeling is required). In practice, access to labeled data is limited, which hinders the use of supervised approaches. There- fore, unsupervised and self-supervised methods are of great importance. Contrastive learning is well-established paradigm of self-supervised representation learning (SSRL), which encourages an encoder to learn similar representations for various augmentations of the same data point, and dis- similar for different data points. Deep InfoMax (DIM) (Hjelm et al. 2019) leverages information- theoretic quantities to construct decent contrastive objective, involving direct maximization of the useful information contained in the embeddings. DIM is universal, flexible, and deeply connected the rigorous information theory, which allows for a good performance on variety of downstream
tasks to be attained (Hjelm et al. 2019
Bachman et al. 2019 Veličković et al.) 2019: Tschannen
al. 2020 2024).
"Correspondence

---

Published as a conference paper at ICLR 2025
Acquiring embeddings admitting specific distribution (i.e., distribution matching, DM) is iliary, yet important task in representation learning. Latent distributions with straightforward sam-
pling procedures tractable densities are crucial for downstream generative modeling (Kingma &
Welling 2014 Makhzani et al. 2016 Larsen et al. 2016] Papamakarios et al., 2021 Addition-
ally, specific distributions (e.g., Gaussian) exhibit properties, which are useful for statistical analysis
(Tipping & Bishop 1999 Duong & Nguyen 2022). disentanglement (Higgins et al. 2017 Balabin
al. 2024). and outliers detection (Bazarova et 2024).
A classical approach to latent DM is to optimize a cheap and imprecise distribution dissimilarity measure during the training or architecture search (gg. 2011, Makhzani Frey 2014, Kingma &
Welling 2014: Heusel et al., [017 Higgins et al., 2017). Methods of this family have to rely on
several strong assumptions, such as embeddings already admitting a Gaussian distribution, which eventually leads to suboptimal results. Another common approach employs adversarial networks to push the learned representations towards a desired distribution (Makhzani et al.] 2016] Hjelm et al. 2019), One can also leverage generative models (e.g., normalizing flows or diffusion mod- els) to explicitly perform DM "post hoc" (Böhm Seljak] 2022] Rombach et al. 2022). These
two approaches yield decent results, but require supplementary networks to match the distributions. Finally, injective likelihood--aal based models can also be used to acquire low-dimensional representa- tions admitting desired distribution (Brehmer & Cranmer [2020] Sorrenson al.) 2024). However, likelihood maximization across dimensions is notoriously problematic due to non-square Jacobi ma- trices.
In contrast to the mentioned approaches, we propose a simple, cost-effective (requiring no additional neural networks) and non-intrusive modification DIM, which allows for an automatic and exact DM of the representations. In the following text, we show that using specific activation functions and noise injections at the outlet of an encoder, combined with the DIM objective, allows for normally and uniformly distributed representations to be learned. Our contributions are the following:
1. We prove that applying normalization and adding a small noise at the outlet of an encoder and maximizing the DIM objective minimizes the Hellbackk-Libble divergence between a Gaussian (or uniform) distribution and the distribution of embeddings.
2. We conduct experiments on several downstream tasks to explore the tradc-off between the downstream performance and the accuracy of DM via our method.
3. We conduct additional experiments assess the quality of DM via our method in the tasks of generative modelling.
The paper is organized as follows. In Section[] the necessary background from information theory and original works on Deep InfoMax is provided. In Section we describe the general method for DM via modified DIM. In Section 4 a connection between the proposed approach and other SSRL methods is established. Section dedicated to the experimental evaluation of our method. Finally, we conclude the paper by discussing our results in Section 6. Complete proofs, additional analysis of tatomax-basedd DM, and technical details are provided in Appendices A C correspondingly.
2 BACKGROUND
In this section, the background necessary to understand our work is provided. We start with the basic definitions from the information theory. Then, the maximum entropy theorems given, which are crucial for understanding our approach. Finally, we outline the general variant of Deep InfoMax representation learning method, which we aim enhance with an automatic distribution matching.
2.1 PRELIMINARIES
Let (R,F,PP be probability space with sample space Ω, o-algebra F. and probability measure P defined on F. Consider an absolutely continuous random vector X Ω + I2d with the probability density function (PDF) denoted as p(z)). The differential entropy of X is defined follows:
((X)=
E EEogg(()) - p(x) log p(x) dr,
supp
2

---

Published as a conference paper at ICLR 2025
where supp X Rd represents the support of X, and log( .) denotes the natural logarithm. Simi- larly, we define the joint differential entropy as htX,, Y) = log p(X, Y) and conditional differ- ential entropy as ((X Y) = p (X Y = Ey (mxx log X Y)). Finally, the mutual information (MI) is given by I(X;Y) h(X) h(x Y), and the following equivalences hold
T(X;Y) l(x) 1(x Y) M(Y) M( X);
I(x:Y) = h(x) + h(Y) h(X,Y), I(X;Y) = DKL (Px.y P
Mutual information can also be defined as an expectation of the pointwise mutual information:
p(æ v) p(x)
(a, = log
I(X;Y) = E x
(I)
The above definitions can be generalized via Radon-Nikodym derivatives and induced densities in case of distributions supports being manifolds, see (Spivak) 1965). We use flexible notation, where entropy and divergence may refer to random vectors or their corresponding distributions.
Our work leverages the maximum entropy properties of Gaussian and uniform distributions:
Theorem 2.1 (Theorem 8.6.5 in Cover & Thomas (2066). Let X be d-dimensional absolutely continuous random vector with probability density function p, mean m and covariance matrix E. Then
1
1(X) = h (Mm,, Σ)) DKL (p (m,, E)), (W E)) = log ((2Te)* det E)
2
where (mn, Σ) is Gaussian distribution of mean and covariance matrix E.
Theorem 2.2. Let X be an absolutely continuous random vector with probability density function p and supp X C S, where S has finite and non-zero Lebesgue measure ((S).. Then
h(X)= h (U(S)) - DRL (p U(S)),
h (U(S)) = log (SS),
where U(S) is uniform distribution on S. Remark 2.3. Note that ((X)= ) + 1(X') Dkr. (px pxx) ) in general.
Finally, we also utilize the following lower bound on the conditional entropy of sum of two condi- tionally independent random vectors:
Lemma 2.4. Let X and Z be random vectors of the same dimensionality, independent under the conditioning vector Y. Then
h(X + Z Y) = x - 2
with equality and only if there exists a measurable function such that X = g
2.2 DEEP INFOMAX
Mutual information (MI) is widely considered as a fundamental measure of statistical dependence between random variables due its key properties, such as invariance under diffeomorphisms, sub- additivity, non- neg-negativ and symmetry (Cover & Thomas] 2006).. These attributes make MI partic- ularly useful in information-theorrttee approaches to machine learning. The concept of maximizing MI between input and output, known as the infomax principle (Linsker) 1988 Bell & Sejnowski 1995), , serves the foundation for Deep InfoMax (DIM), family of self-supervised representation learning methods proposed by Hjelm et al. (2019).
A naive DIM approach suggests learning the most informative embeddings via direct maximization of the mutual information between the original data and compressed representations:
T(x;((X)) + max,
"eeeeinftee,, when comparing random variables, mean equality inequality "almost sure".
3

---

Published as a conference paper at ICLR 2025
where X is the random vector be compressed, and is the encoding mapping (being learned). However, despite its simplicity and intuitiveness, this setting is rendered useless by the fact that [[X; f(X)
for wide range of X and f (Bell & Sejnowski 1995 Hjelm et al. 2019).
To avoid this limitation, it is usually suggested to augment (crop randomly, add noise, etc.) original data to inject stochasticity and make Clormation-theoreel objectives non-degenerate. Consider the following Markov chain:
S(X) x
where X' is augmented data. Now T(X:; f(X)) can be made non-infinite. Moreover, according the data processing inequality (Cover & Thomas 2006), (('; f(x)) I f(X), which connects this new infomax objective with the naïve one.
The only remaining problem is the high dimensionality of X', which makes MI estimation difficult. To resolve this, one can apply an additional (in general, random) dimensionality-rcduinng transfor- mation replace X' with Y':
f(x)
-
with I(Y:; f(X)) < 1XX:; being the new infomax objective. For the sake of computational simplicity, Y' can be defined as f(X'); thus, the same encoder network used to compress both the original and augmented data. Moreover, one can view [(f(X'): f(X)) as a contrastive objective: this value high when ʃ yields similar representations for corresponding augmented and non- augmented samples, and dissimilar for other pairs of samples.
Now, let us consider the task mutual information maximization. As MI is notoriously hard to estimate (McAllester & Stratos 2020), lower bounds are widely used to reparametrize the infomax objective (Belghazi et al. 2021 van den Oord et al. [019 Hjelm al. 2019). this work, we only consider the two most popular approaches based on the variational representations of the Kullback- Leibler divergence:
Donsker-Varuud Varadhan
Donsker Varadhan 1983
I(X:Y) sup T: Q-+R
EFx, T log Ee O exp(T))
(2)
Belghazi et al. 2021)
Jordan
[gguyen al. 2010
I(X:Y)= sup T: Q-
E T E, Ex exp
Belghazi et al. 2021
where Q2 is the sampling space, and T is measurable critic function.
Remark 2.5. Assuming PMIx,y exists, the supremum in 2 and 3 is attained at and only at
T T (a. y) PMIx,y y (z, y) + (a.s. w.r.t. Fx,y),
where a for and can be any real number for (2) (follows from Theorem in Belghazi et al. 202T and Theorem 2.1 in Keziou (2003)).
In practice, f and T approximated via corresponding neural networks, with the parameters being learned through the maximization of the Monte-Carlo estimate of (2) or (3). Although (2) usually yields better estimates, the SGD gradients of this expression biased in mini-batch setting (Belghazi et al. 2021).
3 DEEP INFOMAX WITH AUTOMATIC DISTRIBUTION MATCHING
In the present section, we modify DIM to enable automatic distribution matching (DM) of represen- tations. Specifically, we propose adding independent noise Z to normalized representations of X, thus replacing f(X) ) by f(X) + in the infomax objective from Section 2.2
This corresponds to
the following Markov chain:
f(X) - - -X f(X).
By the data processing inequality, we know that ([(X'):((X) + 2) < I(f(X'): f(X),, which connects method to the family of conventional DIM approaches discussed previously.
As we will demonstrate, the noise injection at the outlet of an encoder, in conjunction with The- orems 2.1 and 2.2 maximum entropy, allows to achieve normally uniformly distributed
4

---

Published as a conference paper at ICLR 2025
embeddings. To prove this, we use the following Lemma 3.1 fomax objective into three components: the entropy of representations with an additive noise Z, the entropy of the noise itself, and the mutual information between f(X) and f + Z conditioned on f(x') -rpresentations of augmented data.
which decomposes the proposed in-
Lemma 3.1. Consider the following Markov chain of absolutely continuous random vectors:
f(X) + Z - X - X' - f(X'),
with Z being independent of (X, X'). Then
[(X + 7) h(f(X) + 7) h(Z) T(((X) 7: f(x) f(x'))
(5)
Lemma 3.1 highlights the importance of both additive noise and input data augmentation. Specifi- cally, if no augmentation is applied, i.e., X = X', then I(f(X) + Z; f(X) f(X)) = this case, maximizing the infomax objective is reduced maximizing the entropy h(f(X)++ + Z). Under the corresponding restrictions on (, this is equivalent to distribution matching, see Theorems and 2.2 In contrast, if no noise is added, I(f(X';; f(X)) is not bounded in general, which may prevent (f(()) from saturation in practical scenarios.
Additionally, DM alone does not guarantee that the learned representations will be meaningful or useful for downstream tasks. That is why we also show that using * X' allows us to recover meaningful embeddings, that are weakly invariant to selected data augmentations:
(X)(X) + Z: fXX) f(X') is set to zero by learning representations
Definition 3.2. We call encoding mapping f weakly invariant to data augmentation X X' if there exists a function such that f(X) (((())) = +(((X')) almost surely.
Lemma 3.3. Under the conditions of Lemma 3.1 let P(X X' x) a> 0. Then I(f(X) + f() f(X)) = precisely when is weakly invariant to X X"
Overall, the signal-to-noise ratio serves as a tradcoff between the distribution matching objective and robustness the data augmentations (the first and the third term in (5) correspondingly> higher magnitudes of Z impose tighter bounds on [ff(X)+Z; (X) fXX)), thus prioritizing maximiza- tion of (X) + Z). In what follows, we formalize the provided reasoning by addressing Gaussian and uniform distribution matching separately. We also briefly discuss how general DM problem can be reduced to the normal or uniform one.
3.1 GAUSSIAN DISTRIBUTION MATCHING
Let us assume Z having finite second-order moments. According to Theorem 2.1 if we restrict f(X) by fixing its covariance matrix, (x + Z) attains its maximal value precisely when f(X)+ + is distributed normally. Now, as Z is independent of X, f(X) + Z is normally distributed if and only the distribution of f(X) is also Gaussian.
achieve the normality of f(X) via restricting second-order moments of f(X) and maximizing the entropy (x((X)= + 2). combination with Lemma 3.1 this forms a basis of the proposed method for Gaussian distribution matching. Moreover, show that the imposed restrictions can be partially lifted via only requiring Var(f(X),, = for every E (1.....d d). This approach is preferable in practice, the widely-used batch normalization (Ioffe & Szegedy 2015) can be employed to restrict the variances. We formalize our findings in the following theorem.
Thus,
Theorem 3.4 (Gaussian distribution matching). Let the conditions of Lemma 3.3 be satisfied. As- sume Z ~ (00,, I2)), E f(X) = and Var (f(X))) = 1 for all i E {1...... dd. Then, the mutual information (Xf(X"); ((X) Z) can be upper bounded as follows
d
log (1+
(U(('):(X) 2)
2
with the equality holding exactly when / is weakly invariant and [(X) N/ I). Moreover, DKL. (f(X) (o) < I(Z; fx) + Z) I(f(X'); (XX) + Z) - dooo
While the first inequality in Theorem 3.4 suggests that exact distribution matching is achieved only asymptotically, the second expression shows that DKL (f(xX) (l.) can be bounded using the value of the proposed infomax objective. Note that f(x) + 2) is typically upper bounded, except the rare cases where the support of f(X) is degenerate.
5

---

Published as a conference paper at ICLR 2025
3.2 UNIFORM DISTRIBUTION MATCHING
Now consider Z distributed uniformly on [-e;e]". We assume that f(X) has bounded support alternative fixing the second-order moments, as in the Gaussian case, and is easy to implement in practice (e.g.. via a sigmoid function). By Theorem 2.2 the entropy h (x) Z) is maximized when f(X) + Z follows uniform distribution. This can be achieved when f(X) conforms to discrete uniform distribution over an equidistant grid within [0, 1]s. As & approaches zero, the distribution of f(X) gradually converges to continuous uniform distribution over [0, 1]6. Thus, injecting uniform noise asympiotically drives f(X) toward uniform representation over [0, 1|a. This reasoning formalized the following theorem.
within [0, 1] This restriction serves
Theorem 3.5 (Uniform distribution matching). Under the conditions of Lemma 3.7 let 7 ~ U(I=E;ald) and supp f(X) S 1j.. Then, the mutual information 1(f(X');
+ D) can
be upper bounded as follows
1 )
I(f(X':: f(X) + Z) s dlog
1+
22
with the equality if and only i/s E N, weakly invariant, and f(X) U(A), where the A = {0, 22, 4s,... contains (1/(2E) + elements.
Moreover
DEL (F(X) (((0:)]1) < f(Z:((X) + Z) T(J(X);; Z) dloog22s).
In sharp contrast Theorem 3.4 the equality in (7) is attained at f(X) conforming discrete distribution. This makes the proposed objective less attractive in comparison to the Gaussian distri- bution matching. However, one still can assured that the standard continuous uniform distribution allows us to approach the equality in (7) with approaching zero:
Remark 3.6 (Butakov et al. 2024a). If fXX U(O:: 1] f is weakly invariant, and < 1/2, then T(f(X';; f(X) + Z) = d log(2e)) dllog (1 + 1/(2s)) (log(1 + 2s) E).
o(E)
3.3 GENERAL CASE
Our approach can be extended wide range of desired distributions of embeddings using the probability integral transform (David & Johnson 1948 Chen & Gopinath 2000). Normalizing flows (lcarnable diffeomorphisms) can also be leveraged to transform the distribution "post-hoc"
due to their universality property (Huang et al. 2018 Jaini et al. 2019 Kobyzev et al. 2020).
The data processing inequality can used to estimate the quality of DM after the transformation:
Statement 3.7 (Corollary 2.18 in (Polyanskiy & Wu
2024 Let be measurable (w.rt IP and DKL (IP Q)) where IP og and P Pog denote the
@ function. Then DEL (IP og - push-forward measures IP and after applying g.
4 CONNECTION TO OTHER METHODS FOR SSRL
In this section, we explore the relation of our approach to other methods for unsupervised and self- supervised representation learning. We argue that the proposed technique for distribution matching can be applied to any other method, given the latter can be formulated in terms of mutual information maximization.
In reconstruction-bon -based generative models (Makhzani et al. 2016), the reconstruc- tion error closely tied to the mutual information (Theorem 8.6.6 (Cover Thomas 2006)):
Autoencoders
(X;Y)= N(X) h(x Y) M(X) - d
2TTe d N log - (
log
Here, X denotes the input, Y the latent representation produced by the encoder, and X (Y) the reconstructed input. The last term involving the expected squared error" is essentially the autoen- coder's loss function. By minimizing this reconstruction loss, the mutual information between the
Here and throughout, ||
denotes the Fuclidean

---

Published as a conference paper at ICLR 2025
input data and its representation is maximized. However, recall that [(X; Y) can diverge infinity, as discussed earlier, so is crucial to introduce augmentations.
InfoNCE Many conventional methods for self-supervised representation learning leverage InfoNCE loss (van den Oord et al. 2019) in conjunction with similarity measure ..... to for-
mulate the following contrastive objective (T'shanneen al. 2020: He et al., 2020 Chen et al.,
2020ah:
e (g.
LInfONCE
[Bp T(Q, K) - log Ep- exp(T(Q, K))],
LIsfNNCE - log
Li=i e ((an)))
where 1+ and P denote the distributions of positive and negative pairs of keys K and queries Q correspondingly. By selecting IP = ((xx)))((x)+ and P = P(xx & (rxx)))(p((+++) P(x)+2
we recover the Donsker- Varadhan bound (2) for our infomax objective I(f(X'); f(X) + 2). However, note that in 2 and (3) the supremum is taken over all measurable functions. In contrast, non-infomax methods typically employ a separable critic: T(g. k) = (ő(4), (*))), where @ and i are projection heads (Tschannen al., 2020 Chen al. 2020a). In special cases o sb Id, similarity of representations is measured via a plain dot product.
Despite this significant difference, our distribution matching paradigm allows us drop the supre- mum in (2) and (3) and establish a direct connection between DIM and traditional non-infomax contrastive SSRL methods:
Theorem 4.1 (Dual form of Gaussian distribution matching). Under the conditions of Theorem 3.4
f((());;(X)) + Z) 2 Exp
T (T
log Ex-
exp
(T (Tro.on))]
|//12 11y =x11
12
|| + |1vll |lyll2/(1 (1)+2
on (z. 2(1 + 0²)
(z, v)
202
o²
2
with the equality holding precisely when f is weakly invariant and f(X) (0.1))
Note that (x, y) is widely used as similarity measure, o² be interpreted as temperature, and the remaining part of the expression serves as regularization term
Covariance-basse based methods Barlow Twins (bontar al. 2021 and VICReg (Bardes et al., 2021) objective functions can be traced to information-theoreiicc terms either through Information Bottleneck (Tishby et 1999) or multi-view InfoMax (Federici et al. 2020) principles:
x(f(')) X) 1-0--Ahh xx(f(x)) + min (BarlowTwins)
((f(');x') 2 A((('))) + Elloga(f(X") ((((')) x")] + max (VICReg)
where > 0, X X" separate augmentation path, independent of + X', and q is probability density function of some distribution. In both cases, the representation entropy h((('') comes into play. Recall that normal distribution the maximum entropy distribution given first two moments (Theorem 2.1 As both methods employ covariance restriction terms in the respective objective functions, one can recover these methods by assuming the normality of f(X').
EXPERIMENTS
In this section, we evaluate our distribution matching approach on several datasets and downstream tasks. To assess the quality of the embeddings, we solve downstream classification tasks and cal- culate clustering scores. To explore the relation between the magnitude of injected noise and the quality DM, a set of statistical normality tests is employed. For the experiments requiring numer- evaluations visualization, other experiments, we use CIFAR- 10, CIFAR-100 datasets (Krizhevsky) 2009) and ImageNet (Rus- sakovsky al. 2015
MNIST handwritten digits dataset LeCun et al. (2010). For

---

Published as a conference paper at ICLR 2025
10
Estimated MI
D' Agosttnon Peanson Shapiro- Wilk Heezz-ZZirkler
0.8
MI. nats
0.6
p-value
4
0.4
2
0.2
1
0.4
0.9875
0.3
Clustering score
ROC AUC AUC
0.975
0.2
Gaussian naive Bayes k- N-Nearest neighbors Multi-layer perceptron
0.9625
0.1
Silhouette clustering score
0.95
0
0
3
5
7
9
10
11
12
13
14
Capacity C' = 2 log (1 - 1/0*). nats
with varying capacity C' _ log (1 + 1/0²). measured in nats (units of information based natural logarithms). The dotted line denotes the minimal capacity required to preserve the information about the class labels in f(X) Z. The dashed line represents the upper bound the mutual information (6) We run experiments for each point and report mean values and 99% asymptotic confidence intervals. InfoNCE loss is used approximate 2
Figure Results for MNIST dataset the Gaussian DM setup for
Multivariate normality and uniformity tests The key part of the experimental pipeline to estimate how much the distribution of embeddings acquired via the proposed method is sim- ilar to the multivariate normal or uniform distribution. To do this, we leverage D Agostino- Pearson (D' Agostino 1971 D' A'gostinno & Pearson 1973). Shapiro-Wilt (González-Estrada al.) 2022) univariate tests, and Henze-Zirkee variate To extend the univariate tests higher dimensions, utilize the fundamental property of the Gaussian distribution, which is the normality of any linear projection.
(Henze & Zirkler 1990, Trujillo-Ortiz et al.77 multi-
In practice, we sample random projectors from d d-dimensional sphere to perform a univariate nor- mality test. We also employ bootsrapping with small subsampling size to get low- oowvvariance averaged p-values and to smooth the transition from low to high p-values. Finally, we reduce the uniform distribution case to the Gaussian via the probability integral transform, see Section 3.3 We report the results the tests in [igure We also visualize the two-dimensional embeddings of the MNIST dataset in Figure Essentially similar experiments are also conducted CIFAR-1 10 dataset in Ap- pendix
Classification and clustering To explore the trade-off between the magnitude of the injected noise and the quality of representations. we evaluate various clustering metrics and perform downstream classification using conventional ML methods, such as Gaussian naive Bayes. k-nearest neighbors and shallow multilayer perceptron. In order to numerically measure the quality of clustering, we compute Silhouette score (Rousseeuw 1987). The results are reported in Figure
We also verify noise injection does not affect typical methods for SSRL such SimCLR (Chen) jet al. 2020a) and VICReg (Bardes et al. (021). To this end, we train both methods on CIFAR-I0 10 and CIFAR-100 datasets with varying degree of noise. The linear probing performance only slightly drops when noise magnitude is increased across a values (0.0, 0.1, 0.3, 0.5) as seen in Table I Finally, to validate this behavior on ImageNet- and ImageNet- during 100 epochs. Linear probing results on respective datasets reported in Table below. See Appendix for implementation and training details.
large-scale dataset, pre- preitrain VICReg with noise injection
8

---

Published as a conference paper at ICLR 2025
unbedding
(a) Normal distribution, 0.1
(b) Uniform distribution, 6 0.05
Figure 2: Visualization of two-dimensional representations of the MNIST handwritten digits dataset.
Table Linear probing: top-1 & top-5 accuracy (in %) on CIFAR-10/100 under noise injection.
CIFAR-10
CIFAR-100
top-1
top-5
top-1
top-5
SimCLR 90.83
99.76
65.64
89.91
SimCLR 0.1
90.96
99.72
67.03
90.49
SimCLR = 0.3
91.56
99.77
65.72
89.76
SimCLR 0.5
90.51
99.74
65.58
89.56
VICReg
90.63
99.67
65.71
88,96
VICReg 0.1
91.09
99.68
68.92
90.50
VICReg σ = 0.3
90.75
99.61
67.31
89.89
VICReg 0.5
91.02
99.75
66.52
89.66
Table 2: Linear probing: top-1 & top-5 accuracy (in %)
ImageNet under noise injection (VICReg).
ImageNet-100
ImageNet-Il 1k
top-1
top-5
top-1
top-5
0
72.18
+ 0.40
92.02 + 0.12
67.41 0.17
87.43 + 0.08
0.05
72.27
- 0.38
91.99 + 0.18
67.29
0.20
87.47 + 0.06
0.1
72.07
0.27
91.65
+ 0.13
67.30
0.13
87.43
+ 0.02
0.2
71.68
+ 0.50
91.61 + 0.24
67.19 + 0.12
87.32 + 0.09
An intuition behind the observed slight accuracy drop lies in that noise injection limits the bandwidth of representations, an effect similarly observed in (Lavoie et al.) 2022) and resolved by increasing the dimensionality of representations. For fair comparison with the base method, we keep the original dimensionality (5 12 and 2048 for ResNet-18 and ResNet-50, respectively).
Generation Generative adversarial networks [Goodfellow et al., 2014) produce samples from the underlying latent distribution, which is usually Gaussian. A common approach to make the gen- eration not purely random to introduce conditioning vector during the generation (Mirza & Osindero 2014 Larsen et al. 2016). However, the distribution of the conditioning vector is usually unknown except for specific cases, which hinders the unconditional (or partially conditional) gener- ation via the same model. Thus, model converting data to the corresponding conditioning vectors with known prior distribution is of great use. In Appendix [ we plug our encoder into conditional GAN and perform conditional and unconditional generation.

---

Published as a conference paper at ICLR 2025
6 DISCUSSION
In this paper, we have proposed and thoroughly investigated a novel and efficient approach to the problem of distribution matching of learned representations. Our technique falls into well- established family of Deep InfoMax self-supervised representation learning methods, and does require solving min-max optimization problems or employing generative models "post-hoc" to match the distributions. The proposed approach is grounded in the information theory, which allows for a rigorous theoretical justification of the method. In our work, we also explore the possibility of applying our technique to other popular methods for unsupervised and self-supprvised represen- tation learning. Consequently, we assert that the proposed approach can be utilized in conjunction with any other method for SSRL, provided that it can be formulated in terms of mutual information or entropy maximization.
To assess the quality of the representations yielded by our method, we (a) visualize the embeddings and run normality tests, and (b) solve a set of downstream tasks in various experimental setups. The results indicate the following:
1. Increasing the noise magnitude facilitates better distribution matching, but only up to a certain point; beyond that, the entire representation learning process begins to deteriorate due to insufficient information being transmitted through the noisy channel.
2. Experiments with MNIST and low-dimensional embeddings indicate the existence of moderate trade-off between the magnitude of the injected noise and the quality in down- stream classification tasks. However, experiments with other infomax-related methods for SSRL and higher embedding dimensionality suggest this influence being negligible.
3. Embeddings acquired via the proposed method can be used to condition generative mod- els, allowing both conditional and unconditional generation due to the distribution of the conditioning vector being known.
Future work As for further research, we consider elaborating on the dual formulation of infomax- based distribution matching. Theorem 4.1 suggests that Gaussian DM can be achieved through specific choice of the critic network T T(r, v)- We plan extending this result to other distributions. Additionally, consider conducting further experiments with generative models to comprehen- sively assess the quality conditional and unconditional generation, utilizing conditioning vectors obtained via the proposed method.
REFERENCES
Philip Bachman, R Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL attps::/prrocedinng.. httpeeedddda4211aaaa22222 htt.:////..........1.....010111.
filess/paper/201/// http://////111//1122221111.1////
Nikita Balabin, Daria Voronkova, Ilya Trofimov, Evgeny Burnaev, and Serguei Barannikov. Disen- tanglement learning via topology. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.). Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learn- ing Research, pp. 2474-2504. PMLR, 21-27 Jul 2024. URL httpss://prceeedinggggg////////t press/v235 ba aaliaiiiaae m244. html
Adrien Bardes, Jean Ponce, and Yann LeCun. Vicreg: o wtw.://mrrananeooovvvinnnnnnanaa regularization for self-supervised learning- arXiv preprint arXiv:2105.044906 2021.
Alexandra Bazarova, Evgenia Romanenkova, and Alexey Zaytsev. Normalizing self-supervised learning for provably reliable change point detection, 2024. URL httt:ttts///xxii..ott
abs/ 0ese/210... 13637
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and R Devon Hjelm. Mine: Mutual information neural estimation, 2021. URL
https://arxiss www.rxwww...va.oooo www.:/owggg..ooeooooooooo
1101104062
10

---

Published as a conference paper at ICLR 2025
A J Bell and T Sejnowski. An information maximiizaimiimiiiinnni approach to blind separation and blind deconvolution. Neural Comput., ((6):1129-113)) November 1995.
Vanessa Böhm and Uroš Seljak. Probabilistic autoencoder. Transactions on Machine Learning Research, 2022. doi: 10448550/ARXIV220600552 URL htttps://operrevieee net/
Johann Brehmer and Kyle Cranmer. Flows for simultaneous manifold learning and density esti- mation. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.). Advances in Neural Information Processing Systems, volume 33, pp. 442 453.. Curran Associates, Inc., 2020. URL https::/prooeesee
newrips pppe fi weweleeppr///0/0 file/ #ile/551 928892884488 F 111.22//111777000000000000000000 pdf
Bradley CA Brown, Anthony L. Caterini. Brendan Leigh Ross, Jesse C Cresswell, and Gabriel Ganem. Verifying the union of manifolds hypothesis for image data. In The Eleventh International Conference on Learning Representations, 2023. https:/: pttttps:///pereviieee /oppneevee..
net/forum?id=Rvee
Ivan Butakov, Alexander Tolmachev, Sofia Malanchuk, A. M. Neopryatnaya, and Alexey Frolov. Mutual information estimation via normalizing flows. ArXiv, abs/2403.02187, 2024a. URL
Ivan Butakov, Alexander Tolmachev, Sofia Malanchuk, Anna Neopryatnaya, Alexey Frolov, and Kirill Andreev. Information bottleneck analysis of deep neural networks via lossy compression. In The Twelfth International Conference on Learning Representations, 2024b. URL https:
Scott Chen and Ramesh Gopinath. Gaussianization. In T. Leen, T. Dietterich, and V. Tresp (eds.), Advances in Neural Information Processing Systems, volume 13. MIT Press, 2000. URL hhttps: /tttps://proceedddt ngs nwwrips cc paperp filess/pper/2000// 2TA.4//7EEE000888822222282288882 http///555dd--Parrepppee
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference Machine Learning, volume 119 of Pro- ceedings of Machine Learning Research, pp. 1597-1607. PMLR, 13-18 Jul 2020a. URL wwwp://che22..ccc .html
Ting Chen, Simon Kornblith, Kevin Swersky, Mohammad Norouzi, and Geoffrey Hinton. Big self- supervised models strong semi-supervised learners. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020b. Curran Associates Inc. ISBN 9781713829546.
Kai Lai Chung. course in probubility theory. Academic Press, 3rd edition, 2001. ISBN 0121741516; 9780121741518, 9780080522982; 008052298X.
Thomas M. Cover and Joy Thomas. Elements of Information Theory (Wiley Series Telecom- munications and Signal Processing). Wiley- Intersintesce,ece,, USA, 2006.
Victor Guilherme Turrisi da Costa, Enrico Fini, Moin Nabi, Nicu Sebe, and Elisa Ricci. solo- learn: A library of self-supervised methods for visual representation learning. Journal of Machine Learning Research, 23(56): 2022. htt.ttp:://...lloooo///////////e html
F. N. David and N. L. Johnson. The probability integral transformation when parameters estimated from the sample. Biometrika, 35(1/2):182-19, 1948. ISSN 00063444. URL http://wwet jstor.rroooo org/ http//tablee.. e 2232 7332638
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HIT
11

---

Published as a conference paper at ICLR 2025
2019, Minneapolis, MN, USA, Jume 2-7, 2019, Volume 1 (Long and Short Papers), pp. 4171-
4186. Association for Computational Linguistics, 2019. doi: 10.18653/V1/199- 1423. URL https://doi httpp/////.bblor//00...////////0
M. D. Donsker and S. R.S. Varadhan. Asymptotic evaluation of certain markov process expectations for large time, iv. Communications on Pure and Applied Mathematics, 36(2):1832112, March 1983. ISSN 0010-3640. doi: 10. (0000//pa.360066//
B. Duong and T. Nguyen. Conditional independence testing via latent representation learning. In 2022 IEEE International Conference on Data Mining (ICDM), pp. 121-130, Los Alamitos, CA. USA, dee 2022. IEEE Computer Society. doi: D0.100//CDMM444422022222
URL https:
R. B. D' Agostino. An omnibus test of normality for moderate and large sample size. Biometrika, 58(1):341334, 1971.
R. B. D'Agostino and E. S. Pearson. Tests for departure from normality. Biometrika, 60(1):613-622, 1973.
Linus Ericsson, Henry Gouk, Chen Change Loy, and Timothy M. Hospedales. Self-supervised rep- resentation learning: Introduction, advances, and challenges. IEEE Signal Processing Magazine, 39(3):42 62, 2022. doi: 10.11109/MSPP22211133333.
Marco Federici, Anjan Dutta, Patrick Forré, Nate Kushman, and Zeynep Akata. Learning robust representations via multi-view information bottleneck. In International Conference on Learning Representations, 2020. URL /ttppe/s/pppeeeeeeeeeeeeeeeereer
Elizabeth Genzzlez-Estradie José Villaseñor, and Rocío Acosta-Pech. Shapiro-wilk test for multivariate kkew-normality. Computational Statistics, 37:1-17, 09 2022. doi: 10.1007/ -00180-021-01188-y:
Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks, 2014. URL https:
//axxiv.rrgaab//2...
2661
Wes Gurnee and Max Tegmark. Language models represent space and time. In The Twelfth Interna- tional Conference on Learning Representations, 2024. URL htttps://operreview net/
Forum?ide 1
Robert M. Haralick, K. Shanmugam, and Its' Hak Dinstein. Textural features for image classification. IEEE Transactions on Systems, Man, and Cybernetics, SMCC336))600-621 1973. doi: 10.1 109/ ISMC.1973.4309314
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition, 2015. Uttps::::://arxs http:///xxxvv.r////// abs /1522. 03385
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In 2020 IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 9726-9735 2020. doi: 10.1109CCYP44220002000..9999
Norbert Henze and Bernd Zirkler. A class of invariant consistent tests for multivariate normality. Communications in Statistics-theors and Methods, 19:3595-36177 1990. URL https::/appi. semanticschole wrg/CorppsID:: :ttp://CopuuIDDD12222222 120328121
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochre- iter. Gans trained by two time-scale update rule converge to local nash equilibrium. In I. Guyon, U. Von Luxburg, Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gamett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL
files/
12

---

Published as a conference paper at ICLR 2025
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerehner. beta-VAE: Learning basic visual concepts with a constrained variational framework. In International Conference Learning Representations, 2017.
net.
Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchildon, Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estima- tion and maximization. In International Conference on Learning Representations, 2019. URL
httts:://ppenreve/ http://wwwee/fooommte/eeme
net forum d= r 3
Jonathan Ho and Tim Salimans. Classifier-free diffusiin guidance. In NeurIPS 202/ Workshop Deep Generative Models and Downstream Applicatton,, 2021. URLitttp:://ppeeeviii
net/forunt (tw.oruuunndddoAAAkxaaeXbbb
Chin-Wei Wei Huang. David Krueger, Alexandre Lacoste, and Aaron Courville. Neural autoregressive
flows. In International conference on machine learning, pp. 2078-2087. PMLR, 2018.
Minyoung Huh, Brian Cheung, Tongzhou Wang. and Phillip Isola. Position: The platonic represen- tation hypothesis. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings Machine Learning Research, pp. 20617-20642. PMLR, 21-27 Jul 2024. URL v235/huh24a. html
htt.eeeesins.mm..reessss/// www.m.rppp... press/
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift, 2015.
Priyank Jaini, Kira A. Selby, and Yaoliang Yu. Sum-of-squares polynomial flow. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Leurning Reseurch, pp. 3009-3018. PMLR, 09-15 Jun 2019.
Amor Keziou. Dual representation of d-divergences and applications.
Comptes Ren- doi: https://oi.orr/10. 1016/31661--033300000022222 URL httpp::///eett https::/wwww.ccieeeee.... www.sciienedirect.t htt.t/ccco/cccccceeeeee
dus Mathemaaique, 336(10) 2003. ISSN 1631-073X.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
Diederik Kingma and Max Welling, Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun (eds.), 2nd International Conference on Learning Representations, ICLR 2014, Banff AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014. URL httt:://appxit/ org/
abs/ Doss/3123 2. 6114
Ivan Kobyzev, Simon Prince, and Marcus A. Brubaker. Normalizing flows: An introduction and review of current methods. IEEE Transactions Pattern Analysis and Machine Intelli- gence, 43:3964-3979, 2020. URL httph: /tttps:://p... wewanticscoola..rr....ooo.. org/ orw.aa/pooooCCoCCI....... [089910744]
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Anders Boesen Lindbo Larsen, Søren Kaae Sønderby, Hugo Larochelle, and Ole Winther. Au- toencoding beyond pixels using a learned similarity metric. In Maria Florina Balcan and Kil- ian Q. Weinberger (eds.), Proceedings of The 33rd International Conference on Machine Learn- ing, volume 48 of Proceedings of Machine Learning Research, pp. 1558-1566, New York, New York, USA, 20-22 Jun 2016. PMLR. URL httph:ttt//cmmm http:///ooooee...............ee
1arsen16. html
Samuel Lavoie, Christos Tsirigotis, Max Schwarzer, Ankit Vani, Michael Noukhovitch, Kenji Kawaguchi, and Aaron Courville. Simplicial embeddings in self-supervised learning and down- stream classification. arXiv preprint arXiv:2204.00 00066, 2022.
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. 2010.
Available: http:///amm.ecccc.oom//e////////
13

---

Published as a conference paper at ICLR 2025
R. Linsker. Self-organization in a perceptual network. Computer, 21(3):105-117, 1988. doi: 10. 1109/2.36.
Alireza Makhzani and Brendan Frey. k-sparse autoencoders. In Yoshua Bengio and Yann Le- Cun (eds.), 2nd International Conference Learning Representations, ICLR 2014, Banff. AB, Canada, April 14-16, 2014, Conference Track Proceedings. 2014. URL http:: org/ hbs/abb//ie 111//31122.5 5663
Alireza Makhzani, Jonathon Shlens, Navdeep Jaitly, and Ian Goodfellow. Adversarial autoencoders In International Conference on Learning Representations, 2016. URL htttp::/axxiv..rg//
abs 11111 05644
David McAllester and Karl Stratos. Formal limitations on the measurement of mutual informa- tion. In Silvia Chiappa and Roberto Calandra (eds.), Proceedings of the Twenty Third Interna- tional Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pp. 875-884. PMLR, 26-28 Aug 2020. URL https: petpp::/prroceedinn:
mlr. w...press//. v ww0. www.mmaaaaeeeeeeeee 2htmmm
Leland Mclnnes, John Healy, Nathaniel Saul, and Lukas Großberger. Umap: Uniform manifold approximation and projection. Journal of Open Source Software, 3(29):861, 2018. doi: 10.
21 0006. (RLttts:: www.doioorgg/// 2et......///ssssse.. 00861
Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word represen- tations in vector space, 2013. htt/://ttpee//rraaa.rrro/////.33
Mchdi Mirza and Simon Osindero. Conditional generative adversarial nets, 2014. URL https:
7/arxiv. org abs /////1111111111111
Michael Moor, Max Hom, Bastian Rieck, and Karsten Borgwardt. Topological autoencoders. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Ma-
chine Learning, volume 119 of Proceedings of Machine Learning Research. pp. 7045-7054. PMLR, 13-18 Jul 2020. URL
mml. htt.://.......reeee/////////0000
html
Andrew Ng. Sparse autoencoder, 2011. 1.
URL
cs294a/ http///sssssuuuoonccer....
XuanLong Nguyen, Martin J. Wainwright, and Michael I. Jordan. Estimating divergence functionals and the likelihood ratio by convex risk minimization. IEEE Trans. Inf. Theor. 56511))5847-8811 nov 2010. ISSN 0018-9448. doi: 10.1109/TT.201002066877 URL / 11C9/TIT. 2010. 2068870
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lak- shminarayanan. Normalizing flows for probabilistic modeling and inference. Journal of Ma-
chine Learning Research, 22(57):1-64, 2021. URL / http:///ll..crr//pppr////////
19-1028 html
Y. Polyanskiy and Y. Wu. Information Theory: From Coding Learning. Cambridge Univer- sity Press, 2024. ISBN 9781 108832908. URL htttps://bookss www.gel.loboooo.e.... CYSoOAEACAAJ
Alec Radford, Karthik Narasimhan, Tim Salimans, and llya Sutskever. Improving language under- standing by generative pre-training, 2018.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and llya Sutskever. Learning transferable visual models from natural language supervision, 2021. URL
httts::/erxiv..erm alw.g///////// /2103 .00200
R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image syn- thesis with latent diffusion models. In 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). pp. 10674-10685. Los Alamitos, CA, USA. jun 2022. IEEE Computer Society. doi: 00.1099CCVP222682222200444
URL https://doi.
01042
14

---

Published as a conference paper at ICLR 2025
Peter J. Rousseeuw. Silhouettes: graphical aid to the interpretation and validation of clus- ter analysis. Journal of Computational and Applied Mathematics, 20:53-65, 1987. doi: 00.101600377-44777779991333
Olga Russakovsky, Jia Deng. Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang. Andrej Karpathy, Aditya Khosla, Michael Bernstein, al. Imagenet large scale visual recognition challenge. International journal of computer vision, 1 11:2111-222, 2015.
Peter Sorrenson, Felix Draxler, Armand Rousselot, Sander Hummerich, Lea Zimmermann, and Ullrich Koethe. Lifting architectural constraints of injective flows. In The Twelfth International Conference on Learning Representations, 2024, / wpwpr//eee/.......eeoooo....o.e. net forum??
id=kBNIx4Biq4
M. Spivak. Calculus On Manifolds: A Modern Approach To Classical Theorems Of Advanced Calculus. Avalon Publishing, 1965. ISBN 9780805390216.
Michael Tipping and Christopher M. Bishop. Mixtures of Probabilistic Principal Component Analyzers. Neural Computation, (12)):43-442 02 1999. ISSN 0899-7667. doi: 10.1162/
0899766699300067733 URL htRl:ttp::///
org. 10. 11t.::..1111196699981881
Naftali Tishby, Fernando C. Pereira, and William Bialek. The information bottleneck method. In Proc. of the 37-th Annual Allerton Conference on Communication Control and Computing, pp. 368-377, 1999. URL /bbrokenrr4ccleseeee nee. com/
tishay99inforraat.om... html
Antonio Trujill-Ortiz, Rafael Hernandez- Walls, and Barba-Rojo. Hzmvntest. henze-zirkler's multivariate normality test., 12 2007.
Michael Tschannen, Josip Djolonga, Paul K. Rubenstein, Sylvain Gelly, and Mario Lucic. On mu- tual information maximization for representation learning. International Conference on Learn- ing Representations, 2020. [ttpsttts/////eerriiiee net htte:/ooommnndddd.om.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding, 2019. URL https:://axxii wwtp///////bb...ooo.
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of Ma- chine Learning Research, 9(86):2579-26) 2008. URL /t jmlr. wrgppapers/v9/
wardermaaten html]
Petar Veličković, William Fedus, William L. Hamilton, Pietro Liò, Yoshua Bengio, and R Devon Hjelm. Deep graph infomax. In International Conference on Learning Representations, 2019.
net Sete:/ooommidddd.233100 URL httpt:::/ http::///pgaaaeeeeeeeee///rsaass
Oriol Vinyals, Alexander Toshev, Samy Bengio, and Dumitru Erhan. Show and tell: A neural image caption generator. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 3156-3164, 2015.
Yang You, Igor Gitman, and Boris Ginsburg, Large batch training of convolutional networks, 2017. URL https::/aaxxi......tmm wrw.vv..rrrgaaaa /ogg/bbb//70 03888
Xuechu Yu. Leveraging superfluus information in contrastive representation learning, 2024. URL https::/aaxiv..rr//asssss
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 learning, pp. 12310- 12320. PMLR, 2021.
15

---

Published as a conference paper at ICLR 2025
A COMPLETE PROOFS
Theorem 2.1 (Theorem 8.6.5 in Cover & Thomas (2006), Let X be a d-dimensional absolutely continuous random vector with probability density function p, mean m and covariance matrix E. Then
I(X) h (m,, Σ)) DRL. (p (mm,, Σ)),
le (W Σ)) = -2 log ((2me)* del Σ),
where (m, 2) is a Gaussian distribution of mean and covariance matrix Σ.
Proof of Theorem 2.1 As for any m E Rd it holds h(X - m) = (,(), let us consider a centered random vector X. Denoting the probability density function of (0,, 22) by øx, we have
[ p(ar) log p(z) dx -(()) / p(x) log ox (2) dr. (5(
DKL (p (O Σ))
Now, consider the second term:
J (()) log (x((x) de = const + Ex X
1 (xiE-"E(XXX") const + (T(E-11)) const +
const + 2
const + 2 EA X- / ggs (xr(2)) log (x)( (x)) dx.
Here, in the second line, the cyclic property of the trace is used.
Theorem 2.2. Let X be an absolutely continuous random vector with probability density function p and supp X C S, where S has finite and non-zero Lebesgue measure U(S). Then
h(X)= h (U(S)) - DRLL. (p U(S)),
h (U(S)) log (SS),
where U(S) is uniform distribution on S.
Proof of Theorem 2.2
DKL (p U(S)) / p(x) log p(x) dr L/p(S)
= =-h(X) + / p(x) p(z)) dx = -h(X) + log u(S) = -h(X) A(U(S)).
Lemma 2.4. Let X and Z be random vectors of the same dimensionality, independent under the conditioning vector Y. Then
h(X + Z (Y) h(z ( ((;;
1(Z YY),
with equality if and only if there exists a measurable function g such that X = g (Y))
Proof of Lemma 2.4
By the definition of mutual information between and X + Z conditioned on
Y. we have
I(X; Z Y) h(x + Z Y) h(X + Z X,
where 1(x + Z X,Y) h(Z Y) due to the independence of X and Z given Y.
Now, let us prove that [(X; X Z Y) = if and only if ig X = gY)) Firstly, X = gY)),
I(X: + Z Y) = (Vg(Y g 2 Y) = [(0; Y) = I(0) =
Hereiiafter, when comparing random variables, mean equality inequality "almost sure".
16

---

Published as a conference paper at ICLR 2025
Next, consider [(X; X + Y)) =0. Thus, X is conditionally independent both of X Z and Z, which allows for the multiplicative property of characteristic functions to be used:
E (1
y) = E ( (cll1-x+(2+x)))
x) E
(e (e y) E () r) E (elt..x)) r)
= E (clt..:)) (cl(x)) E (elte.2)
= E ( (te) v) | (tt2 r)
Note that there exists =(Y) > such that E (eftt.2) Y) * 0 for ||||| < =(Y) (non-vanishing property). Consequently, [E (elte.x) Y) = holds for ||t|| < eY)) This implies that the con- ditional distribution of X given Y is 6-distribution (Theorem 6.4.7 in (Chung 200T)), leading to
X E(X Y), where g(Y) & IE(X Y) -measurablee function.
Lemma 3.1. Consider the following Markov chain of absolutely continuous random vectors:
f(X) 7 - X X" - {(x'),
with Z being independent of (X, X'). Then
f(X)) + Z) h(f(X) + Z) h(Z) T(((X) + f:f(X) ((X'))
(5)
Proof of Lemma 3.7
From the definition of mutual information, we have
[() Z) -X((X) + Z) M(((X)((X))))))))))))(( + Z ((X'))
We apply Lemma
2.4 to rewrite the second term
xx + Z f(x)) = h(Z
+ + + 220)))+11((7)) ffx))) =M(Z) + (X((X):((X)+ 2 (Ix))),
where the independence of Z and f(X') is used.
Lemma 3.3. Under the conditions of Lemma 3.7 let P(X X' X) 2 > 0. Then, I(f(X) +
fx) f(()) = 0 precisely when f is weakly invariant to X + X'.
Proof of Lemma 3.3 According to Lemma 2.4 l(((X) + 2: | f(X') = 0 if and only
f(X) = f(((X )) for some g. Thus, weak invariance implies L(f(X) + | = 0.
On the other hand, if f(X) = x(f(x')).
P(f(X) E (g fx) XX) 2 xx
2a 0.
Note that the predicate P(X) & "f(X) E (g- f(x)" is not random when conditioned on X. Thus, (()(X) X) = o(X), where is a function taking values in {0, As x(X) a O.
x(X) = 1. This implies (x)( = f(X) almost surely.
Theorem 3.4 (Gaussian distribution matching). Let the conditions of Lemma
3.3 be satisfied. As- sume 2 ~ (O0 I21), (X) = and Var (f(X)) = 1 for all i E {l...., d}. Then, the mutual information (X'(X'); f(X) + Z) can be upper bounded as follows
d
T((('); ((X) Z) log
1+
(6)
with the equality holding exactly when f is weakly invariant and f(X) ~ ~N(0,1) I). Moreover, Z) 1(f(x';; + Z) d logo.
DKL (f(x) (IN(O,1) I(2;
Proof of Theorem
3.4
Applying the result from Lemma B.1 gives us:
DKL (f(x) Z (0, (1 + +)))) _ d log 1+
1+ (ラ) -
-IX(X));((xx+ Z).
2
Since KL-divergence is non-negative, we obtain the desired bound. Furthermore, equality holds exactly when ʃ is weakly invariant and f(X) (O,OT), I), as discussed in the proof of Lemma B.I
The second inequality involving DKL. (x(X) (O, T)), follows from Corollary B.3
17

---

Published as a conference paper at ICLR 2025
Theorem 3.5 (Uniform distribution matching). Under the conditions of Lemma 3.1 let Z U(I=E;El") and supp f(X) - [0:1]". Then, the mutual information f( + Z) can be upper bounded as follows
1 2E
I(f('); f(X) + =) dlog
1+
with the equality if and only 1/e E N. is weakly invariant, and f(X) U(A), where the set A = {0. 22, Se...... contains (1/(2e) + elements.
Moreover,
DKL U(((1))) <
- dlog(2e).
Proof of Theorem
3.5
Proceeding similarly to the Gaussian case, we apply Lemma B.4 yielding
1 - I(() + 2),
DKL. (J(X) + 2 2| + +|() 5 d
25
where the left-hand side is non-negative, so we obtain the claimed inequality. The conditions for equality also established through Lemma B.4
To prove the second upper-bound, concerning DKL (f(X) U((0; 111)), suffices to use the result from Corollary B.6
Theorem 4.1 (Dual form of Gaussian distribution matching). Under the conditions of Theorem 3.4
I(((X); + Z) 2 Ept
T (Tr...m)
log Ep-
[exp (Tr
||
11
||0||² + |1y1| [1]l2((1+
T (I.
(z. v)
2(1+00 02)
202
with the equality holding precisely when f weakly invariant and ~ (ox)~N(0.1)))
Proof of Theorem 4.1 By Remark 2.5
know that the equality holds if and only if T2 (I, y) = PMI x. 3) const.
Now, for independent Y ~ (0. and Z o2)
(M log PY Pr
x)) log Py +2
|//|
1 x1 log
1 0²
1+
2(1 02)
20
(1, + I(Y;Y + 7).
Thus, the equality holds if and only if f(X) precisely when weakly invariant and f(X) (0, (see Theorem 3.4).
f(X)+2) = _ log( + +1/02), which, in turn, holds
B SUPPLEMENTARY RESULTS ON DISTRIBUTION MATCHING
In this section, we explore the connection between the proposed infomax objective and the problem of distribution matching, formulated in terms of the Kullback-Leible divergence.
GAUSSIAN DISTRIBUTION MATCHING
Lemma B.I. Assume the conditions of Theorem 3.4 are satisfied, then for Gaussian distribution matching, have
DKL (f(X) Z (0, (1+ I d
log 1+ 0² -2) -T(((X'): 2),
with equality holding exactly when f is weakly invariant and S(X) N(O.1).
18

---

Published as a conference paper at ICLR 2025
Proof. From Lemma 3.1 obtain
I(f()): f(X)++ 2) h(f(x) + 2) M(V(0, o²1) -(f()))((((())))))))))) + 2:f(x) f(()))
Using Theorem 2.1 we can rewrite the first term, which yields
I((('); Z) h(V(m, E)) - M()((0,021)
DKL. (ʃ(X) + Z [1 +
22/+
where and the mean and covariance matrix of f(X) +
To bound the KL-divergence, note that the conditional mutual information is non-negative:
DKL (f(X) Z (m,, Σ)) - ((M(m, E)) h(N((,, [(f(X"); + Z).
Equality holds exactly when T( Z; (X) f(X)) 0, which is equivalent to f being
weakly invariant (see Lemma 3.3
Next, we estimate the difference between the entropies by observing that
hm Σ))
Σ h(N(m, Var(f(X)) of)) d (NN(0,1
with the equality holding if and only if 2 is diagonal, which implies K = 1 since Var(f(X); for all E (I.....d.. Finally,
d [log(1 + + - -log 0²1 = log [ (1+
d (A[(0,++0)) NMV(), +2)
which proves the claimed inequality.
Lemma B.2. Under the conditions of Theorem
3.4 the following holds:
DKL (f(X) (1, (1+ + DAL (f(X)
[I (1,
= DRL (f(x) (I,
Z) log (1 (1+o²).
Proof. The left-hand inequality follows directly from the data processing inequality for the Kullback-Leibler divergence (Theorem 2.17 in (Polyanskiy & Wu. 2024)).
To establish the right-hand side, we first apply Theorem 2.1 and then use the independence of (X) and Z, yielding (x())) ((f(X)== + / 2). Thus, we have
DKL (I = h(N((,1) h(f(x)((x)))))))))))))(()))))h + 2 Z).
Next, by using the definition of mutual information and again applying Theorem 2.1 one can write
h(f(X) + Z Z) h(f(X) Z) (X) Z; Z)=
1(N (1+ DRL. (f(X) [I (1+ (((
Finally, substituting this into the expansion for DKL (f(X) (O) and noting that h(N((0,
\ (1+ d²)1)) - log(1 + o²) completes the proof.
Corollary B.3. In the Gaussian distribution matching setup (Theorem 3.4 we have
DKL (f(X) (I,(O,)) & (X) + - + Z) - d a.
Proof. This follows directly from combining Lemma B.1 and Lemma B.2
19

---

Published as a conference paper at ICLR 2025
UNIFORM DISTRIBUTION MATCHING
Lemma B.A. Let the conditions of Theorem 3.5 hold, then the following bound applies for uniform
distribution matching:
DKL (S(X) + ( 1((())))+(())) S d dlog 1+ 1
-100)11))) +x)+
25 -10(()));
with the equality if and only 1/s N. weakly invariant, and f(X) U(A), where the set A = (0, 2e, 4.. contains (1/(2E) elements.
Proof. Similarly to the proof of Lemma B.I we use the decomposition of the infomax objective
from Lemma 3.1 Afterward, we apply Theorem 2.2 to the term h(f(X)+ 2):
I(f(X);;()) + Z) = [[U([[=; MACK--EE))))))))(
DKL (f(X) + Z [|-[[ - - - f(x'))
Therefore, the KL-diieegeence can be bounded as:
DKL (f(X) + | U([ 174) < > h(U([=e;e + 1j²)) - M[U([- T(J(X';; 2)
1 2s ((((x)) +
dlog 1
with equality achieved if and only if the mutual information between f(X) + Z and f(X) condi- tioned on f(X') is zero, which occurs precisely when f is weakly invariant (see Lemma
Next, show when the equality holds. The probability density function of the sum of independent random vectors f(X) and Z is given by the convolution:
d aLto
J P)o0(())p2(z-x)ds II
1
Σ aEA
8(4 -a) a)
dæ
dz.
1/(22)
2s
(1+
Here, we used the independence the components, with Z; uniformly distributed on [-a,e] and f(X); uniformly distributed over the discrete set A = (0, 2s, 4s.
Therefore, (f(X) + Z) ~ U[[=e;1 + a)d). Given this, one can calculate the mutual information explicitly:
1
IU(();;(X)) + 2) = +((())((X))))))))))))))()) 2) M(2) = d log
1
25
This concludes the proof.
Lemma B.5. Under the conditions of Theorem
3.5 the following holds:
DKL (f(X) U([0; 1)() = DRL +
Proof. One can build on the reasoning from Lemma B.2 by applying Theorem 2.2 to express the KL-divergence between f(X) and U( 1]d) follows:
DKL [I ) = ((U(( 1}1)) h(f(x) + Z -hf(((() + Z).
Using Theorem 2.2 again, the conditional entropy can be expressed as
M((X)+ 2 Z) -- (UU([=;; + 1)4) - DRL. (f(x) + U[=
T(f(X) + + 2:
To conclude it remains to calculate h(U[[-E; 1) dlog + 2z).
Corollary B.6. In the uniform distribution matching setup (Theorem 3.5 have
DeL (x(X || U((0 -
Z) dllog(2e).
Proof. Applying Lemma B.4 alongside Lemma
B.5 is enough.
20

---

Published as a conference paper at ICLR 2025
C DETAILS OF IMPLEMENTATION
For experiments on MNIST dataset, we use simple ConvNet with three convolutional and two fully connected layers. A three-layer fully-connected perceptron InfoNCE loss. We provide the details 3 We use additive Gaussian noise with a = 0.6 as an input augmentation. Training hyperparameters are follows: batch size 1024, 2000 epochs, Adam optimizer (Kingma & Ba 2017) with learning rate 10-3
critic network for the
Table 3: The NN architectures used to conduct the tests on MNIST images in Section
NN
Architecture
Conv2d(1, 32, ks=3), MaxPool2d((2), BatchNorm22, Lea.k
ConvNet,
x1:
Conv2d(32, 64, ks=3), MaxPool2d(2), BatchNorm2d, LeakyRe).U(0.00) Conv2d(64, 128, ks=3), MaxPool2d(2), BatchNorm2d, LeakyReLU(00001) Dense(128, 128), LeakyReLU(0.00), Dense(128, dim)
24 24
x1:
images
x1:
Critic NN, pairs of vectors
Dense(dim dim, 256), LeakyReLU(0.00)) Dense(256, 256). LeakyReL.U(0.00), Dense(256, 1)
The results on CIFAR datasets (Krizhevsky) 2009) in Table 1 were obtained with the standard con- figuration of SSL methods. Namely, we use ResNet-18 18 (He et al. 2015) backbone. Projection head for SimCLR consists of two linear layers, for VICReg - 3 layers. Respective configurations are [2048, 256] and [2048, 2048, 2048), meaning embedding dimensions are 256 and 2048, respectively. We apply standard set of augmentations:
Pretrain Transform: Compose RandomResizedCron(
size= 32)
scale =0.08 1.0)
ratio =0.75. 1.3333333333).
interolatioon//teepllttiiiool Btt.ttooooooee.BBBBB.C... BICUBIC,
antialias=True
RandomApply(
ColorJitter(
1igghtees=(0.66 1.4),
contrast=(0.6, 1.4).
saturation= 1.2).
hue=(-0.1, 0.1)
RandomGrayscalc (p=022)
Gaussian (p=0.0) (pz
RandommHorizonntaliiipppp:.3:3 ToTensor () Normalize
mean=[0.499 0.4822, 0.4465].
std=[0.247, 0.2435, 0.2616].
inplace=False)
TestTransform: Compose ToTensor Normalize
mcan=(0.4914. 0.4822, 0.4465),
std=[0.247, 0.2435, 0.2616]. inplace=False
21

---

Published as a conference paper at ICLR 2025
Training hyperparameters are as follows: batch size 256, 800 epochs, LARS optimizer (You et al., 2017) with clipping, base learning rate 0.3, momentum 0.9, trust coefficient 0.02, weight decay 10 For SimCLR, we use temperature 0.2, for VICReg standard hyperparameters (25, 25, 1).
The source code available on GitHub repository
D CONDITIONING GENERATIVE ADVERSARIAL NETWORKS
In this section, we leverage our method to generate conditioning vectors for a conventional cGAN setup (Mirza & Osindero 2014). We use the two-dimensional Gaussian embeddings of the MNIST dataset, acquired from the noise level 0.1. For conditioned generation, we get embeddings from a batch of original images. For unconditioned generation, embeddings are sampled from N(,, I). The results are presented Figures and
E ADDITIONAL EXPERIMENTS ON CIFAR-10
Similarly to embedding MNIST into two-dimensional space, we run experiments with CIFAR- 10. We use a ResNet-18 architecture augmented with an additional linear layer to map the 512- dimensional output to R², along with batch normalization layer without learnable parameters (i.e., "affineFFalse"). We experimented with both contrastive (SimCLR) and non-contrastive (VICReg) objectives, observing similar results across both methods. The experimental pipeline closely follows the theoretical setup:
Unaugmented Input: The unaugmented image X is fed to the encoder, yielding f(X), followed by noise injection to obtain f(X) + Z.
Augmented Input: The augmented image X' is fed to the encoder to obtain f(X'), without noise injection.
Both outputs are then processed by the projection head and loss function (c.g.. InfoNCE). The noise Z is sampled from a normal distribution with standard deviation We conducted pre-training across various noise magnitudes with E (0.0, 0.01, 0.025, 0.05, 0.1, 0.2, 0.2658, 0.5). For each a, we ran training with 5 different random seeds with base__r r 0.1 for 400 epochs (visualized Figure 5). To speed up convergence we additionally trained models (a E (0,0.00 0.05, 0.13) with base_lr 0.5 for 800 epochs (visualized in Figure
After pre-training, we evaluated downstream performance using ROC AUC and performed normal- ity tests on the learned 2D representations from total of 40 trained models (that were trained for 400 epochs) (see Figure depicting similar results to our previous MNIST experiments (Figure [
F ADDITIONAL EXPERIMENTS ON IMAGENET
To better assess the accuracy-DM trade-off in complex setups, additional experiments with the Im- ageNet datasets are conducted with VICReg loss.
For ImageNet-100, as backbone we use ResNet-18 followed by batch normalization layer without Projection head consists
learnable parameters, i.e. BatchNormid (512 2, affine=Falsed of standard MLP sequence: Lineir (512, 2048), nn Bat chNormld (2048), nn. Rn. (), nn. Linear (2048, 2048), , nn Bn..atchNormle (2048), nn. ReLR () nn. Linear (2048, 2048).
For ImageNet-1k, we ResNet-50 as backbone architecture, followed by batch normalization layer without learnable parameters (dim=2048) The projection head is comprised similar to the above setup of 3-layer MLP with hidden dimension 8192.
In both setups, we run pre-training for 100 epochs and conduct linear probing afterwards closely following standardized protocol for augmentation and training. for ImageNet- batch size and
22

---

Published as a conference paper at ICLR 2025
7
7
7
7
2
7
6
Le
6
4
/
0
0
0
O
O
4
9
4
4
9
9
9
9
9
9
(a) Encoded images
(b) Generated images
Figure 3: Results of conditional generation.
9
8
8
4
4
9
0
S
2
7
8
5
3
3
9
6
7
S
O
3
7
0
4
8
7
5
9
3
Figure 4: Results of unconditional generation.
23

---

Published as a conference paper at ICLR 2025
(a) No noise injection
(b) Normal distribution, 0.05 (c) Normal distribution. 0.1
Figure 5: Visualization of 2D representations of the CIFAR-10 dataset, 1r=0.1, 400 epochs.
(a) No noise injection
(b) Normal distribution, a = 0.05 (c) Normal distribution, o = 0.1
Figure 6: Visualization of 2D representations of the CIFAR-10 dataset, Ir=0.5, 800 epochs.
D' A'Aosstin- Pearson Shapiro-Wilk
0.8
Henze-Zirkler
p-value
0.6
0.4
0.2
0
0.88
0.86
ROC AUC
0.84
8
0.82
k n-neersst neighbors
0.8
3
4
5
6
7
8
9
10
Capacity C o log + 1/02),
Figure 7: Results for CIFAR1O dataset in the Gaussian DM setup for d = 2 with varying capacity C of log (1 + 1/0²), measured in nats (units of information based on natural logarithms). The dotted line denotes the minimal capacity required to preserve the information about the class labels in f(X)= + Z. We run 5 experiments for each point and report mean values and 99% asymptotic confidence intervals. InfoNCE loss is used to approximate
24

---

Published as a conference paper at ICLR 2025
other (e.g. optimizer, data augmentation) hyperparameters set to specified in the original paper Bardes et al. 2021), however, evaluation batch size is set to 2048 for faster results. For ImageNet- 100, batch size 256, optimizer's trust coefficient 0.02, base learning rate = 0.5, weight decay 1 1e-4, 4, inspired by da Costa et al. (2022). The resulting accuracies are reported in Table across random seeds.
25


Reviews:
[{"review_id": "qFWBJcy2Fc", "reviewer": "Reviewer_Vzb6", "summary": "This work introduces  noise injection to shape a self-supervised learning represenation as Gaussian or Uniformed representation. In short, by injecting a Guassian noise and restricting the convariance of f(x), f(x) tends to be gaussian distirbutioned. By injecting an uniformed distribution within a narrow range $[-\\epsilon, \\epsilon]$ and restricting the range of f(x) (e.g. within [0,1]), f(x) tends to be uniformly distributioned.\n\nThe reason for the gaussian distribution comes from the fact that the Gussian distribution has the max entropy of all distributions with a fixed covariance. The reason for the uniform distribution comes from uniform distribution U(a, b) has the max entropy of a distribution within range [a,b]. \n\nApart from the theorical analysis, this work also present interesting plots to demenstrate the two kinds of learned representation, i.e. Gaussian and Uniformed.", "questions": "- what is the meaning of \"nats\" in Figure 1 caption \"Capacity, nats\"? Does it refer to $C = d/2 log(1 + 1 / \\sigma^2)$?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The intuition of two noise injection, Gaussian and Uniform, methods are resonable. The reason for the gaussian distribution comes from the fact that the Gussian distribution has the max entropy of all distributions with a fixed covariance. The reason for the uniform distribution comes from uniform distribution U(a, b) has the max entropy of a distribution within range [a,b]. \n\n- This work provides a theorical proof of the effect of these two noise injections. \n\n- This paper provides rich background knowledges. \n\n- The experiments in this paper (cifar, mnist) demenstrate the usefulness of the two noise injection methods.", "weaknesses": "- line 236-239 explains why injecting a gaussian noise helps learn a gaussian representation. This sentence is helpful for understanding. It would be good to have a similar short explaination before the proof of uniform noise injection. For example, the reason of small $\\epsilon$, discrete support $A$, and the entropy of uniformed distribution. \n\n- Large experiments beyond cifar and mnist are helpful to convince readers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gkCFXncayO", "reviewer": "Reviewer_y7PL", "summary": "The paper propose to add gaussian or uniform noise to the output of the encoder in contrastive methods.", "questions": "* Is the linear probe trained on the output of the encoder or on some intermediate layer?\n* Are the gains observed in Table 1 transferrable when training on ImageNet?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper defines the notation and robustly use it throughout the paper making it easy to follow along as a    reader. I appreciate the clarity and reading the method and the experimental section. I was also pleasantly  surprised of the novelty of the idea proposed, while being very simple.", "weaknesses": "The main weakness of the paper is that the experiments are conducted on simple and sometimes trivial           benchmarks. The raw input space of MNIST can already be well represented using t-SNE [1] and thus any SSL      representation learning experiments performed on this dataset are dubious. I encourage the authors to          reproduce the visualization obtained on MNIST with CIFAR-10 or CIFAR-100.\n\nOne potential issue with adding noise to the output representation is that we limit the bandwith of the        representation. Therefore, while we may see improvement on small datasets such as CIFAR-100, this improvement  may vanish as we scale up the size of the dataset. I encourage the authors to consider running their           experiments on ImageNET to validate the their method scales up to larger and more diverse datasets. \n\n[Minor] There are some issues with the formatting. E.g. in page 1, the citation overflows to the second page   breaking the reference.\n\n[1] https://www.kaggle.com/code/venkatkrishnan/t-sne-decomposition-on-mnist-dataset", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JXpY7Akicg", "reviewer": "Reviewer_UXaV", "summary": "This paper proposes an improved Deep Infomax (DIM) method for self-supervised representation learning. Overall, an independent noise $Z$ is added to the normalized representations of the data. The authors demonstrate that the method is effective since it keeps the same training objective as the DIM. Experimental results also show the effectiveness.", "questions": "If the authors can clarify the significance of the additive noise in the theoretical view, this work will be complete, in my opinion. Then, I will consider changing my rating.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The research has been well accomplished since both the theoretical effectiveness and experimental results are provided.\n\n2. The background and fundamental method are well described, so readers can easily follow the idea.", "weaknesses": "The only improvement in this work is the additive noise in DIM, which may be not very overwhelming. However, if this small improvement can be thoroughly discussed and analyzed, the contribution of this work could still be significant. Therefore, my concerns mainly focus on the discussions and analysis:\n\n1. Since the authors claim in the title that the proposed method is an efficient matching method, it should be clarified why the improved approach is efficient. However, I only see that adding noise will not impact the DIM objective while improving the performance. The efficiency is not discussed.\n\n2. The advantage of adding noise is also not well discussed. It is stated that Lemma 3.1 highlights the importance of both additive noise and input data augmentation. Still, I only see the analysis of the significance of augmentation and do not understand why adding noise is also necessary.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "j15ipSPfgR", "reviewer": "Reviewer_RR5m", "summary": "The paper modifies the Deep InfoMax approach to learning representations that match specific distributions (uniform and normal). Specifically, it suggests adding noise to the normalized representations encoded from the inputs, before training them to maximize the mutual information with encoded input augmentations. It provides theoretical upper bounds for the mutual information between the noisy representation and the representation of the augmentation. It shows empirical results demonstrating the distribution matching, and the trade-off regarding the amounts of noise - a high noise level facilitates better distribution matching, but might result in worse representation learning.", "questions": "My main questions are about scaling: \n\nWill the representation learning improve over the baseline InfoMax in more conventional large-scale settings - ImageNet linear-probing/fine-tuning? \n\nHow stable is the training with noise under large-scale settings?\n\nFor generation and sampling - can you show interpretable individual directions in the learned representation distribution (again - on more advanced datasets like FFHQ/ImageNet)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper is very clear and well-written. It is easy to follow the derivations and the main ideas. \n\n2. It provides a good theoretical and empirical analysis of the method and the trade-offs between representation learning and distribution matching.  \n\n3. The presented modification is simple, and yet it provides a non-negligible improvement in usefulness (sampling).", "weaknesses": "1. The quality of conditional and unconditional sampling, the main property that can be achieved from a distribution matching is not tested on larger scale data. The paper can significantly improve with generation results on a larger scale of data (e.g. ImageNet/Cifar100).\n\n2. Even for the representation learning tests, this paper can improve with testing on larger datasets. The results of Figure 1, for example, are mostly saturated on MNIST and it will be useful to verify the same on larger more diverse datasets. \n\n3. While the paper has a section that argues that the same method can be applied to other SSRLs, it lacks empirical results that can demonstrate the usefulness/fragility of the training regime with noise.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)