Efficient Distribution Matching of Representations via Noise-Injected Deep InfoMax

Dear Reviewer RR5m,

Thank you for your valuable feedback and insightful comments. We have carefully addressed your concerns and provided updates based on new experiments and analysis.

Conditional and Unconditional Sampling

We are currently finalizing ImageNet-1K pre-training and conducting experiments on conditional and unconditional sampling for CIFAR-10. While we are eager to evaluate the proposed approach in practical generative settings, we emphasize that this experiment is auxiliary to the main contribution of the paper. In our submission, generation is considered as a supplementary validation effort, but conducting a comprehensive analysis requires considerable computational resources. We will provide an update with results as soon as these experiments are complete.

Representation Learning

We extended our representation learning experiments to CIFAR-10, analogous to the MNIST experiment in the original submission.

Experimental setup. We use a ResNet-18 architecture augmented with a linear layer (mapping 512-dimensional output to $\mathbb{R}^2$) and a batch normalization layer without learnable parameters (affine=False). Both contrastive (SimCLR) and non-contrastive (VICReg) objectives are tested, showing consistent results. 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.

The outputs are then processed by the projection head and loss function (e.g., InfoNCE). Noise $Z$ is sampled from a normal distribution with standard deviation $\sigma$.

Models were trained for 400 epochs across $\sigma \in {0.0, 0.01, 0.025, 0.05, 0.1, 0.2, 0.2658, 0.5}$ with five random seeds and base learning rate base_lr=0.1. To accelerate convergence, we also trained models with $\sigma \in {0, 0.05, 0.1}$ for 800 epochs with base_lr=0.5. Downstream performance was evaluated using ROC AUC, with normality tests on the learned 2D representations. Results are shown in the updated manuscript's Appendix E and F (Figures 5–7). Figures 5 and 6 demonstrate that CIFAR-10 embeddings into 2D space increasingly resemble a Gaussian distribution as $\sigma$ increases, consistent with MNIST results. The overall results demonstrate the approach generalizes to more complex datasets than MNIST.

Q1: Performance in Large-Scale Settings

We acknowledge the original submission lacked experiments on large-scale datasets due to computational constraints. Motivated by your feedback and that of others, we conducted experiments on ImageNet-100 (a 100-class subset of ImageNet-1K) and are awaiting completion of full ImageNet-1K training.

We used VICReg with $\sigma \in {0.0, 0.05, 0.1, 0.2}$ and a ResNet-18 encoder (output dimension = 512). Models were trained for 100 epochs with standard hyperparameters across 5 random seeds. Performance is reported below (also in Appendix F, Table 3):

Noise injection minimally affects performance for small $\sigma$, suggesting the method scales well to larger datasets. Full ImageNet-1K results will be shared once training is complete.

Q2: Stability in Large-Scale Settings

Training on ImageNet-100 revealed sensitivity to higher noise magnitudes ($\sigma = 0.2$), with top-1 accuracy showing slight degradation and increased variance. However, training remains robust with small noise injection, and no collapse was observed across all runs.

Transferability to Other SSRL Methods

This is a fair concern as there are a great number of SSRL methods, e.g. contrastive, non-contrastive, self-distillation, clusterisation-based approaches to name a few. In our work, we focused on representative methods from contrastive (SimCLR) and non-contrastive (VICReg) families. SimCLR’s InfoNCE loss closely aligns with the InfoMax principle underlying our theoretical setup. VICReg also has information-theoretic interpretations, making it a suitable candidate for empirical validation. As for the broader applicability, given the connections between joint-embedding methods in existing literature, we expect our approach to generalize to similar SSRL methods. Our current experiments serve as representative case studies.

Thank you once again for your constructive feedback! We look forward to your further comments and suggestions.

Dear Reviewer UXaV, thank you sincerely for reviewing our article! In the following text, we provide responses to your concerns. We hope that all of them are addressed properly.

Weaknesses:

1. Acquiring embeddings conforming to a specific distribution (i.e., distribution matching, DM) is an auxiliary, yet important task in representation learning due to several applications which we list on lines 50-55 (generative modelling, statistical analysis, disentanglement, outliers detection). Conventional approaches to DM are either imprecise (lines 56-60), or almost as expensive as full-fledged generative modelling, requiring a setup similar to GANs, normalizing flows or diffusion models (lines 60-68). In contrast, we propose achieving DM via normalization and noise injection only, which is cheap (as no additional NN is required), and justified both theoretically (Section 3) and experimentally (Section 5). Thus, we claim our method to be efficient.

2. We highlight the importance of noise injection in Theorems 3.4-3.5 and Remark 3.6. Note that without the noise injection, the bounds from the aforementioned theorems deteriorate, becoming $+\infty$, which renders the maximum entropy properties inapplicable. Thus, one can not be assured that maximization of DIM objective yields normally or uniformly distributed embeddings, even if the normalization is still performed. Moreover, note that without the noise injection, the $I(f(X) + Z; f(X) \mid f(X'))$ term in Lemma 3.1 is either $0$ (when $f$ is weakly invariant to $X \to X'$), or $+\infty$ (otherwise), rendering this decomposition useless for the later analysis.

   Additional intuition on the importance of noise injection can be acquired from line 818 in the proof of Lemma 3.1:

   $I(f(X');f(X) + Z) = h(f(X) + Z) - h(f(X) + Z \mid f(X'))$

   Note that if $Z = 0$, and some kind of normalization is still employed (so $h(f(X))$ is upper-bounded), $h(f(X) \mid f(X'))$ can get arbitrary small, up to $-\infty$ for weakly invariant $f$. This makes the minimization of this term preferable to the maximization of $h(f(X))$ during the gradient descent (it is even possible that $I(f(X');f(X)) \to +\infty$, but $h(f(X))$ stays far from the maximal value). However, if $Z$ is non-zero, $h(f(X) + Z \mid f(X'))$ saturates, with the limiting value being lower for higher noise-to-signal ratio. Thus, the magnitude of the noise serves as a trade-off between the distribution matching and the weak invariance. This explanation is also backed by the experimental results from Figure 1, which illustrate the aforementioned trade-off.

Questions: in the response above, we stress out the importance of the noise injection from the theoretical point of view. In short, the desired distribution matching through DIM and embedding normalization only is not achievable; injecting the noise is crucial.

This clarification will be added to the upcoming revision of our manuscript. We hope that you will find our answer satisfactory. If there are some parts of our response which require further clarification, please, let us know.

We again sincerely thank Reviewer UXaV for carefully reading our article and pointing out the parts of our theoretical framework which require additional explanation.

Dear Reviewer y7PL,

Thank you for your thoughtful review and valuable feedback. We are pleased to address your concerns below.

Weaknesses:

2D Experiment on CIFAR-10

We appreciate your observation regarding the triviality of embedding MNIST into 2D space. Our primary goal with the MNIST experiment was to validate that the representations align with the theoretical predictions for noise injection. To address your suggestion and validate our approach on a more complex dataset, we have conducted additional experiments on CIFAR-10.

In these experiments, we used a ResNet-18 architecture augmented with an additional linear layer to map the 512-dimensional output to $\mathbb{R}^2$, along with a batch normalization layer without learnable parameters (i.e., affine=False). 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:

1. Unaugmented Input: The unaugmented image $X$ is fed to the encoder, yielding $f(X)$, followed by noise injection to obtain $f(X) + Z$.
2. 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 (e.g., InfoNCE). The noise $Z$ is sampled from a normal distribution with standard deviation $\sigma$. We conducted pre-training across various noise magnitudes with $\sigma \in \{0.0, 0.01, 0.025, 0.05, 0.1, 0.2, 0.2658, 0.5\}$, and for each $\sigma$, we ran training with 5 different random seeds with base_lr=0.1 for 400 epochs (visualized in the updated Appendix E, Figure 5) To speed up convergence we additionally trained 3 models ($\sigma \in \{0, 0.05, 0.1\}$) with base_lr=0.5 for 800 epochs (visualized in Appendix E, Figure 6).

After pre-training, we evaluated downstream performance using ROC AUC and performed normality tests on the learned 2D representations (see Appendix F, Figure 7), analogous to our previous MNIST experiments (Figure 1). Consistent with the MNIST results, the CIFAR-10 samples were successfully embedded into 2D space, exhibiting a distribution that increasingly resembles a Gaussian as $\sigma$ increases (Figures 5 and 6). This demonstrates that our approach generalizes to more complex datasets than MNIST.

Large-Scale Datasets

We acknowledge that the original submission lacked experiments on large-scale datasets due to computational constraints. Motivated by your feedback and that of other reviewers, we have secured additional resources to conduct experiments on larger datasets. While we await the completion of full ImageNet-1K training, we present preliminary results on ImageNet-100, a 100-class subset of ImageNet-1K.

Using the VICReg objective with noise injection, we experimented with $\sigma \in \{0.0, 0.05, 0.1, 0.2\}$ (note that $\sigma=0$ corresponds to the original VICReg setup), employing a ResNet-18 encoder (output dim = 512) and standard hyperparameters. We trained the models for 100 epochs and evaluated the downstream performance of the learned representations (encoder outputs). The results are summarized in the table below (Appendix F, Table 3):

These results indicate that noise injection has a negligible impact on performance for small noise magnitudes, suggesting that our method scales to larger and more diverse datasets without significant degradation in accuracy. Once training finishes on full ImageNet-1K, we will update the table.

Questions:

1. Is the linear probe trained on the output of the encoder or on some intermediate layer?

   The linear probe is trained on the output of the encoder. This approach is standard for classification benchmarks in self-supervised learning methods within the vision domain, particularly when using ResNet-type encoders. Most methods, such as SimCLR, VICReg, Barlow Twins, BYOL, and SimSiam, follow this pattern.

2. Are the gains observed in Table 1 transferable when training on ImageNet?

   While we are awaiting the completion of full ImageNet-1K pre-training, we can provide preliminary insights into the transferability of our method based on the ImageNet-100 experiments mentioned above. The results suggest that our approach maintains comparable performance on larger datasets. Please refer to the table above (also included in the updated manuscript) for detailed results.

We hope that these additional experiments and clarifications address your concerns. Thank you again for your constructive feedback, which has helped us strengthen our work.

Dear Reviewer Vzb6,

Thank you for your comprehensive review and for taking the time to carefully evaluate our submission. Below, we provide detailed responses to your comments and questions, and we hope that all your concerns are properly addressed.

Weaknesses:

1. We will add a concise clarification to this section, similar to the one provided for Gaussian noise, to enhance clarity. Below we provide the text to be included in the upcoming revision of our article.

   We consider $Z$ distributed uniformly on $[-\varepsilon; \varepsilon]^d$, and assume that $f(X)$ has bounded support on $[0; 1]^d$. By Theorem 2.2, the entropy $h(f(X) + Z)$ reaches its maximum precisely when $f(X) + Z$ is uniformly distributed. While the independence of $Z$ and $f(X)$ does not imply that $f(X)$ is uniform, the distribution of $f(X)$ that maximizes $h(f(X) + Z)$ is a discrete uniform distribution over a finite grid within $[0; 1]^d$. As $\varepsilon$ approaches zero, the distribution of $f(X)$ converges to a continuous uniform distribution over $[0; 1]^d$. Thus, injecting uniform noise asymptotically drives $f(X)$ towards a uniform representation over $[0; 1]^d$. This reasoning is formalized in Theorem 3.5

2. Thank you for your suggestion. We are already in the process of running our model on larger datasets, specifically ImageNet, and hope to include these results before the end of the rebuttal period.

Questions:
Thank you for pointing out the term "nats" in Figure 1. The quantity $C = \frac{d}{2} \log \left( 1 + 1/\sigma^2 \right)$, derived in Theorem 3.4, represents the capacity mentioned in the caption of Figure 1. It is measured in 'nats', units of information based on natural logarithms. We will update the Figure 1 caption to explicitly define "nats" and also include the mathematical expression for the capacity in the axis label.

We sincerely thank you once again for your constructive feedback, which has been valuable in refining our work. If any further concerns arise, we will be glad to address them.

Dear Reviewer Vzb6,

Once again, thank you very much for your detailed review. As we are nearing the end of the discussion period, we would like to ask if the concerns you raised have been addressed.

To remedy the first weakness, a concise clarification was added; please, refer to our response or to the beginning of Section 3.2 in our revised manuscript.

To address the second weakness, we've updated the manuscript to include experiments on larger scale datasets (including 2d embedding and representation learning with noise injection on CIFAR and ImageNet datasets; please see Appendix E and F in the updated manuscript). For your convenience, the resulting accuracies are also reported below; the experiments were conducted across 5 random seeds for ImageNet-100 and on 1 seed on ImageNet-1k due to compute limitations. We hope the other runs to finish before the discussion period ends.

Overall, the additional experiments are in line with the results obtained in the original empirical study.

We hope this update addresses your concerns, we thank you again for your review and look forward to your further comments and suggestions. We understand that the discussion period is short, and we sincerely appreciate your time and help!

Thank authors for the explaination and additional experiments on ImageNet. My concerns are addressed and I would like to raise my score.

Here is a summary of my suggestions:

1. give an (intuitive) explaination about how uniform representation is learned (gaussian representation is well explained).

2. move ImageNet experiment into main text, because it is more convincing.

Dear Reviewer Vzb6,

Thank you once again for your review and valuable suggestions. We appreciate your positive feedback. Below, we provide responses to your specific suggestions:

1. We have revised our explanation of how uniform representation is learned as follows:

   Now consider $Z$ distributed uniformly on $[-\varepsilon; \varepsilon]^d$. We assume that $f(X)$ has bounded support within $[0, 1]^d$. This restriction serves as an alternative to 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(f(X) + Z)$ is maximized when $f(X) + Z$ follows a uniform distribution. This can be achieved when $f(X)$ conforms to a discrete uniform distribution over an equidistant grid within $[0, 1]^d$. As $\varepsilon$ approaches zero, the distribution of $f(X)$ gradually converges to a continuous uniform distribution over $[0, 1]^d$. Thus, injecting uniform noise asymptotically drives $f(X)$ toward a uniform representation over $[0, 1]^d$. This reasoning is formalized in Theorem 3.5.

   We will replace the paragraph in the next revision.

2. The results from linear probing for ImageNet will be included in Table 2 of the main body, with additional details provided in Appendix F.

We are delighted to hear that your concerns have been addressed. We sincerely appreciate the time and effort you have put into reviewing our paper.