Statistically Valid Information Bottleneck via Multiple Hypothesis Testing

TL;DR

A statistically valid solution to the information bottleneck problem, referred to as IB via multiple hypothesis testing (IB-MHT), which ensures that the learned features meet the IB constraints with high probability, regardless of the size of the available dataset is introduced.

Abstract

The information bottleneck (IB) problem is a widely studied framework in machine learning for extracting compressed features that are informative for downstream tasks. However, current approaches to solving the IB problem rely on a heuristic tuning of hyperparameters, offering no guarantees that the learned features satisfy information-theoretic constraints. In this work, we introduce a statistically valid solution to this problem, referred to as IB via multiple hypothesis testing (IB-MHT), which ensures that the learned features meet the IB constraints with high probability, regardless of the size of the available dataset. The proposed methodology builds on Pareto testing and learn-then-test (LTT), and it wraps around existing IB solvers to provide statistical guarantees on the IB constraints. We demonstrate the performance of IB-MHT on classical and deterministic IB formulations, including experiments on distillation of language models. The results validate the effectiveness of IB-MHT in outperforming conventional methods in terms of statistical robustness and reliability.

Figures (9)

• Figure 1: Illustration of the information bottleneck (IB) setup.
• Figure 2: Illustration of the operations of IB-MHT: ① The calibration data set $\mathcal{D}$ is split into two disjoint subsets $\mathcal{D}_\text{OPT}$ and $\mathcal{D}_\text{MHT}$. ② The Pareto frontier in the plane $(I(T;Y),I(X;T))$ is estimated by using the mutual information estimates $\hat{I}^\lambda_{\mathcal{D}_\text{OPT}}(T;Y)$ and $\hat{I}^\lambda_{\mathcal{D}_\text{OPT}}(X;T)$ to obtain the ordered subset $\Lambda_{\text{OPT}}$. ③ FST, a sequential FWER-controlling MHT algorithm, is applied to the subset $\Lambda_{\text{OPT}}$ to form the subset $\Lambda_{\text{MHT}}\subseteq\Lambda_{\text{OPT}}$ of hyperparameters $\lambda \in \Lambda_{\text{MHT}}$ that are likely to satisfy the constraint (\ref{['eq:relaxed_constraint']}). Finally, the hyperparameter $\lambda^*$ is chosen as the vector in $\Lambda_{\text{MHT}}$ that minimizes the estimate $\hat{I}^\lambda_{\mathcal{D}_\text{MHT}}(X;T)$.
• Figure 3: Illustration of the operation of IB-MHT for the experiment in Section \ref{['sec:simulations']}: (a) Estimated Pareto front using the estimated mutual informations $\hat{I}^\lambda_{\mathcal{D}_\text{OPT}}(T;Y)$ and $\hat{I}^\lambda_{\mathcal{D}_\text{OPT}}(X;T)$; (b) Sequential MHT using the estimated mutual informations $\hat{I}^\lambda_{\mathcal{D}_\text{MHT}}(T;Y)$ and $\hat{I}^\lambda_{\mathcal{D}_\text{MHT}}(X;T)$.
• Figure 4: Joint distributions of the mutual informations $I^{\lambda^*}(T;Y)$ and $I^{\lambda^*}(X;T)$ obtained by using a conventional IB solver (Section \ref{['sec:conventional_IB']}) and IB-MHT for the classical IB problem (\ref{['eq:classic_IB']}) using 50 trials of Algorithm \ref{['alg:Pareto']}. The outage probability for conventional IB and IB-MHT are reported to be 0.27 and 0.06, respectively.
• Figure 5: Joint distributions of the mutual informations $I^{\lambda^*}(T;Y)$ and $I^{\lambda^*}(X;T)$ obtained by using a conventional IB solver (Section \ref{['sec:conventional_IB']}) and IB-MHT for the deterministic IB problem (\ref{['eq:deterministic_IB2']}) using 50 trials of Algorithm \ref{['alg:Pareto']}. The outage probability for conventional IB and IB-MHT are reported to be 0.26 and near zero, respectively.

Theorems & Definitions (4)

• Lemma 1: stefani2014confidence
• Proposition 1
• proof
• Proposition 2: laufer2022efficiently