SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Future work can leverage PID for targeted dataset creation, representation learning optimized for PID values, and applications of information theory to higher-dimensional data. More broadly, there are several exciting directions in investigating more applications of multivariate information theory in modeling feature interactions, predicting multimodal performance, and other tasks involving feature interactions such as privacy-preserving and fair representation learning from high-dimensional data [28, 42]. Being able to provide guarantees for fairness and privacy-preserving learning can be particularly impactful.
This material is based upon work partially supported by Meta, National Science Foundation awards 1722822 and 1750439, and National Institutes of Health awards R01MH125740, R01MH132225, R01MH096951 and R21MH130767. PPL is supported in part by a Siebel Scholarship and a Waibel Presidential Fellowship. RS is supported in part by ONR grant N000142312368 and DARPA FA87502321015. One of the aims of this project is to understand the comfort zone of people for better privacy and integrity. Any opinions, findings, conclusions, or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the sponsors, and no official endorsement should be inferred. Finally, we would also like to acknowledge feedback from the anonymous reviewers who significantly improved the paper and NVIDIA's GPU support.
[1] Akshay Agrawal, Robin Verschueren, Steven Diamond, and Stephen Boyd. A rewriting system for convex optimization problems. Journal of Control and Decision, 5(1):42-60, 2018.
[2] Paul D Allison. Testing for interaction in multiple regression. American journal of sociology, 1977.
[3] Brandon Amos. Tutorial on amortized optimization for learning to optimize over continuous domains. arXiv preprint arXiv:2202.00665, 2022.
[4] Galen Andrew, Raman Arora, Jeff Bilmes, and Karen Livescu. Deep canonical correlation analysis. In International conference on machine learning, pages 1247-1255. PMLR, 2013.
[5] Stanislaw Antol, Aishwarya Agrawal, Jiasen Lu, Margaret Mitchell, Dhruv Batra, C Lawrence Zitnick, and Devi Parikh. Vqa: Visual question answering. In Proceedings of the IEEE international conference on computer vision, pages 2425-2433, 2015.
[6] MOSEK ApS. MOSEK Optimizer API for Python 10.0.34, 2022.
[7] Benjamin Auffarth, Maite López, and Jesús Cerquides. Comparison of redundancy and relevance measures for feature selection in tissue classification of ct images. In Industrial conference on data mining, pages 248-262. Springer, 2010.
[8] Maria-Florina Balcan, Avrim Blum, and Ke Yang. Co-training and expansion: Towards bridging theory and practice. Advances in neural information processing systems, 17, 2004.
[9] Tadas Baltrusaitis, Chaitanya Ahuja, and Louis-Philippe Morency. Multimodal machine learning: A survey and taxonomy. IEEE transactions on pattern analysis and machine intelligence, 2018.
[10] Reuben M Baron and David A Kenny. The moderator-mediator variable distinction in social psychological research: Conceptual, strategic, and statistical considerations. Journal of personality and social psychology, 51(6):1173, 1986.
[11] Jan Beirlant, Edward J Dudewicz, László Györfi, Edward C Van der Meulen, et al. Nonparametric entropy estimation: An overview. International Journal of Mathematical and Statistical Sciences, 1997.
[12] Anthony J Bell. The co-information lattice. In Proceedings of the fifth international workshop on independent component analysis and blind signal separation: ICA, volume 2003, 2003.
[13] Nils Bertschinger, Johannes Rauh, Eckehard Olbrich, Jürgen Jost, and Nihat Ay. Quantifying unique information. Entropy, 16(4):2161-2183, 2014.
[14] Avrim Blum and Tom Mitchell. Combining labeled and unlabeled data with co-training. In Proceedings of the eleventh annual conference on Computational learning theory, pages 92-100, 1998.
[15] Cameron W Brennan, Roel GW Verhaak, Aaron McKenna, Benito Campos, Houtan Noushmehr, Sofie R Salama, Siyuan Zheng, Debyani Chakravarty, J Zachary Sanborn, Samuel H Berman, et al. The somatic genomic landscape of glioblastoma. cell, 155(2):462-477, 2013.
[16] Nancy Ann Oberheim Bush, Susan M Chang, and Mitchel S Berger. Current and future strategies for treatment of glioma. Neurosurgical review, 40(1):1-14, 2017.
[17] Santiago Castro, Devamanyu Hazarika, Verónica Pérez-Rosas, Roger Zimmermann, Rada Mihalcea, and Soujanya Poria. Towards multimodal sarcasm detection (an _obviously_ perfect paper). In ACL, 2019.
[18] CDC. CDC WONDER: Underlying cause of death, 1999-2019. 2020.
Figure 6: The PID framework formalizes the nature of feature interactions for a task. Through our proposed scalable estimators for redundancy, uniqueness, and synergy, we demonstrate their usefulness towards quantifying multimodal datasets and multimodal model predictions, which together yield principles for model selection.
A Partial Information Decomposition
Information theory formalizes the amount of information that a single variable  $(X)$  provides about another  $(Y)$ , and is quantified by Shannon's mutual information [86]  $I(X;Y)$  which measures the amount of uncertainty reduced from  $H(Y)$  to  $H(Y|X)$  when given  $X$  as input. However, the direct extension of information theory to 3 or more variables, such as through total correlation [106, 34, 90, 97] or interaction information [70, 93, 12, 35], both have significant shortcomings. In particular, the three-way mutual information  $I(X_{1};X_{2};Y)$  can be both negative and positive, leading to considerable difficulty in its interpretation by theoreticians and practitioners.
Partial information decomposition (PID) [108] was proposed as a method to generalize information theory to multiple variables, by positing a decomposition of the total information 2 variables provide about a task  $I(X_{1},X_{2};Y)$  into 4 quantities: redundancy  $R$  between  $X_{1}$  and  $X_{2}$ , unique information  $U_{1}$  in  $X_{1}$  and  $U_{2}$  in  $X_{2}$ , and synergy  $S$  that should together satisfy the following consistency equations:
$$
R + U _ {1} = I \left(X _ {1}; Y\right), \tag {7}
$$
$$
R + U _ {2} = I \left(X _ {2}; Y\right) \tag {8}
$$
$$
U _ {1} + S = I \left(X _ {1}; Y \mid X _ {2}\right), \tag {9}
$$
$$
U _ {2} + S = I \left(X _ {2}; Y \mid X _ {1}\right) \tag {10}
$$
$$
R - S = I \left(X _ {1}; X _ {2}; Y\right) \tag {11}
$$
A definition of  $R$  was first proposed by Williams and Beer [108] and was subsequently improved by Bertschinger et al. [13], Griffith and Koch [41] giving:
$$
R = \max  _ {q \in \Delta_ {p}} I _ {q} \left(X _ {1}; X _ {2}; Y\right) \tag {12}
$$
$$
U _ {1} = \min  _ {q \in \Delta_ {p}} I _ {q} \left(X _ {1}; Y \mid X _ {2}\right) \tag {13}
$$
$$
U _ {2} = \min  _ {q \in \Delta_ {p}} I _ {q} \left(X _ {2}; Y \mid X _ {1}\right) \tag {14}
$$
$$
S = I _ {p} \left(X _ {1}, X _ {2}; Y\right) - \min  _ {q \in \Delta_ {p}} I _ {q} \left(X _ {1}, X _ {2}; Y\right) \tag {15}
$$
where  $\Delta_p = \{q \in \Delta : q(x_i, y) = p(x_i, y) \forall y, x_i, i \in \{1, 2\}\}$  and  $\Delta$  is the set of all joint distributions over  $X_1, X_2, Y$ . We call this the set of marginal-matching distributions. The core intuition of these definitions and how they enable PID lies in optimization over  $\Delta_p$  instead of simply estimating information-theoretic measures for the observed distribution  $p$ . Specifically, synergy  $S$  is the difference between multimodal information in  $p$  and total multimodal information in  $q^* \in \Delta_p$ , the non-synergistic distribution (see Figure 7 for an illustration). While there have been other definitions of PID that also satisfy the consistency equations, Bertschinger et al. [13] showed that their proposed measure satisfies several desirable properties. Optimizing over  $\Delta_p$  can be seen as going beyond purely observed data and discovering latent redundancy, uniqueness, and synergy via representation learning.
According to Bertschinger et al. [13], it suffices to solve for  $q$  using the following max-entropy optimization problem, the same  $q$  equivalently solves any of the remaining problems defined for
Unimodal marginal-matching distributions:
$$
\Delta_ {p} = \left\{q \left(x _ {1}, x _ {2}, y\right): q \left(x _ {1}, y\right) = p \left(x _ {1}, y\right), q \left(x _ {2}, y\right) = p \left(x _ {2}, y\right) \right\}
$$
Figure 7: We illustrate the set of marginal-matching distributions  $q \in \Delta_p$ , and an intuition behind the definition of synergy  $S$  as the difference between total multimodal information in  $p$  and total multimodal information in  $q^* \in \Delta_p$ , the non-synergistic distribution.
redundancy, uniqueness, and synergy.
$$
q ^ {*} = \underset {q \in \Delta_ {p}} {\arg \max } H _ {q} (Y | X _ {1}, X _ {2}) \quad (\text {m a x - e n t r o p y}) \tag {16}
$$
$$
q ^ {*} = \underset {q \in \Delta_ {p}} {\arg \max } I _ {q} \left(X _ {1}; X _ {2}; Y\right) \quad (\text {r e d u n d a n c y}) \tag {17}
$$
$$
q ^ {*} = \underset {q \in \Delta_ {p}} {\arg \min } I _ {q} (\{X _ {1}, X _ {2} \}; Y) \quad \text {(s y n e r g y)} \tag {18}
$$
$$
q ^ {*} = \underset {q \in \Delta_ {p}} {\arg \min } I _ {q} \left(X _ {1}; Y \mid X _ {2}\right) \quad (\text {u n i q u e i n} X _ {1}) \tag {19}
$$
$$
q ^ {*} = \underset {q \in \Delta_ {p}} {\arg \min } I _ {q} \left(X _ {2}; Y \mid X _ {1}\right) \quad (\text {u n i q u e i n} X _ {2}) \tag {20}
$$
B Scalable PID Estimators
This section provides algorithmic and implementation details on the estimation of PID values for high-dimensional continuous datasets.
B.1 Estimator 1: CVX for Discrete Representations
Our first proposed estimator, CVX, is exact, based on convex optimization, and is able to scale to problems where  $|\mathcal{X}_i|$  and  $|\mathcal{Y}|$  are around 100. We show a high-level diagram of CVX in Figure 8 and provide several implementation details here:
Implementation Details: We implemented (5) using CVXPY [25, 1]. The transformation from the max-entropy objective (16) to (5) ensures adherence to disciplined convex programs [40], thus allowing CVXPY to recognize it as a convex program. All 3 conic solvers, ECOS [27], SCS [75], and MOSEK [6] were used, with ECOS and SCS being default solvers packaged with CVXPY. Our experience is that MOSEK is the fastest and most stable solver, working out of the box without any tuning. However, it comes with the downside of being commercial. For smaller problems, ECOS and SCS work just fine.
Potential Numerical Issues: Most conic solvers will suffer from numerical instability when dealing with ill-conditioned problems, typically resulting from coefficients that are too large or small. It is fairly common, especially in the GMM experiment that a single bin contains just one sample. Since  $P$  is estimated from an empirical frequency matrix, it will contain probabilities of the order 1/num-samples. When the number of samples is too large (e.g.,  $\geq 1e8$ ), this frequency is too small and causes all 3 solvers to crash.
Approximating Continuous Representation using Histogramming: The number and width of bins can affect quality of PID estimation [77]. For example, it is known that with a fixed number of samples, making the width of bins arbitrarily small will cause KL estimates to diverge. It is known that the number of bins should grow sub-linearly with the number of samples. Rice [84] suggest setting the number of bins to be the cubed-root of the number of samples. In our synthetic experiments, we sampled at least 1e6 samples and used no more than 100 bins. Note that increasing the number of samples ad infinitum is not an option due to potential numerical issues described above.
Figure 8: Overview of CVX estimator. Appropriate histogramming of each  $X_{i}$  or clustering in preprocessed feature space is first performed to constrain  $|\mathcal{X}_i|$ . By rewriting the entropy objective as a KL divergence (with auxiliary variables) which is recognized as convex, this allows us to use conic solvers such as SCS [75], ECOS [27], and MOSEK [6] in CVXPY to solve for  $q^{*}$  and obtain PID values.
B.2 Alternate Methods for Working with Continuous  $X_{1}, X_{2}$ :
Another method for discretization is based on nearest neighbors [80, 11, 104], which has the advantage of performing density estimation as an intermediate step, but is ill-suited for our purposes of approximating  $q^*$ . Pakman et al. [76] attempt to optimize  $q^*$  directly in the space of continuous distribution by approximating an optimal copula, which canonically captures dependencies between  $X_1, X_2$  in  $q$ . However, this is nontrivial due to the space of possible joint distributions, and the authors restrict themselves to Gaussian copulas.
B.3 Estimator 2: BATCH for Continuous Representations
Our next estimator, BATCH, is suitable for large datasets where the support of  $X_{1}$  and  $X_{2}$  are high-dimensional continuous vectors and not enumerable. Instead, we propose an end-to-end model parameterizing the joint distribution and a training objective whose solution allows us to obtain PID. Here  $|\mathcal{X}_i| = |\mathcal{Y}| = \infty$ . However, we wish to estimate the PID values given a sampled dataset  $\mathcal{D} = \{(x_1^{(j)}, x_2^{(j)}, y^{(j)})\}$  of size  $n$ .
In this appendix section, we provide full details deriving our BATCH estimator, see Figure 9 for a high-level overview. We first rewrite the optimization objective in (18) before defining neural estimators to approximate high-dimensional variables that we cannot compute exactly, before optimizing the (18) using learned models and obtaining the PID values.
Objective function: First, we observe that the optimization objective in (18) can be written as:
$$
\begin{array}{l} q ^ {*} = \underset {q \in \Delta_ {p}} {\arg \min } I _ {q} \left(\left\{X _ {1}, X _ {2} \right\}; Y\right) (21) \\ = \arg \min  _ {q \in \Delta_ {p}} \mathbb {E} _ {x _ {1}, x _ {2}, y \sim q} \left[ \log \frac {q \left(x _ {1} , x _ {2} , y\right)}{q \left(x _ {1} , x _ {2}\right) q (y)} \right] (22) \\ = \arg \min  _ {q \in \Delta_ {p}} \mathbb {E} _ {x _ {1}, x _ {2}, y \sim q} \left[ \log \frac {q \left(x _ {1} , x _ {2} \mid y\right)}{q \left(x _ {1} , x _ {2}\right)} \right] (23) \\ = \arg \min  _ {q \in \Delta_ {p}} \mathbb {E} _ {x _ {1}, x _ {2}, y \sim q} \left[ \log \frac {q \left(x _ {2} \mid x _ {1} , y\right) q \left(x _ {1} \mid y\right)}{q \left(x _ {2} \mid x _ {1}\right) q \left(x _ {1}\right)} \right] (24) \\ = \arg \min  _ {q \in \Delta_ {p}} \mathbb {E} _ {\substack {x _ {1}, y \sim q (x _ {1}, y) \\ x _ {2} \sim q (x _ {2} | x _ {1}, y)}} \left[ \log \frac {q \left(x _ {2} | x _ {1} , y\right) q \left(x _ {1} | y\right)}{\sum_ {y ^ {\prime} \in Y} q \left(x _ {2} | x _ {1} , y ^ {\prime}\right) q \left(y ^ {\prime} | x _ {1}\right) q \left(x _ {1}\right)} \right] (25) \\ \end{array}
$$
Our goal is to approximate the optimal  $\tilde{q}$  that satisfies this problem. Note that by the marginal constraints,  $q(x_{1},y) = p(x_{1},y)$  and  $q(y^{\prime}|x_{1}) = p(y^{\prime}|x_{1})$ . The former can be obtained by sampling from the dataset and the latter is an unimodal model independent of  $q$ .
Unnormalized joint distribution
Figure 9: Overview of BATH estimator: We define an end-to-end model that takes in batches of features  $\mathbf{X}_1\in \widetilde{\mathcal{X}}_1^m$ $\mathbf{X}_2\in \widetilde{\mathcal{X}}_2^m$ $\mathbf{Y}\in \mathcal{V}^m$  and learns a matrix  $A\in \mathbb{R}^{m\times m\times |\mathcal{V}|}$  to represent the unnormalized joint distribution  $\tilde{q}$ . Following ideas in multimodal alignment [99] and attention methods [103] that learn approximate distributions over input data, we parameterize  $A$  by computing the outer-product similarity matrix over learned features. To ensure that  $\tilde{q}\in \Delta_p$ , we use an unrolled version of Sinkhorn's algorithm [24] which projects  $A$  onto  $\Delta_p$  by iteratively normalizing all rows and columns to sum to 1 and rescaling to satisfy the marginals  $\hat{p}$ . We plug in the learned  $\tilde{q}$  into our objective in (25) and update trainable neural network parameters  $\phi$  using projected gradient descent. Upon convergence, we extract PID by approximating  $I_{\tilde{q}}(\{X_1,X_2\};Y)$  by sampling and plugging into (1)-(3). This use of a mini-batch of size  $m$  can be seen as an approximation of full-batch gradient descent over the entire dataset of size  $n\gg m$ . While it is challenging to obtain an unbiased estimator of the full-batch gradient since computing the full  $A$  is intractable, we found our approach to work in practice for large  $m$ . Our approach can also be informally viewed as performing amortized optimization [3] by using  $\phi$  to implicitly share information about the full batch using subsampled batches.
Neural architecture: To scale to continuous or unseen  $x_{i}$ 's, we define a function approximator that takes in  $x_{1}, x_{2}$  and  $y$  and outputs an unnormalized joint density  $\tilde{q}(x_{1}, x_{2}, y)$ . Given batches  $\mathbf{X}_{1} \in \widetilde{\mathcal{X}}_{1}^{m}, \mathbf{X}_{2} \in \widetilde{\mathcal{X}}_{2}^{m}, \mathbf{Y} \in \mathcal{Y}^{m}$ , we define learnable encoders  $f_{\phi(1)}$  and  $f_{\phi(2)}$  that output a matrix  $A \in \mathbb{R}^{m \times m \times |\mathcal{Y}|}$  to represent the unnormalized joint distribution  $\tilde{q}$ , i.e., we want  $A[i][j][y] = \tilde{q}(\mathbf{X}_{1}[i], \mathbf{X}_{2}[j], y)$ . Following ideas in multimodal alignment [99] and attention methods [103] in deep learning that learn approximate distributions over input data, we also parameterize our neural architecture by computing the outer-product similarity matrix over learned features:
$$
A = \exp \left(f _ {\phi (1)} \left(\mathbf {X} _ {\mathbf {1}}, y\right) f _ {\phi (2)} \left(\mathbf {X} _ {\mathbf {2}}, y\right) ^ {\top}\right) \tag {26}
$$
Constraining  $A$  using Sinkhorn's algorithm: To ensure that the learned  $A$  is a valid pdf and respects the marginal constraints  $\tilde{q}(x_i, y) = p(x_i, y)$ , we use the Sinkhorn-Knopp algorithm [24], which projects  $A$  into the space of non-negative square matrices that are row and column-normalized to satisfy the desired marginals  $p(x_1, y)$  and  $p(x_2, y)$ . Sinkhorn's algorithm enables us to perform this projection by iteratively normalizing all rows and columns of  $A$  to sum to 1 and rescaling the rows and columns to satisfy the marginals. Since  $p(x_i, y)$  is not easy to estimate for continuous  $x_i$ , we compute the marginals using Bayes' rule by expanding it as  $p(x_i)p(y|x_i)$ . Note that by sampling  $x_i$  from the dataset, the rows and columns of  $A$  are already distributed as  $p(x_i)$ . The only remaining term needed is  $p(y|x_i)$ , for which we use unimodal models  $\hat{p}(y|x_i)$  trained before running the estimator and subsequently frozen. Finally, Sinkhorn's algorithm is performed by normalizing each row to  $\hat{p}(y|x_1)$  and normalizing each column to  $\hat{p}(y|x_2)$ .
**Optimization of the objective:** We obtain the value of  $\tilde{q}(x_2|x_1,y)$  by marginalizing from the above distribution  $\tilde{q}(x_1,x_2|y)$ . Given matrix  $A$  representing  $\tilde{q}(x_1,x_2,y)$ , the remaining terms can all be computed in closed form through appropriate summation across dimensions in  $A$  to obtain  $\tilde{q}(x_i)$ ,
$\tilde{q}(x_1, x_2)$ ,  $\tilde{q}(x_i | y)$ , and  $\tilde{q}(x_1, x_2 | y)$ . We now have all the terms for the objective in (25). We optimize the objective via gradient descent.
Our final algorithm for training BATCH is outlined in Algorithm 1.
Extracting PID values from the learned  $\tilde{q}$ : To extract the PID values from a learned distribution of  $\tilde{q}$ , we first expand the definition of mutual information under  $p$  and  $\tilde{q}$ , and calculate the PID values from their definitions in (1)-(3):
$$
I _ {p} (Y; X _ {i}) = \underset {x _ {i}, y \sim p} {\mathbb {E}} \left[ \log \left(\frac {\hat {p} (y \mid x _ {i})}{\hat {p} (y)}\right) \right] \tag {27}
$$
$$
I _ {p} (Y; X _ {1}, X _ {2}) = \underset {x _ {1}, x _ {2}, y \sim p} {\mathbb {E}} \left[ \log \left(\frac {\hat {p} (y \mid x _ {1} , x _ {2})}{\hat {p} (y)}\right) \right] \tag {28}
$$
$$
I _ {\tilde {q}} (Y; X _ {1}, X _ {2}) = \underset {x _ {1}, x _ {2}, y \sim \tilde {q}} {\mathbb {E}} \left[ \log \left(\frac {\tilde {q} \left(x _ {2} \mid x _ {1} , y\right) \hat {p} \left(y \mid x _ {1}\right)}{\hat {p} (y) \sum_ {y ^ {\prime}} \tilde {q} \left(x _ {2} \mid x _ {1} , y ^ {\prime}\right) \hat {p} \left(y ^ {\prime} \mid x _ {1}\right)}\right) \right] \tag {29}
$$
$$
R = I _ {\tilde {q}} (Y; X _ {1}, X _ {2}) \tag {30}
$$
$$
U _ {1} = I _ {\tilde {q}} (Y; X _ {1}, X _ {2}) - I _ {p} (Y; X _ {2}) \tag {31}
$$
$$
U _ {2} = I _ {\tilde {q}} (Y; X _ {1}, X _ {2}) - I _ {p} (Y; X _ {1}) \tag {32}
$$
$$
S = I _ {p} \left(Y; X _ {1}, X _ {2}\right) - I _ {\tilde {q}} \left(Y; X _ {1}, X _ {2}\right) \tag {33}
$$
Algorithm 1 BATCH algorithm.
Require: Multimodal dataset  $\mathbf{X}_1\in \tilde{\mathcal{X}}_1^n$ $\mathbf{X}_2\in \tilde{\mathcal{X}}_2^n$ $\mathbf{Y}\in \mathcal{Y}^n$
Initialize feed-forward networks  $f_{\phi (1)},f_{\phi (2)}$    
while not converged do for sampled batch  $\mathbf{X_1}\in \widetilde{\mathcal{X}}_1^m$ $\mathbf{X}_2\in \widetilde{\mathcal{X}}_2^m$ $\mathbf{Y}\in \mathcal{V}^m$  do  $h_1[:,y]\gets f_{\phi (1)}(x_1,y),h_2[:,y]\gets f_{\phi (2)}(x_2,y)$ $A(:,.,y]\gets \exp (h_1(:,y]h_2(:,y)]^\top)$ $\tilde{p} (y|x_i)\gets$  unimodal predictions.  $A\gets \mathrm{SINKHORN}_{\hat{p}}(A)$  Marginalize and sample  $(x_{1},x_{2},y)$  from  $A$  Calculate the loss from (25) using  $(x_{1},x_{2},y)$  Perform a gradient step on the loss. end for end while return network weights in  $f_{1,j},f_{2,j}$  for  $j\in [C]$
Important assumptions and caveats: As a note, the above formulation assumes the existence of unimodal models  $\hat{p}(y|x_i)$ . We train these models independently and freeze them when training the above estimator. After obtaining  $\tilde{q}$ , we extract the PID values. To do so, we need an additional multimodal model to estimate  $\hat{p}(y|x_1,x_2)$  due to (28). Because part of our goal is multimodal model selection, we must assume that the existing multimodal model might not be expressive enough and could underestimate (28). However, this value is only subsequently used to calculate  $S$ , leading to a possibly underestimated  $S$ , but leaving all other PID values unaffected by the choice of the multimodal model. In practice, the possibility of  $S$  being larger than estimated should be considered, and more complex models should be investigated until  $S$  stops increasing. Empirically, we have observed that measured value of  $S$  by BATCH is generally larger than that measured by CVX; the underestimation is not as large as in the case of clustering in CVX.
In addition to the above, we note that the learned joined distribution  $\tilde{q}(x_1, x_2, y)$  is limited by the expressivity of the neural network and similarity matrix method used in our model. Empirically, the PID values are not sensitive to hyperparameters, as long as the network is trained to convergence. Finally, our approximation using batch sampling limits the possible joined distributions to alignments within a batch, which we alleviates by using large batch size (256 in our experiments).
Implementation Details: We separately train the unimodal and multimodal models from the dataset using the best model from Liang et al. [60] and freeze them when running Algorithm 1. For the