Connecting Jensen–Shannon and Kullback–Leibler Divergences: A New Bound for Representation Learning

We thank the Reviewer for their positive and detailed feedback. We are glad to see that the key strengths and contributions of the paper, especially the theoretical connection between JSD and KLD and the tightness of the proposed bound, were clearly appreciated. Below, we address the concerns and suggestions raised.

Weaknesses:

On the assumption of an optimal discriminator. We agree that the tightest bound assumes access to an optimal discriminator and recognize that, in practice, discriminators are trained with limited capacity and finite data. In the revised manuscript, we have added a discussion (Section 4.3) on how the tightness of the bound may degrade under suboptimal discriminators. Specifically, we note that even when the discriminator is not optimal, its output remains a lower bound on JSD and thus still leads to a valid (though looser) lower bound on mutual information. Moreover, we show in the Appendix that we can estimate the gap between the true JSD lower bound and the one obtained from the cross-entropy $\mathcal{L}_{CE}$ as:

$ \Xi(\text{D}_{\text{JSD}}\left[p_{UV} || p_{U} \otimes p_{V}\right]) - \Xi(\log2 - \mathcal{L}_{CE}) \approx \delta \cdot\Xi'(\text{D}_{\text{JSD}})>0

$

where $\delta=\mathbb{E}\_{(u,v)\sim \tilde{p}\_{UV}} \text{D}\_{\text{KL}}[\tilde{p}\_{Z|UV}(u,v)\;||q\_{\theta}(z| u,v)]$ is the KL divergence between the true posterior of the mixture model defined in Eq. (21) and the parametric approximate posterior $q\_{\theta}$.

Real-world settings and representation learning outcomes To demonstrate the utility of our JSD-LB in practical representation learning tasks, we have added a new set of experiments in the revised Appendix D.3.2 using the Information Bottleneck (IB) framework. Specifically, we use the JSD information as a surrogate for MI. Following the experimental setting from Pan et al (2021) [1], we evaluate its impact on downstream tasks such as generalization, adversarial robustness, and out-of-distribution detection using the MNIST dataset.

Our results, reproduced below in Table 1, show that replacing the MI estimator with JSD-LB leads to state-of-the-art performance in terms of generalization, adversarial robustness, and out-of-distribution robustness. Specifically, it outperforms other Variational Information Bottleneck objectives [2,3,4] and the Disentangled Information Bottleneck approach [1]. This set of experiments shows that maximizing JSD has practical value for representation learning.

Computational cost and training efficiency We thank the Reviewer for pointing this out. In the revised Appendix D.2.5, we now report the time requirements between JSD-LB and other estimators such as MINE and NWJ. Table 2 from the response to Reviewer KvyC shows the results. As in Letizia et al. (2024), we found that the influence of the MI lower-bound objective functions on the algorithm’s time requirements is minor.

Questions

Can the authors provide a sensitivity analysis or theoretical discussion on how the bound evolves under suboptimal discriminators?

Please find above a discussion of the sensitivity and a theoritical discussion on how the bound evolves under suboptimal.

Would the bound remain tight and competitive when used in real-world settings, such as self-supervised representation learning based model pretraining.

Evaluating the tightness and competitiveness of our bound in real-world scenarios, such as self-supervised representation learning, would require access to the intractable true mutual information. While our main objective is to provide a theoretical justification for why maximizing JSD-based information effectively increases mutual information, such empirical evaluation is left for future work.

To validate the practical benefits of JSD-LB, I suggest the authors conduct controlled experiments where it replaces conventional MI estimators in IB-based methods (e.g., [1,2]).

We thank the Reviewer for this suggestion, which is indeed an excellent use case to demonstrate the practical value of our work. We conducted additional experiments where the JSD-LB was replaced with conventional MI estimators in IB-based methods, reported in Tables 1, 2, 3, and 4 below. The results show the effectiveness of our approach, outperforming other MI estimators and reaching comparable performance to disentanglement approaches. These results will be included in our revised manuscript.

Please consider adding wall-clock comparisons, memory usage analyses, etc, especially in complex higher-dimensional representation space.

See our response above and Table 2 from the response to Reviewer KvyC.

Variational Information Bottleneck Experiments

Information bottleneck results, where the objective is to compress the source random variable $U$ to a "bottleneck" random variable $T$, keeping the information relevant for predicting the target random variable $V$. This corresponds to the following maximization objective: $I [ T; V ] - I [T;U]$

In this set of experiments, we reproduce the experiments from [1] on MNIST. Specifically, the source random variable $U$ corresponds to an image and $V$ to its label. We employ the same model architecture to encode the variable $U$ and obtain the variable $T$, and directly report the corresponding results of [1].

In our setting (JSD-LB), $I [ T; V ]$ is maximized by maximizing the JSD via the cross-entropy loss.

• Information compression

  Table 1: Comparison of $I[T;V]$ obtained using non-disentanglement bottleneck approaches [2,3,4], with disentanglement [1] and via JSD optimization (ours).

  Dataset	Training	Testing
  H(V)	2.30	2.30
  VIB [2]	2.01	1.99
  NIB [3]	1.95	1.97
  squared-VIB [4]	1.94	1.95
  squared-NIB [4]	1.98	1.99
  DisenIB [1]	2.25	2.17
  JSD-LB	2.30	2.29

• Generalization

  Generalization performance in Table 2 is evaluated by the classification mean accuracy on MNIST test set after training the model on MNIST training set

  Table 2: Generalization performance (%) on MNIST dataset.

  Method	VIB [2]	NIB [3]	squared-VIB [4]	squared-NIB [4]	DisenIB [1]	JSD-LB
  Testing	97.6	97.2	96.2	93.3	98.2	98.8

• Robustness to adversarial attack

  To evaluate adversarial robustness, we adopt a standard one-step gradient-based attack method following [5]. After training the model on the MNIST training set, we assess robustness by computing the average classification accuracy on adversarially perturbed inputs derived from both the training and test sets. These adversarial examples are created by applying a single gradient ascent step on the input with respect to the loss. The perturbation strength is controlled by the parameter $\epsilon$, which we vary over the values {0.1, 0.2, 0.3}, corresponding to different pixel-wise perturbation magnitudes, as commonly done in prior work [5].

  Table 3: Adversarial robustness performance (%) on MNIST dataset.

  	Method	VIB [2]	NIB [3]	squared-VIB [4]	squared-NIB [4]	DisenIB [1]	JSD-LB
  a) Training set
  	ε = 0.1	74.1	75.2	42.1	61.3	94.3	99.5
  	ε = 0.2	19.1	21.8	8.7	24.1	81.5	96.0
  	ε = 0.3	3.5	3.2	5.9	9.3	68.4	91.4
  b) Testing set
  	ε = 0.1	73.4	75.2	42.7	62.0	90.2	94.6
  	ε = 0.2	20.8	23.6	9.2	24.5	80.0	89.6
  	ε = 0.3	4.2	3.4	5.9	9.9	67.8	86.1

• Out-of-distribution detection

  In line with prior work [1], we evaluate out-of-distribution (OOD) detection using a synthetic dataset composed of 10,000 samples drawn independently from a standard Gaussian distribution $\mathcal{N}(0, 1)$. The goal is to determine whether the model can reliably distinguish in-distribution data from these OOD noise samples. For evaluation, we adopt common OOD detection metrics including False Positive Rate at 95% True Positive Rate (FPR95), Detection Error, Area Under the ROC Curve (AUROC), and Area Under the Precision-Recall Curve (AUPR).

  Table 4: Distinguishing in- and out-of-distribution test data for MNIST image classification (%). ↑ (resp., ↓) indicates that a larger (resp., lower) value is better.

  Method	VIB [2]	NIB [3]	squared-VIB [4]	squared-NIB [4]	DisenIB [1]	JSD-LB
  TPR95 ↓	27.4	34.4	49.9	47.5	0.0	0.00
  AUROC ↑	94.6	94.2	86.6	85.6	99.4	100.0
  AUPR In ↑	94.8	95.2	83.5	83.3	99.6	99.9
  AUPR Out ↑	93.7	91.8	83.2	83.3	98.9	99.9
  Detection Err ↓	11.5	11.9	20.0	15.0	1.7	0.1

[1] Pan, Z.; Niu, L.; Zhang, J.; and Zhang, L. 2021. Disentangled Information Bottleneck. In AAAI

[2] Alemi, A. A.; Fischer, I.; Dillon, J. V.; and Murphy, K. 2017. Deep variational information bottleneck. In ICLR.

[3] Kolchinsky, A.; Tracey, B. D.; and Wolpert, D. H. 2019. Nonlinear information bottleneck. Entropy 21(12): 1181.

[4] Rodrıguez Galvez, B.; Thobaben, R.; and Skoglund, M. 2020. The Convex Information Bottleneck Lagrangian. Entropy 22(1): 98.

[5] Goodfellow, I. J.; Shlens, J.; and Szegedy, C. 2015. Explaining and Harnessing Adversarial Examples. In ICLR.

We thank the Reviewer for their thoughtful and detailed review. We are pleased that the theoretical rigor and practical relevance of our work were appreciated. In particular, we are glad that the Reviewer acknowledged that this work closes a gap between the practice and the theory. We address the raised concerns through additional experiments and clarifications, detailed below.

Weaknesses

Approximation of the Bound

We are happy to clarify how we derived our approximation of the $\Xi$ function. We began by observing that the equality $x = \Xi^{-1}(e^{-y})$, where $\Xi^{-1}$ is known analytically, leads to:

$$\mathbb{L}\left(\frac{1}{2} \cdot (\frac{x}{\log 2} + 1)\right)=\log \left( \frac{\frac{x}{\log 2} + 1}{1 - \frac{x}{\log 2}} \right)= \log \left( \frac{4 \log 2}{(1 + e^{-y}) \log(1 + e^{-y}) + y e^{-y}} - 1 \right) .$$

Then, we found that the first-order Taylor polynomial of the right term near $y=0$ is $y\mapsto y$. This led us to approximating $\Xi$ with a logit. We then conducted a grid search and found that scaling the logit function by a factor of 1.15 led to the best median approximation error over $x \in (0,1)$, i.e.:

$

\Xi(x) \approx 1.15 * \mathbb{L}\left(\frac{1}{2} \cdot (\frac{x}{\log 2} + 1)\right) .

$

Importantly, $\Xi$ is a strictly increasing function. As a result, maximizing $\Xi(\text{I}\_{\text{JS}}[U, V])$ is equivalent to directly maximizing $\text{I}_{\text{JS}}[U, V]$. For this reason, the approximation of $\Xi$ is not used during optimization; it is included solely for the reader’s interest and for interpretability.

In the revised manuscript:

• We explain how this approximation arises from the Taylor expansion of the KL divergence in terms of JSD.
• We include a plot comparing the approximation to the numerically evaluated $\Xi$, highlighting the error.
• We explicitly state that the optimization doesn't rely on the approximation of $\Xi$.

High Variance in the Student’s T Experiments (Fig. 4)

We agree that the Student’s T distribution poses challenges due to its heavy tails. While our method does exhibit higher variance in this specific setting, we note that this behavior is shared across estimators, and the challenge is characteristic of the distribution rather than the method alone.

Usefulness for Representation Learning

To better support the applicability of our method in real-world scenarios, we included a new representation learning benchmark based on the Information Bottleneck (IB) framework (Appendix D.3.2). IB seeks to learn compressed representations that retain information relevant to predicting a target variable.

Following prior work (e.g., Pan et al., 2018), we evaluate the resulting representations in terms of generalization, adversarial robustness, and out-of-distribution detection on MNIST.

The new results replicated in response to Reviewer ifQp show that:

• Replacing conventional MI estimators with JSD-LB within the IB framework leads to superior performance on various downstream tasks, including generalization, adversarial robustness, and out-of-distribution detection.
• This supports our claim that maximizing JSD is both a theoretically grounded and practically effective strategy for representation learning.

Clarifications and Corrections

• Eqs. (79) and (80): Thank you for pointing this out. We have corrected the reversed definitions in Eq. (79). Regarding Eq. (80), we clarified that $Z$ is a Bernoulli variable with parameter $\frac{1}{2}$, ensuring equal marginal probabilities $p(z=0) = p(z=1)$, which justifies the cancellation in the expression.

• Fig. 2: The purpose of this figure is to illustrate how tight our bound is for pairs of uniform categorical variables of varying dimensionality $k$, where both $(\text{I}\_{\text{JS}}, \text{I}\_{\text{KL}})$ can be computed analytically. Specifically, we show how this pair evolves as the joint distribution becomes more dependent. The color gradient encodes the number of categories $k$: higher $k$ values correspond to trajectories that reach higher values of both $\text{I}\_{\text{KL}}$ and $\text{I}\_{\text{JS}}$. Each colored curve represents the trajectory of $(\text{I}\_{\text{JS}}, \text{I}\_{\text{KL}})$ as the joint distribution $p\_{U,V}^{(\alpha)}$ transitions from the independent case ($\alpha = 0$) to the fully dependent case ($\alpha = 1$). The red curve corresponds to the theoretical lower bound we derive between $\text{I}\_{\text{JS}}$ and $\text{I}\_{\text{KL}}$, which is the tightest for arbitrary distributions. This figure demonstrates that the end of each trajectory lies exactly on the boundary. As suggested by Reviewer zy6E, we will include an analytical proof of this fact. Therefore, our bound is also tight in the mutual information setting, i.e., when comparing a joint distribution to the product of its marginals. We will revise the figure caption and the main text to make this interpretation more explicit.

• L43: We revised the language to avoid the phrase “KLD-based MI,” instead referring to “standard MI, defined via KL divergence.”

• L64: Corrected to “nonparametric,” without a hyphen.

• L167: Thank you for pointing this out. The correct full domain is $[0,1] \times (0,1) \cup \\{ (0,0),(1,1) \\}$. Therefore, the the lower triangle is $\Omega=\\{(\mu,\nu) \in [0,1]\times(0,1] \mid \nu \leq \mu\\}$. We corrected this in L165, L167, L169, L420, L444 and L446.

• L212: The definition of the distribution $\tilde{p}(z|u,v)=\frac{\tilde{p}(z)\tilde{p}(u,v|z)}{p(u,v)}$, where each term is defined in Eq. (21) and (22) has been added when first introduced.

Questions

It is unclear what the authors mean by "optimal bound".

In this work, we derive the tightest possible lower bound between the Jensen–Shannon divergence (JSD) and the Kullback–Leibler divergence (KLD) by analyzing the joint range between KLD and JSD. This means that our bound is saturated. i.e., equality is achieved for a family of distributions (right edge $S_3$ of our triangle). In contrast to classical inequalities such as Pinsker’s inequality, which provides a loose bound, our bound is attainable and therefore the tightest.

Although our bound is derived in full generality for arbitrary pairs of distributions $p$ and $q$, one might expect that restricting to structured pairs, such as joint and product-of-marginals (as in mutual information), could allow for a tighter bound. However, through empirical investigations on discrete distributions, where both JSD and KLD between joint and product-of-marginals can be computed exactly, we observed experimentally that such pairs already lie near the boundary of our derived lower bound (Figure 2). In our response to Reviewer zy6E, we proved that we can find an infinite family of joint and marginal discrete distributions on this bound. This suggests that our bound is also the tightest in the context of mutual information estimation.

We sincerely thank the Reviewer for the time and effort spent on reviewing our paper. We appreciate that the Reviewer recognizes the novelty of the connection between KL divergence, JSD, and cross-entropy, and its application to mutual information (MI) estimation. Below, we respond to the specific concerns and clarify the theoretical and empirical contributions of our work.

Weaknesses

“It is not clear whether the estimation of the MI by exploitation of such bounds can lead to superior performance with respect to state-of-the-art estimators.”

Our primary goal is to theoretically justify a practice that is already widespread in representation learning: maximizing JSD-based information. While JSD has been used empirically as a proxy for MI (e.g., in the popular Deep InfoMax framework), our contribution shows that maximizing JSD does maximize a tight lower bound on MI. This connection clarifies the mathematical rationale behind JSD-based information methods, which often work well in practice.

As a result, our initial experiments evaluated how this JSD-based bound performs compared to other plug-in MI estimators (MINE, NWJ, and CPC) compatible with representation learning.

Importantly, while alternative two-step MI estimators, such as $f$-DIME or SMILE, have been shown to demonstrate state-of-the-art bias-variance tradeoffs, they require retraining whenever the underlying joint distribution changes. Therefore, once trained, they cannot be reused in representation learning, since the distribution of the latent variables is evolving during optimization. For this reason, we initially excluded them from our experiments. However, we agree with the Reviewer that they offer useful context and therefore conducted additional experiments as detailed below.

Questions

"Possibly provide some analysis on how tight the bound is..."

Yes, the JSD–KLD lower bound we derive is the tightest possible between these two divergences. Specifically, it corresponds to the lower envelope of the joint range of $(\mathrm{KL}, \mathrm{JSD})$ pairs over Bernoulli distributions, and we provide an explicit analytic expression via its inverse mapping (Eq. 4).

While our result holds for arbitrary distributions, one might expect a tighter bound when restricted to the mutual information setting, i.e., when comparing a joint distribution to the product of its marginals. Remarkably, as shown in Figure 2 of our manuscript, there exist discrete joint-marginal pairs that lie very close to our general bound. To address the Reviewer’s comment, we performed an analytical derivation of the KL and JSD divergences between the joint and product of marginals for fully dependent uniform categorical variables of dimension $k$ (i.e., when $p_{UV}$ is diagonal and $\alpha = 1$). We found that these pairs lie on our bound for any $k\in\mathbb{N}^{+}$, providing an infinite family of distributions within the mutual information setting on our bound. This result will be included in the Appendix and referenced in the main text. We thank the Reviewer for this suggestion, which has helped us strengthen the contribution of our work.

"Is the function E in Theorem 4.1 unique?"

We thank the Reviewer for this question. Yes, the function $\Xi$ is unique and defined via its inverse $\Xi^{-1}$, which maps JSD to a lower bound on KL. We added a proof that $\Xi^{-1}:\mathbb{R}^+\mapsto[0,\log2)$ is bijective in the revised manuscript. We clarify this point in the revised manuscript (Section 4.1).

"Better describe Section 4.3 and its relation to GAN-DIME." [...] "GAN-DIME is said to be the best estimator, but not shown in plots."

Thank you for pointing this out. Section 4.3 indeed proposes a two-step estimator, which turns out to be equivalent to GAN-DIME, up to a linear transformation. We now clarify this equivalence explicitly in the revised manuscript.

"The comparison of MI estimation does not consider all estimators (e.g., SMILE)."

We have now added comparisons to SMILE and KL-DIME. While these cannot be directly optimized during the training of representation methods (as they require a two-step optimization procedure), we agree that they offer useful context. These new results have been added in Appendix D.2.4 and are reproduced below in Table 1. Our results show that the two-step estimator derived from our bound performs on par with or better than KL-DIME and SMILE, especially in high-MI or high-dimensional settings.

"No analysis of variance is reported."

We have added a detailed variance and MSE analysis in the revised version (Appendix D.2.3), which is reported in rebuttal to Reviewer KvyC. The key findings are:

• Our lower bound (JSD-LB) shows significantly lower variance than MINE and NWJ, especially at higher MI values and larger dimensions.
• The MSE of our lower bound is lower than that of CPC and NWJ.

"What is the advantage of this method over others?"

The main advantage is practical usability: many state-of-the-art MI estimators (e.g., GAN-DIME, SMILE) require a two-step training pipeline that is not compatible with end-to-end training in representation learning. Our bound can be directly optimized via JSD, using a single discriminator, enabling stable, end-to-end training. This provides a principled rationale for existing JSD-based methods and introduces a theoretically grounded alternative to common variational estimators like MINE.

"No results are reported about the bound’s value in representation learning."

To strengthen the bound's value in representation learning, we have added new experiments in the Information Bottleneck (IB) framework. Results in revised Appendix D.3.2 and reproduced in our response to Reviewer ifQp show that maximizing JSD as a proxy for MI in IB leads to superior performance compared to standard MI estimators.

"Minor: Equation (6): please check if it should be dp/dq."

It should indeed be $dp/dq$. For example, for the KL-divergence: $f:t\mapsto t \ln(t)$.

Conclusion

We hope the clarification, additional proofs, and new experiments address the Reviewer's concerns.

Comparison with two-step MI estimators

Table 1: Bias, variance, and mean squared error (MSE) of the two-step MI estimators using the joint architecture, when $b \in$ {64,258}, for the Gaussian setting with varying dimensionality $d \in$ {5,20}.

Tightness analysis

Let $(U,V)$ be a pair of uniform categorical distributions of dimension $k$ that are fully dependent, i.e. $P_{U,V} = \text{diag}(P_U)$, where $P_U=\frac{1}{k}\mathbf{1}$. Then $\text{I}\_{\text{KL}}\left[U,V\right] = \Xi\left( \text{I}\_{\text{JS}}\left[U,V\right] \right)$.

Proof:

Let $k\in\mathbb{N}^{+}$. We assume that $U$ and $V$ are fully dependent (i.e. $U=V$) uniform categorical distributions of dimension $k$. Therefore, the matrix representation $P\_{U}$, $P\_{V}$ and $P\_{UV}$ are respectively $\frac{1}{k}\mathbf{1}\_{k}$, $\frac{1}{k}\mathbf{1}\_{k}$ and $\frac{1}{k} I\_{k}$ . Consequently, $P\_{U}\otimes P\_{V} = \frac{1}{k^2}\mathbf{1}\_{(k,k)}$.

Let $M = \frac{1}{2} ( P\_{U} \otimes P\_{V} + P\_{U,V})$ be the mixture. $M_{i,j}=\frac{1/k+1/k^2}{2}$ if $i=j$ and $M_{i,j}=\frac{1}{2k^2}$ otherwise.

Let's calculate $\text{I}\_{\text{KL}}\left[U,V\right]$ and $\text{I}\_{\text{JS}}\left[U,V\right]$.

$\text{I}\_{\text{KL}}\left[U,V\right] = \sum_{i=1}^{k} \frac{1}{k} \log \frac{1/k}{1/k^2} = \log(k)$

$\mathrm{I}\_{\text{JS}}[U, V]= \frac{1}{2} \left( \sum\_{i=1}^{k} \frac{1}{k} \log \frac{1/k}{\frac{1}{2}(1/k + 1/k^2)} + \sum\_{i=1}^{k} \frac{1}{k^2} \log \frac{1/k^2}{\frac{1}{2}(1/k + 1/k^2)} + \sum\_{i=1}^{k} \sum\_{j=1\text{ s.t. } j \neq i}^{k} \frac{1}{k^2} \log \frac{1/k^2}{\frac{1}{2k^2}} \right)$ $\quad \quad \quad \quad = \frac{1}{2} \left(\log \frac{2}{1 + 1/k} + \frac{1}{k} \log \frac{2/k}{1 + 1/k} + \frac{k-1}{k} \log(2) \right)$ $\quad \quad \quad \quad = \log(2) - \frac{1}{2} \left( \log(1 + 1/k) - \frac{1}{k} \log \left( \frac{1/k}{1 + 1/k} \right) \right)$ $\quad \quad \quad \quad = \Xi^{-1}(\log(k))$

Therefore: $\text{I}\_{\text{JS}}\left[U,V\right] = \Xi^{-1}(\text{I}\_{\text{KL}}[U, V]).$

Dear Reviewer,

Thank you for your excellent follow-up questions. Please find our responses below:

a) Our theoretical contribution is to: 1) show that maximizing the JSD-based information \text{I}\_{\text{JS}}\ is a principled surrogate for maximizing the mutual information ($\text{I}_{\text{KL}}$); 2) derive and prove the tightest possible bound between JSD and KL for arbitrary distributions (Theorem 4.1.).

We agree with the reviewer that MI is unbounded, unlike $\text{I}_{\text{JS}}$. Therefore, our JSD lower-bound is not always tight across the full range of pairs ($\text{I}_{\text{JS}},\text{I}_{\text{KL}}$). However, our experiments on discrete distributions and continuous distributions show that our bound $\Xi(\text{I}_{\text{JS}})$ tracks MI reasonably well in practice.

b) SMILE and $f$-DIME were designed for MI estimation, not for MI maximization in representation learning. These methods rely on a two-step procedure: first training a network $T_\theta$ using a variational objective (based on $f$-divergences), and then estimating MI from the learned $T_\theta$ via a separate expression.

To illustrate this, in $f$-DIME (each two-step estimator is detailed in Appendix A.2), $T_\theta$ is trained by maximizing:

$$\mathcal{J}\_f(T) = \mathbb{E}\_{(u,v) \sim p\_{UV}} \left[ T(u,v) - f^{*} \left( T\left(x,\sigma(y)\right) \right) \right] \quad (1)$$

where $f^{*}$ is the Fenchel conjugate of $f$. Once trained, MI is estimated via:

$$I\_{f\text{-DIME}} = \mathbb{E}\_{(u,v) \sim p\_{UV}} \left[ \log \left( (f^*)'(\hat{T}(u,v)) \right) \right] \quad (2)$$

In representation learning (RL), the distribution over $(U, V)$ changes as the encoder evolves. Using (2) as a training objective would require retraining $T_\theta$ at every step, a computationally expensive strategy.

The Reviewer interestingly suggests using the optimization objective (1) directly as the training RL loss. In the case of GAN-DIME and SMILE, this is effectively equivalent to maximizing JSD-based information, which is indeed commonly used. However, (1) was not originally designed as a surrogate for MI; it was constructed to obtain a MI estimator in (2). Therefore, there were no theoretical guarantees that optimizing (1) would increase MI or improve the learned representations.

Our work closes this gap by showing that when (1) corresponds to JSD (as in Deep InfoMax / SMILE or GAN-DIME), optimizing it does maximize a lower bound on MI, thereby providing a principled justification for these practices.

c) As noted above, since estimators such as SMILE and $f$-DIME require training a separate network ($ T_\theta \ $) for each estimation step, they are less suitable for representation learning, where the data distribution evolves.

In contrast, methods like MINE, NWJ, CPC, and our JSD-LB are variational lower bounds that do not require an optimization for estimation. Therefore, they can be used as direct training objectives. In our case (which matches the Deep InfoMax objective), it simply involves training an encoder $ f_\phi $ and a discriminator $ D_\theta $ jointly to discriminate between positive pairs from the joint $(U, V=f_\phi(U))$ and negative pairs sampled from the product of marginals.

d) The results of our two-step estimator and GAN-DIME are indeed equivalent, since they share an equivalent discriminative optimization objective (up to a linear transformation) and a similar procedure for two-step estimations. We omitted $f$-DIME and SMILE from the main paper, as our focus was to compare direct MI estimators and to justify that our JSD lower bound is both practically tight and stable. We will add these additional results in the Appendix.

Let us emphasize once again that our main contribution is not a new estimator, but a theoretical justification for an existing and popular representation learning objective (JSD/Deep InfoMax). We show that it maximizes a lower bound on MI, a theoretical guarantee previously missing in the literature that helps explain the empirical success of JSD-based methods in representation learning.

We sincerely appreciate your constructive feedback and hope these clarifications address your remaining concerns. We are happy to continue the discussion if needed.

We thank the Reviewer for their time. We are pleased to clarify the practical implications of our theoretical contributions.

Weaknesses:

Practical consequences of the tight lower bound on machine learning applications.

Our primary objective is to provide a theoretical foundation for a practice that is already widespread in representation learning: maximizing JSD-based information using discriminators. For instance, this optimization strategy is a cornerstone of frameworks like Deep InfoMax [1], which has been successfully applied to a wide range of downstream tasks, such as image classification [1], graph node classification [2], molecular property prediction [3], recommender systems [4] and remote sensing [5]. While JSD has often been used empirically as a proxy for mutual information, its theoretical justification remained unclear.

Our work closes this gap: we demonstrate that maximizing JSD corresponds to maximizing a tight lower bound on the true mutual information. This result strengthens the theoretical underpinnings of many existing methods and provides assurance that optimizing JSD discriminative losses is a principled and effective strategy for maximizing mutual information.

Questions

Does it lead to better representations? Can you show a real-world application that benefits from it in terms of accuracy?

We thank the reviewer for this important question. To further demonstrate the practical relevance of our results, we conducted additional experiments using the Information Bottleneck (IB) framework. The IB framework aims to learn compact yet informative representations from an input variable that are predictive of a target variable. We included new experiments where the objective is to maximize JSD rather than alternative variational bounds. Following prior work [6], we assess the learned representations obtained via JSD maximization on generalization performance, adversarial robustness, and out-of-distribution detection using the MNIST dataset.

Our findings show that:

• Maximizing JSD information outperforms alternative variational IB baselines that rely on alternative mutual information bounds.

• These empirical results support our theoretical claims and confirm that maximizing JSD not only has solid theoretical grounding but also leads to practically robust representations.

What are the real-world implications of this work?

The real-world implication of our work is that methods aiming to maximize mutual information, particularly in representation learning, can confidently rely on JSD-based objectives as a mathematically justified and effective surrogate.

[1] Hjelm, R. D., Fedorov, A., Lavoie-Marchildon, S., Grewal, K., Bachman, P., Trischler, A., & Bengio, Y. (2018). Learning deep representations by mutual information estimation and maximization. In: ICLR.

[2] Veličković, P., Fedus, W., Hamilton, W. L., Liò, P., Bengio, Y., & Hjelm, R. D. (2019). Deep Graph Infomax. In: ICLR

[3] Stärk, H., Beaini, D., Corso, G., Tossou, P., Dallago, C., Günnemann, S., & Liò, P. (2022). 3D Infomax improves GNNs for Molecular Property Prediction. In: ICML.

[4] Yu, J., Yin, H., Xia, X., Chen, T., Li, J., & Huang, Z. (2023). Self-supervised learning for recommender systems: A survey. IEEE Transactions on Knowledge and Data Engineering, 36(1), 335-355

[5] Wang, Y., Albrecht, C. M., Braham, N. A. A., Mou, L., & Zhu, X. X. (2022). Self-supervised learning in remote sensing: A review. IEEE Geoscience and Remote Sensing Magazine, 10(4), 213-247.

[6] Pan, Z.; Niu, L.; Zhang, J.; and Zhang, L. 2021. Disentangled Information Bottleneck. In AAAI

We thank the Reviewer for the positive and constructive feedback. We appreciate that the Reviewer clearly identified the key contributions of our work. As suggested, we have conducted additional experiments and revised the manuscript to include:

• Bias-variance and time complexity analyses
• Extended experiments to higher dimensions
• Comparisons to additional baselines, including two-step estimators such as SMILE and $f$-DIME

Below, we address each of the Reviewer’s points in detail.

Weaknesses

On bias-variance results and time complexity. We appreciate the suggestion. In the revised manuscript, we include a detailed analysis of the bias, variance, and mean squared error (MSE) of our JSD lower bound, following protocols from Letizia et al. (2024) and Poole et al. (2019). These new results in the revised appendix D.2.3 (and reproduced in Tables 1 and 2) show that our JSD-LB consistently provides a stable lower bound on the true MI across dimensions:

• While MINE and NWJ exhibit high variance, JSD-LB and CPC remain more stable. In particular, MINE’s instability can lead to significant overestimation of its lower bound, sometimes even exceeding the true MI. While CPC provides low-variance estimates, they are upper bounded by the contrastive bound $\log(b)$, resulting in higher MSE and bias in high MI regimes. We experimentally found that the bias of JSD-LB increases with dimension $d$, but its variance remains stable, making it a reliable surrogate for MI maximization.
• When repurposed for two-step MI estimation, our method surpasses other variational lower-bound estimators across dimensions.
• As in Letizia et al. (2024), we found that the influence of the MI lower-bound objective functions on the algorithm’s time requirements is minor.

High-dimensional experiments. In response to the Reviewer’s concern about dimensionality, we extended our staircase experiments to higher dimensions (up to $d=100$). Results in Table 1 show that our method maintains a favorable bias-variance trade-off, while estimators such as MINE and NWJ have variance and stability that degrade with dimensionality.

More challenging benchmarks. To further demonstrate practical relevance, we added a new representation learning benchmark using the Information Bottleneck (IB) framework. Results presented in revised Appendix D.3.2 and reproduced in our response to Reviewer ifQp show that using JSD as a surrogate to MI leads to state-of-the-art downstream performance, outperforming other variational IB approaches. These experiments demonstrate that, in practice, maximizing JSD is an effective strategy for representation learning.

Comparison to Additional Baselines. While our initial focus was on variational bounds that can be directly used for MI maximization, we agree that including comparisons with two-step estimators provides better context. In revised Appendix D.2.4 (referenced in our response to Reviewer zy6E), we now report comparisons to SMILE, KL-DIME, and GAN-DIME. The results show that our derived two-step MI estimator, which is equivalent to the GAN-DIME (up to a linear operation), allows reaching competitive performance.

Other Comments.

• On the Poole reference: Thank you for pointing this out. We have corrected the citation to include all co-authors.
• On limitations: We added a short comment on the numerical verification of Conjecture B.3 in the main text (Section 4.1).

Questions

“Do you have any characterizations, either theoretical or experimental, of the bias-variance trade-off for the presented bound/MI estimator compared to the baselines? - or the time-complexity?”

As mentioned above, we provide an experimental characterization of the bias-variance trade-off of our bound. We found that maximizing JSD leads to an increase in the lower bound of MI that is stable. A more competitive bias-variance tradeoff can be obtained by reusing the discriminator for two-step MI estimation.

“How does the approach compare to baselines in higher-dimensional versions of the problem you consider - or, ideally, on a more challenging problem?”

See our response above for higher-dimensional cases and a challenging Information Bottleneck problem.

“Do you have a sense of how the presented bound (and the estimator) compares to, e.g., f-DIME or SMILE (even if the comparison should be done with due discussion of the limitations you point out)?”

See our response above.

Rating

We appreciate the Reviewer recognizing the theoretical contributions. We hope that our answers and changes address the Reviewer’s concerns and can contribute to increasing their score.

Additional experiments

Table 1: Bias, variance, and mean squared error (MSE) of the MI estimators using the joint architecture, when N = 64, for the Gaussian setting with varying dimensionality $d \in$ {5,10,20,100}.

The computational time analysis is developed on a server with CPU ”Intel Xeon Platinum 8468 48-Core Processor” and an NVIDIA GPU H100.

Table 2: Comparison of the time requirements (in minutes) to complete the 5-step staircase MI ($d=5$) over the batch size $b$ using