SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\sup  _ {P _ {X, Y}} \mathbb {E} [ | \mathcal {S I} ^ {*} (X; Y) - \widehat {\mathcal {S I} ^ {*}} _ {n, m} (X; Y) | ] \leq \delta (n) + \frac {U}{2 \sqrt {m}},
$$
where  $U\propto (d_x^{-1} + d_y^{-1})^{1 / 2}$  is a constant factor.
Proof. See Appendix A.4 for details.
We note that the estimation rate depends on the dimensions of the problem  $d_x, d_y$  only up to the constant factor  $U$ . The explicit dependence of  $U$  on  $d_x, d_y$  has been a recent challenge which is now solved in Theorem 4. By imposing additional regularity on  $P_{X,Y}$ , we can obtain  $\delta(n) \leq C\left((n\log n)^{-\frac{s}{s+2}}\log n^{1-\frac{2}{p(s+2)}} + n^{-1/2}\right)$  where  $s, p:0 < s \leq 2 \leq p$ , and  $C > 0$  (see Han et al. (2020); Goldfeld and Greenewald (2021)), which is significantly faster than the MI estimation rate, i.e.  $n^{-1/(d_x+d_y)}$ .
4 Neural Estimator for  $\mathcal{ST}^*$
In this section, we introduce a practical implementation of  $\mathcal{ST}^*$  (based on Corollary 2) by leveraging a neural network estimator and connect our theoretical results in Section 3 with modern machine learning frameworks. Let  $\mathcal{D} \stackrel{\mathrm{def}}{=} \{(X_k, Y_k)\}_{k=1}^m$  be a mini-batch with  $X \in \mathbb{R}^{d_x}$ ,  $Y \in \mathbb{R}^{d_y}$ , and  $n = \lfloor \frac{m}{2} \rfloor$ . We take  $(X, Y) \sim \{(X_k, Y_k)\}_{k=1}^n$ ,  $(\overline{X}, \overline{Y}) \sim \{(X_k, Y_{k+n})\}_{k=1}^n$ , and  $(\Psi_k, \Upsilon_k) \sim \pmb{\gamma}(\Omega_{d_x,d_y})$ . Then, the neural network estimator  $\mathcal{ST}_{\mathbb{W}}^*$  is computed as:
$$
\begin{array}{l} \sup  _ {T, f _ {1}, f _ {2}} \left\{\frac {1}{n} \sum_ {k = 1} ^ {n} T \left(\Psi_ {k}, \Upsilon_ {k}, f _ {1} ^ {(k)} ^ {\top} X _ {k}, f _ {2} ^ {(k)} ^ {\top} Y _ {k}\right) - \log \frac {1}{n} \sum_ {k = 1} ^ {n} \exp \left(T \left(\Psi_ {k}, \Upsilon_ {k}, f _ {1} ^ {(k)} ^ {\top} \overline {{X}} _ {k}, f _ {2} ^ {(k)} ^ {\top} \overline {{Y}} _ {k}\right)\right) \right. \\ + \lambda_ {1} \left(\frac {1}{n ^ {2}} \sum_ {k, j} \arccos  | f _ {1} ^ {(k)} ^ {\top} f _ {1} ^ {(j)} | - \omega_ {X}\right) + \lambda_ {2} \left(\frac {1}{n ^ {2}} \sum_ {k, j} \arccos  | f _ {2} ^ {(k)} ^ {\top} f _ {2} ^ {(j)} | - \omega_ {Y}\right) \Bigg \}, \tag {6} \\ \end{array}
$$
where  $f_{i}^{(k)} = f_{i}(\Psi_{k},\Upsilon_{k})$  for  $i = 1,2;k\in [n]$ .  $\lambda_1,\lambda_2 > 0$ , and  $\mathbb{W}\stackrel {\mathrm{def}}{=}[w_T,w_1,w_2]$  are the parameters of neural networks  $T$ ,  $f_{1},f_{2}$ , respectively. We also follow Song and Ermon (2020) and account for the possible high-variance issues by using a smoothed version of Equation (6), where we clip the  $\exp (\cdot)$  term in the second line between  $\exp (-\rho)$  and  $\exp (\rho)$  where  $\rho >0$ . We refer to Appendix C for the pseudocode of  $S\mathcal{I}_{\mathbb{W}}^*$ . We later analyze the computational complexity of the estimator.
2See detailed mathematical explanation in our proof of Corollary 2.
${}^{3}$  Alternatively,one can take  $\left( {X,Y}\right)  \sim  {\left\{  \left( {X}_{k},{Y}_{k}\right) \right\}  }_{k = 1}^{n}$  ,and  $\left( {\bar{X},\bar{Y}}\right)  \sim  \left\{  \left( {{X}_{k},{Y}_{\sigma \left( k\right) }}\right) \right\}  {}_{k = 1}^{n},\sigma \left( k\right)  \neq  k,\sigma  \in  \mathcal{S}$  (the group of  $n$  -permutations).
Figure 2: Statistical efficiency of dependence measures with structures of varying complexity  $X \in \mathbb{R}^3, Y \in \mathbb{R}^3$ . The AUC ROC quantifies the measure's discriminative ability, where a perfect classifier has AUC ROC of 1.0 while a random classifier has AUC ROC of 0.5. The figures are computed from 100 random runs.
The main idea in quantifying dependence between random variables  $X, Y$  is to characterize the distance between their joint distribution  $P_{X,Y}$  and the product of its marginals  $P_X \otimes P_Y$ . Naturally, this problem has been studied in the area of optimal transportation theory (Mordant and Segers, 2022; Wiesel, 2021; Chuang et al., 2023). For example, Nies et al. (2021) defined the distance as the optimal transport cost on the spaces of the aforementioned distributions for a restricted set of cost functions. Other popular approaches study the distance covariance, which is the weighted Euclidean distance between the joint characteristic function and the product of the marginals' characteristic functions (Székely et al., 2007). Also, Reshef et al. (2011) introduced MIC, a measure of dependence for the 2-variable relationships based on testing multiple grids with varying dimensions and searching for the grid with maximum information between the two variables. Notably, (Liu et al., 2021) benefit from the optimal transportation theory to calculate Squared-loss Mutual Information with a small number of paired samples (drawn from the joint distribution) and a large number of unpaired samples. Lastly, MID (Sugiyama and Borgwardt, 2013) uses the concept of fractal dimensions (Ott, 2002) to define the information dimension and shows benefits in various applications.
In the context of sliced statistical distances, Nguyen et al. (2021) proposed a similar solution to alleviate the problem of using a large number of slices and applied it to the Wasserstein (W-) distance by reformulating the distance as an optimization problem. While we share a similar motivation, the two slicing approaches have significant differences stemming from the fact that Sliced W-distances have only been studied in the context of generative modeling; they only require one slicing direction; and their cost function (1-Wasserstein distance) is continuous, which makes their problem a vanilla expectation maximization (EM) problem. In a different domain, Nguyen et al. (2020) proposed an improvement to relational regularized autoencoders by finding an important area of projections characterized by a von Mises-Fisher (vMF) distribution (Jupp and Mardia, 1979). Optimizing over the family of vMF distributions helps to identify important directions on the sphere ( $\epsilon$  parameter) and how much weight to put there compared to other directions ( $\kappa$  parameter). We, however, do not restrict our optimization to a particular family of distributions since the optimal slicing policy might not take the form of vMF.
We conduct extensive experiments to demonstrate  $\mathcal{S}\mathcal{I}^*$ 's efficacy. We compare  $\mathcal{S}\mathcal{I}^*$ 's dependence quantification quality to a diverse set of competitive baselines. Through careful analyses, we show  $\mathcal{S}\mathcal{I}^*$ 's scalability, computation speed, sample efficiency, and slicing efficiency. We further provide empirical evidence that  $\mathcal{S}\mathcal{I}^*$  works excellently across challenging representation learning and rein
Figure 3: Statistical efficiency of  $\mathcal{S}\mathcal{I}^*$ , SI, and MI with dimension, number of samples  $(n)$ , and number of slices  $(m)$ . The figures with respect to  $m$  do not include MI because this method does not use slices.
forcement learning tasks. We refer to Appendix B for further analysis experiments and Appendix D for implementation details and hyperparameters provided at https://bit.ly/3foLke2.
6.1 Effectiveness of  $S\mathcal{I}^*$  as a Dependence Measure
We first illustrate the superiority of  $\mathcal{S}\mathcal{I}^*$  as a dependence measure by comparing against the following state-of-the-art measures: Mutual Information (MI) (Shannon, 1948; Kraskov et al., 2004), Sliced Mutual Information (SI) (Goldfeld and Greenewald, 2021), distributional sliced Wasserstein Dependence (WD) (Nguyen et al., 2021), Transport Dependency Coefficient (TD) (a.k.a Transport Correlation) (Nies et al., 2021), Mutual Information Dimension (MID) (Sugiyama and Borgwardt, 2013), Maximal Information Coefficient (Reshef et al., 2011), and Distance Correlation (dCor) (Székely et al., 2007). We perform ROC curve analysis to examine precision and recall in recognizing complex relationship types with the existence of uniformly distributed noise. We test on synthetic data that comprises positive samples that depict extremely complicated structures of  $P_{X,Y} \in \mathcal{P}(\mathbb{R}^3 \times \mathbb{R}^3)$ . In this setting, we conduct 10 statistical tests and cover a wide range of structures, including sinusoidal features, elliptical surfaces, spherical harmonic surfaces, and hyper-hypocycloids. We provide the full details (e.g. the mappings formulae) in Appendix D.2.
Results are reported in Figure 2 as a function of the noise ratio. Notably,  $S\mathcal{I}^*$  surpassed all other dependency measures in discerning structure from noise and sometimes by a significant margin (up to  $30\%$  improvement in accuracy). Although  $S\mathcal{I}^*$  has shown similar performance to some of the baselines in the linear and parabolic settings, a clear superiority of  $S\mathcal{I}^*$  is manifested in the eight complicated geometries, such as the spherical harmonics and sinusoidals. As a result, Figure 2 clearly shows the advantages of our slicing optimality in highly non-trivial settings.
6.2 Sample and Slicing Efficiency of  $\mathcal{S}\mathcal{I}^*$
Supporting our theory in Section 3.3, we show that  $S\mathcal{I}^*$  scales efficiently to higher dimensions. We test against SI and MI estimators (Kriskov et al., 2004) on two structures of common signal and elliptical. We calculate the AUC of the ROC while varying the number of samples or projections. Figure 3 shows that  $S\mathcal{I}^*$  attains the perfect accuracy with less number of slices compared to the uniform-sampling-based baselines. As such, this empirical result supports our claim that  $S\mathcal{I}^*$  succeeds in detecting the dependence between the variables with an improved sample and projection efficiency. Appendix B.2 provides additional experiments in higher dimensions.
6.3 Result on Representation Learning
Representation learning employs MI to discover useful representations by training a neural network encoder to maximize MI between its inputs and outputs. For example, Hjelm et al. (2019) introduces Deep InfoMax (DIM), a popular method for learning unsupervised representations, which improves the learning by including knowledge about the local structure of the input into the objective. We follow DIM settings and substitute MI with  $S\mathcal{I}$  and  $S\mathcal{I}^*$  as the information measure, respectively. We test these three methods along with BiGAN (Donahue et al., 2016) on the STL-10 (Coates et al., 2011) and CIFAR10 (Krizhevsky et al., 2009) datasets, which consist of high-dimensional images. Results in Tables 1 and 2 show that  $S\mathcal{I}^*$  outperforms the baselines for all classifiers: high-level vector
representation  $(Y)$ , the output of the previous fully-connected layer (fc) and the last convolutional layer (conv). As such, this experiment highlights the scalability aspect of  $\mathcal{SL}^*$  with high-dimensional variables not to mention its excellent performance on a challenging ML task.
Table 1: Classification accuracy (%) on STL-10
conv	fc	Y
BiGAN	71.53	67.18	58.48
DIM (MI)	69.15	63.81	61.92
DIM (SI)	74.54	71.34	68.90
DIM (SI*)	76.89	71.67	70.04
Table 2: Classification accuracy (%) on CIFAR10
conv	fc	Y
BiGAN	62.57	62.74	52.54
DIM (MI)	72.66	70.66	64.71
DIM (SI)	74.37	70.23	65.99
DIM (SI*)	77.01	70.39	69.04
6.4 Result on Reinforcement Learning
Previous frameworks have studied MI objectives as regularizers to the reinforcement learning (RL) objective that involves high-dimensional variables of states, actions, and rewards (Nachum et al., 2019; Schwarzer et al., 2021). Since it is practically difficult to estimate MI in high dimensions, most works resort to learning representations of the high-dimensional data by projecting them into low-dimensional embeddings. Thanks to the slicing technique, these representations are no longer needed when using  $S\mathcal{I}^*$ . We empirically evaluate the scalability of  $S\mathcal{I}^*$  with an RL task, where we adapt the forward information objective in Rakelly et al. (2021) as:
$$
\mathbb {J} _ {\mathcal {S I} ^ {*}} ^ {\pi} = \max  _ {\pi} \mathcal {S I} ^ {*} \left(S _ {t + 1}; \left[ A _ {t}, S _ {t} \right]\right), \tag {7}
$$
where  $\pi$  denotes the agent's policy,  $A_{t} \sim \pi(S_{t})$ ,  $S_{t+1} \sim \mathbb{T}(S_{t}, A_{t})$ , and  $\mathbb{T}$  is the transition dynamics. For SI, we use the same objective (by replacing  $\mathcal{ST}^*$  with the dependence measure  $\mathcal{ST}(:, \cdot)$ ). We run tests on three challenging Gym environments (Brockman et al., 2016), where the environment dimensions (observation dim, action dim) are as follows: Humanoid-v3: (378, 17); Ant-v3: (113, 8); Hopper-v3: (12, 3). Further details can be found in Appendix D.3. We report the results comparing  $\mathcal{ST}^*$ , SI, and the original paper objective (MI). Figure 4 shows clear gains (up to  $25\%$ ) with a noticeable learning speed of  $\mathcal{ST}^*$  compared to the baselines.
(a) Humanoid-v3
(b) Ant-v3
(c) Hopper-v3
Figure 4: Results on high-dimensional Reinforcement Learning environments show a cutting-edge performance of  $\mathcal{S}\mathcal{I}^*$  on complex high-dimensional control problems. The mean and variance computed for 12 seeds are shown in the figures.
6.5 Behavior Analysis
Convergence of  $S\mathcal{I}^*$ . We validate the convergence rate of  $S\mathcal{I}^*$  (Theorem 4). In Figure 5, we show the RMSE between the ground truth  $S\mathcal{I}^*$  and the estimated in the setting where  $X$  and  $Y$  are normal random variables with 5 overlapping entries:  $Z \sim \mathcal{N}(0, I_{15})$ ,  $\Lambda \sim \mathcal{N}(0, 0.1I_{10})$ ,  $X = Z_{[1,10]}$ ,  $Y = Z_{[6,15]} + \Lambda$ . The heat map in Figure 5a shows convergence when  $n, m$  (# samples and # slices, respectively) vary independently, while Figure 5b shows when one parameter is varying and the other is set to  $10^3$ . The MI estimator used is the Kozachenko-Leonenko estimator (Kraskov et al., 2004). We also test with  $X, Y \in \mathbb{R}^d$  with linear dependence to study the effect of the dimension on the convergence. Results are reported in Figure 5c with  $n$  and  $m$  varying together.
Computational complexity. Results in Figure 6 (left) show that the computational complexity of  $S\mathcal{I}^*$  is of an order similar to that of  $S\mathcal{I}$ . Note that, while  $S\mathcal{I}^*$  is slower than  $S\mathcal{I}$ , the difference is
(b)
Figure 5: Figure shows the RMSE with varying  $n$ ,  $m$  and variables dimensions.
insignificant not to mention the clear advantage of  $\mathcal{S}\mathcal{I}^*$  over other baselines given its superiority in detecting complex dependencies. We also analyze the effect of number of gradient update steps of  $f_{1}, f_{2}$  on the computation speed (Figure 6 (right)). Further results can be found in Appendix B.
Figure 6: (Left) Computation speed of measures based on the number of samples and slices (dependency is linear and noise ratio  $= 0.6$ ). (Right) Measures with respect to the number of update steps ( $m = 700$ ).
We presented a novel dependence measure, called  $S\mathcal{I}^*$ , that is scalable to high dimensions while being efficient regarding time, sample, and slicing complexity. We discussed its theoretical properties, proved that its estimation error depends on the problem dimensions only up to a constant factor, and empirically validated the competence of  $S\mathcal{I}^*$  in detecting complicated dependencies against state-of-the-art dependence measures. We further placed  $S\mathcal{I}^*$  into modern ML, where we proved the adequacy of our dependence measure on a more onerous set of tasks.
Limitations. Since the aim of this paper was to introduce optimality into slicing methods,  $\mathcal{S}\mathcal{I}^*$  is formalized to use two slicing variables. Although the current setting has proved effective and rather sufficient in complex machine learning scenarios, the extension into multivariate settings is encouraged for future works.
Broader impact. The findings and methodologies developed through this work can significantly improve the way we understand and analyze systems with complex relations. We believe this work will have noticeable impact on broader communities, since the study of quantifying information is active in research areas including physics, statistics, computational biology, economics, and neuroscience. At its current form, this work is mainly theoretical and does not impose any negative societal impact.
Reshef, D. N., Reshef, Y. A., Finucane, H. K., Grossman, S. R., McVean, G., Turnbaugh, P. J., Lander, E. S., Mitzenmacher, M., and Sabeti, P. C. (2011). Detecting novel associations in large data sets. science, 334(6062):1518-1524.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. (2017). Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347.
Schwarzer, M., Anand, A., Goel, R., Hjelm, R. D., Courville, A., and Bachman, P. (2021). Data-efficient reinforcement learning with self-predictive representations. In International Conference on Learning Representations.
Shannon, C. E. (1948). A mathematical theory of communication. The Bell system technical journal, 27(3):379-423.
Sonderby, C. K., Caballero, J., Theis, L., Shi, W., and Huszar, F. (2016). Amortised map inference for image super-resolution. arXiv preprint arXiv:1610.04490.
Song, J. and Ermon, S. (2020). Understanding the limitations of variational mutual information estimators. In International Conference on Learning Representations.
Sugiyama, M. and Borgwardt, K. M. (2013). Measuring statistical dependence via the mutual information dimension. In Twenty-Third International Joint Conference on Artificial Intelligence. Citeseer.
Székely, G. J., Rizzo, M. L., and Bakirov, N. K. (2007). Measuring and testing dependence by correlation of distances. The annals of statistics, 35(6):2769-2794.
Todorov, E., Erez, T., and Tassa, Y. (2012). Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 5026-5033.
Villani, C. (2009). Optimal transport: old and new, volume 338. Springer.
Wiesel, J. (2021). Measuring association with Wasserstein distances. arXiv preprint arXiv:2102.00356.
Yang, P. and Chen, B. (2018). Robust kullback-leibler divergence and universal hypothesis testing for continuous distributions. IEEE Transactions on Information Theory, 65(4):2360-2373.
Yoo, A. B., Jette, M. A., and Grondona, M. (2003). Slurm: Simple linux utility for resource management. In Workshop on job scheduling strategies for parallel processing, pages 44-60. Springer.