Exploring the Trade-Off in the Variational Information Bottleneck for Regression with a Single Training Run

Abstract

An information bottleneck (IB) enables the acquisition of useful representations from data by retaining necessary information while reducing unnecessary information. In its objective function, the Lagrange multiplier $\beta$ controls the trade-off between retention and reduction. This study analyzes the Variational Information Bottleneck (VIB), a standard IB method in deep learning, in the settings of regression problems and derives its optimal solution. Based on this analysis, we propose a framework for regression problems that can obtain the optimal solution of the VIB for all $\beta$ values with a single training run. This is in contrast to conventional methods that require one training run for each $\beta$. The optimization performance of this framework is theoretically discussed and experimentally demonstrated. Our approach not only enhances the efficiency of exploring $\beta$ in regression problems but also deepens the understanding of the IB’s behavior and its effects in this setting.

1. Introduction

1.1. Information Bottleneck

Information extraction refers to obtaining a new representation by retaining the necessary information from a given input while reducing unnecessary information. An information bottleneck (IB) [1] formalizes this process from an information-theoretic perspective. Consider a source random variable X and a target random variable Y, with their joint distribution assumed to be known. The objective of the IB is to derive a random variable Z from X that compresses the information contained in X while retaining as much information as possible about Y. Formally, with mutual information $I(·;·)$, this can be expressed as follows:

$$\underset{Z\in \Delta }{max}I(Z;Y)s.t.I(X;Z)\le r$$

(1)

Here, $\Delta$ represents the set of all random variables Z that satisfy the Markov chain $Y\leftrightarrow X\leftrightarrow Z$. We call this the IB objective. In many practical cases, to avoid constrained optimization, instead of directly optimizing the IB objective, the following objective is used. This is the Lagrangian relaxation [2] of the IB objective and is called the IB Lagrangian [3].

$${\mathcal{L}}_{IB}(Z;\beta )=I(Z;Y)-\beta I(X;Z)$$

(2)

$\beta >0$ is called the Lagrange multiplier. The advantage of the IB method lies in its ability to control the trade-off between compression $I(X;Z)$ and prediction $I(Z;Y)$ through $\beta$. When $\beta$ is small, the resulting representation tends to be more predictive, while a larger $\beta$ leads to a more concise representation.

1.2. Methods of IB

Mutual information often involves integrals, which are difficult to calculate, making it challenging to solve the IB Lagrangian in general situations. Consequently, early IB methods were proposed for limited scenarios, such as when X and Y take on a relatively small number of discrete values [1] or when they follow a joint Gaussian distribution [4]. Later, the Variational Information Bottleneck (VIB) [5] introduced a method to solve the IB Lagrangian in more general situations using variational approximations and deep neural networks (DNNs). This approach has become a standard method of solving IBs in DNN-based applications. The details of the VIB are introduced in Section 3. While this paper focuses on analyzing the VIB, it is important to note that there are other methods for IBs [6,7,8,9,10], as well as other objective functions inspired by IBs [11,12,13,14,15].

1.3. Effects and Applications of IB in DNNs

There has been extensive research on the use of IBs in supervised learning. As we will see later, maximizing $I(Z;Y)$ is related to likelihood maximization. In addition to this likelihood maximization, an IB can remove unnecessary information, which can sometimes be detrimental to robustness, through compression. This can have a beneficial effect on deep learning models, which can occasionally become overly flexible. Theoretically, ref. [16] has explored IB-based statistical learning theory using DNNs, and [17] has provided a generalization gap between true mutual information and empirical mutual information in probabilistic models. In particular, the IB framework has shown many experimental benefits in classification problems with DNNs. These include improvements in generalization performance [5,8], invariance to nuisance factors [18], robustness against adversarial attacks [5,8,9], out-of-distribution detection [9,19], domain generalization [20,21], and calibration [19]. There are some works that have applied IBs in regression tasks [6,10,22,23,24]. Ngampruetikorn and Schwab [22] studied the quantification of overfitting in the context of IBs specifically in regression tasks with a linear data generation process. Recently, ref. [10] proposed a superior IB method for regression tasks using Cauchy–Schwarz divergence. While the specific effects of compression in regression problems have not been explored as extensively as in classification, this area holds promise for future research. In addition to supervised learning, the IB has also been used in unsupervised learning [25,26] and as a tool to understand DNNs [27,28,29].

1.4. Efficient $\beta$ Exploration

In general, the optimal value of $\beta$ to be used during training is not known beforehand. This is because the optimal trade-off between prediction and compression, which maximizes generalization, depends on the task-specific distribution [17]. Therefore, practitioners usually follow a procedure whereby they train multiple models with varying $\beta$ values and select the model with the best properties. However, the drawback of this approach is that it requires multiple training runs and is computationally expensive.

To make this search more efficient, several studies have been conducted. Ref. [30] showed, that even within the range $\beta \in (0,1)$, inappropriate values of $\beta$ can render the IB Lagrangian unlearnable (i.e., the optimal Z becomes independent of X). They theoretically established sufficient conditions for $\beta$ to be learnable and proposed an algorithm to estimate the range of $\beta$ values that are likely to be learnable. This research is valuable in narrowing down the range of $\beta$ values to be explored from the perspective of the search efficiency. However, finding a useful $\beta$ still requires multiple training runs. Ref. [31] developed a method to achieve the desired compression rate r in a single training run by establishing a one-to-one mapping between $\beta$ and r. However, the challenge of identifying the appropriate compression rate ultimately still requires trial and error. Additionally, ref. [9] used supervised disentangling to obtain a maximally compressed representation Z from the training data without reducing $I(Z;Y)$. Although this training process does not depend on $\beta$, it is important to note that this representation is not necessarily the most useful. This is because our goal is often to maximize the true $I(Z;Y)$, not the empirical $I(Z;Y)$ observed in the training data. In fact, similar to the bias–variance trade-off, the true $I(Z;Y)$ is determined by the trade-off between the empirical $I(Z;Y)$ and $I(X;Z)$ in the training data [17]. Previously, we proposed a framework called the Flexible Variational Information Bottleneck (FVIB) [32] that enables the VIB objective to be learned for all values of $\beta$ in a single training run for classification tasks. The FVIB allows models with all $\beta$ values to be obtained without requiring additional parameters or training time, as compared to a single training process for the VIB, thereby enabling efficient $\beta$ exploration.

1.5. Contributions and the Structure of This Paper

In this paper, we propose FVIB-Regression (FVIB-R) for regression tasks, which allows the VIB objective to be learned for all $\beta$ values in a single training run. First, in Section 3, we derive the closed-form solution for the VIB in regression tasks. This provides insights into the behavior of the VIB in regression problems and is expected to be useful for future theoretical analysis and improvements to the IB or VIB in this setting. Next, in Section 4, we use this analysis to design FVIB-R and theoretically discuss its optimization performance. Finally, in Section 5.1.1, we experimentally demonstrate the optimization performance of FVIB-R.

2. Related Work

2.1. Theory of IB

This study analytically derives the optimal solution of the VIB, and its findings have the potential to contribute to the theoretical understanding of both the VIB and IB in the future. Therefore, we discuss the relationship between this work and previous theoretical studies on IBs. Ref. [33] discovered that, in IBs, when $\beta$ is varied, there is a point at which a significant change in the prediction accuracy occurs, which they referred to as the IB phase transition. They also demonstrated that new correlations are learned during this phase transition. Our study elucidates the behavior of the VIB optimal solution concerning $\beta$, which could potentially connect with the theoretical research on the IB phase transition. Kolchinsky et al. [34] studied IBs in the settings of stochastic models and deterministic scenarios. While our research is constrained to regression problems, it does not rely on the deterministic scenario constraint, meaning that it also addresses datasets with multiple different labels for the same input data. Amjad and Geiger [35] theoretically demonstrated several issues when training classification problems using IBs with deterministic DNNs, and they pointed out that stochastic DNNs such as the VIB can avoid these problems. In contrast, our study focuses on the VIB in the context of regression problems. There are many works that study the relationship between information-theoretic quantities and generalization [16,36,37,38,39,40,41,42,43]. In particular, ref. [16] explored statistical learning theory using mutual information in the context of deterministic deep neural networks. Ref. [17] provided a generalization gap between the true mutual information and the empirical mutual information in stochastic models. This research confirms that the VIB is effective for generalization as a stochastic model.

2.2. Variational Autoencoders

Variational autoencoders (VAE) [44] and $\beta$-VAE [45] have loss functions that include both distortion (prediction) and rate (compression) terms, and they can be interpreted as special cases of the VIB [5,46,47]. Therefore, our analysis is related to the existing analyses of VAE. There have been many studies conducted to understand the characteristics of the optimal solutions in $\beta$-VAE [48,49,50]. Our analysis revisits these studies in the context of the VIB, which is a supervised learning framework. By doing so, our analysis extends the insights gained from VAE to the domain of IBs and has the potential to provide a profound understanding of how the interplay between distortion and rate terms can be leveraged in supervised learning scenarios.

3. Analysis of VIB in Regression Tasks

In this section, we analyze the VIB in the settings of regression tasks. We begin by explaining the VIB, followed by the introduction of a specific model setup for our analysis. Finally, we derive the optimal solution for the VIB with this model.

3.1. Variational Information Bottleneck

The VIB enables the learning of IBs in general settings by providing a lower bound on Equation (2) through variational approximation. When predicting Y from X, the random variable Z is obtained via the feature extractor $p_{\theta }\left(Z\right|X)$. First, let us consider the prediction term, i.e., the first term in Equation (2). By using a new model $q_{\varphi }\left(Y\right|Z)$ as a variational approximation of $p_{\theta }\left(Y\right|Z)$, the following variational lower bound is obtained:

$$\begin{array}{cc}\hfill I(Z;Y)\ge \int dxdydzp(x,y)p_{\theta }\left(z\right|x)& log{q}_{\varphi }\left(y\right|z)+H\left(Y\right)\hfill \end{array}$$

(3)

Here, $H\left(Y\right)$ represents the entropy of Y (or the differential entropy in the case of a regression problem). Estimating $H\left(Y\right)$ from observed data is not a straightforward task when Y is a continuous variable. However, since this value remains constant throughout the learning process, it can be ignored during training [51]. Next, regarding the compression term $I(X;Z)$ in Equation (2), by using $r\left(Z\right)$ as a variational approximation of $p_{\theta }\left(Z\right)$, an upper bound can be obtained:

$$I(X;Z)\le \int dxp\left(x\right)D_{KL}[p_{\theta }\left(Z\right|x)\left|\right|r\left(Z\right)]$$

(4)

In practice, a fixed distribution is used for $r\left(Z\right)$. By combining these elements, a lower bound on the IB Lagrangian can be obtained. Given data {$(x_1,y_1),\cdots ,(x_N,y_N)$}, and by using the empirical distribution as the joint distribution of X and Y, the objective function of the VIB is derived:

$$\begin{array}{cc}\hfill {\mathcal{L}}_{VIB}(\theta ,\varphi ;\beta ):=& {\displaystyle \frac{1}{N}}\sum _{i=1}^N\left[{\mathbb{E}}_{p_{\theta }\left(Z\right|x_i)}[log{q}_{\varphi }\left(y_i\right|Z\left)\right]-\beta D_{KL}[p_{\theta }\left(Z\right|x_i)\left|\right|r\left(Z\right)]\right]\hfill \end{array}$$

(5)

The first term, which is the prediction term, corresponds to the expected value of the log-likelihood, revealing the relationship between maximizing $I(Z;Y)$ and maximum likelihood estimation.

3.2. Model Settings for Analysis

Below, the model setup for the analysis is introduced. All subsequent lemmas and theorems are based on the following setup. Consider a d-dimensional regression task, using input data $x\in \mathcal{X}$ and the continuous label $\mathit{y}\in {\mathbb{R}}^d$. The random variable Z is defined as $\mathit{z}\in {\mathbb{R}}^{\kappa }$ with $\kappa \ge d$. We define $p_{\theta }\left(Z\right|x)$ as follows:

$$p_{\theta }\left(Z\right|x)=\mathcal{N}(\mu \left(x\right),\Sigma \left(x\right))$$

(6)

Here, let $\mu :\mathcal{X}\to {\mathbb{R}}^{\kappa }$ and $\Sigma :\mathcal{X}\to \{A\in {\mathbb{R}}^{\kappa \times \kappa }|A\succ 0\}$, where $\Sigma$ represents the covariance matrix. Although the covariance matrix is typically restricted to a diagonal matrix, our analysis does not impose this restriction. Based on these settings, the parameters are defined as $\theta =\{\mu ,\Sigma \}$. The distribution $r\left(Z\right)$ is typically set as $r\left(Z\right)=\mathcal{N}(\mathbf{0},I)$, and we adopt this in our study as well.

The predictor $q_{\varphi }$ is often composed of a single fully connected layer in practice. With this in mind, in our analysis, we define the fully connected layer as $v\left(\mathit{z}\right):=W\mathit{z}+\mathit{b}$, where $W\in {\mathbb{R}}^{d\times \kappa }$ and $\mathit{b}\in {\mathbb{R}}^d$. Using this, $q_{\varphi }$ is defined as follows:

$$q_{\varphi }\left(\mathit{y}\right|\mathit{z})=\mathcal{N}(v\left(\mathit{z}\right),{\displaystyle \frac{1}{2}}I)$$

(7)

This setting corresponds to the mean squared error (MSE):

$$log{q}_{\varphi }\left(\mathit{y}\right|\mathit{z})=-||{v\left(\mathit{z}\right)-\mathit{y}\left|\right|}_2^2-{\displaystyle \frac{d}{2}}log\pi$$

(8)

Therefore, $\varphi =\{W,\mathit{b}\}$. It should be noted that the second term is constant during training.

3.3. The Optimal Solution

Under the above settings, we analyze the VIB objective. To this end, we introduce several new definitions. We split the data indices into equivalence classes defined by

$$\left[i\right]:=\{j\in \left\{1,2\cdots ,N\}\right|x_j=x_i\}$$

(9)

This means that we are grouping data indices with the same input into a single set. Using this, for each label, we reassign the average of the labels whose indices belong to the same equivalence class:

$${\tilde{\mathit{y}}}_i:={\displaystyle \frac{1}{\left|\right[i\left]\right|}}{\Sigma }_{j\in \left[i\right]}{\mathit{y}}_j$$

(10)

We define the mean and the covariance matrix for ${\tilde{\mathit{y}}}_i$:

$${\mathit{m}}_{\tilde{\mathit{y}}}:={\displaystyle \frac{1}{N}}\sum _{i=1}^N{\tilde{\mathit{y}}}_i$$

(11)

$$S_{\tilde{\mathit{y}}}:={\displaystyle \frac{1}{N}}\sum _{i=1}^N({\tilde{\mathit{y}}}_i-{\mathit{m}}_{\tilde{\mathit{y}}})({{\tilde{\mathit{y}}}_i-{\mathit{m}}_{\tilde{\mathit{y}}})}^{\top }$$

(12)

Additionally, in the following, $\mathrm{diag}\left(\mathit{a}\right)$ denotes a diagonal matrix with the vector $\mathit{a}$ as its diagonal elements, and ${\left[a_i\right]}_i$ represents a vector with its i-th component being $a_i$. Using these definitions, we derive the optimal solution for the VIB objective in the context of regression tasks.

Lemma 1.

Consider the model settings in Section 3.2. For any $\beta >0$, the feature dimension $\kappa \ge d$ does not affect the value of ${max}_{\theta ,\varphi }{\mathcal{L}}_{VIB}(\theta ,\varphi ;\beta )$. Below, we consider the case where $\kappa =d$. For any $\beta >0$, if the parameters satisfy the following conditions, ${\mathcal{L}}_{VIB}(\theta ,\varphi ;\beta )$ is maximized.

• $\mu \left(x_i\right)={\displaystyle \frac{2}{\beta }}({I+{\displaystyle \frac{2}{\beta }}{W}^{\top }W)}^{-1}{W}^{\top }({\tilde{\mathit{y}}}_i-{\mathit{m}}_{\tilde{\mathit{y}}})$ for all $i=1,2\cdots ,N$.
• $\Sigma (x_i)=({I+{\displaystyle \frac{2}{\beta }}{W}^{\top }W)}^{-1}$ for all $i=1,2\cdots ,N$.
• $W=\mathit{Pdiag}\left([{\pm \sqrt{max(0,{\lambda }_i-{\displaystyle \frac{\beta }{2}})}]}_i\right)R$, where $S_{\tilde{\mathit{y}}}=:\mathit{Pdiag}([{{\lambda }_i]}_i)P^{\top }$ by orthogonal diagonalization and R is an arbitrary orthogonal matrix.
• $\mathit{b}={\mathit{m}}_{\tilde{\mathit{y}}}$.

Refer to Appendix A.1 for the proof. Below, the properties of the optimal solution of the VIB in regression tasks are considered based on Lemma 1. We discuss how well the optimal solution of the VIB optimizes the IB objective, as given by Equation (1). It is important to note that the joint distribution of X and Y considered here is not the true distribution but the empirical distribution of the training data. The IB framework assumes a Markov chain $Y\leftrightarrow X\leftrightarrow Z$, and, according to the data processing inequality, we have $I(Z;Y)\le I(X;Z)$. This means that when the IB objective maximizes $I(Z;Y)$ with the limitation of a given $I(X;Z)$, the feasible region is constrained by $I(Z;Y)\le I(X;Z)$. Furthermore, the equality holds if and only if $I(X;Z|Y)=0$. Now, consider Lemma 1. We see that the representation Z in the optimal solution becomes determined solely by Y and is conditionally independent of X given Y, when X and Y lie within the support of the empirical distribution. This implies $I(X;Z|Y)=0$, and thus $I(Z;Y)=I(X;Z)$ is achieved. This indicates that the set of optimal solutions for the VIB in the regression tasks includes, at the very least, the optimal solutions of the IB objective under the empirical distribution. This is particularly interesting when considering that the VIB is designed through Lagrangian relaxation and the variational approximation of the original IB objective.

Furthermore, as we show in Section 4, FVIB-R achieves the condition of Lemma 1 through training, thereby reaching the optimal solutions not only of the VIB objective but also of the IB objective.

4. Methods

4.1. FVIB-R

Based on the above analysis, we design a framework called FVIB-R, which allows the optimization of the VIB objective for all values of $\beta$ in a single training run for regression problems. In this section, we first introduce the model structure and objective function of FVIB-R. Then, we theoretically discuss FVIB-R’s ability to maximize the VIB objective.

For FVIB-R, we train a model $h_{\psi }:\mathcal{X}\to {\mathbb{R}}^d$ with parameters $\psi \in \Psi$ by maximizing the following objective function:

$$J_{FVIB-R}\left(\psi \right):=-{\displaystyle \frac{1}{N}}\sum _{i=1}^N||{h_{\psi }(x_i)-({\tilde{\mathit{y}}}_i-{\mathit{m}}_{\tilde{\mathit{y}}})||}_2^2$$

(13)

Note that, unlike the VIB objective, this objective function does not depend on $\beta$. Next, using the trained $h_{\psi }$, for any $\beta >0$, we set the VIB parameters as follows:

• ${\tilde{\mu }}_{\beta ,\psi }\left(x\right):={\displaystyle \frac{2}{\beta }}({I+{\displaystyle \frac{2}{\beta }}{W}^{\top }W)}^{-1}{W}^{\top }{h}_{\psi }\left(x\right)$;
• ${\tilde{\Sigma }}_{\beta }\left(x\right):=({I+{\displaystyle \frac{2}{\beta }}{W}^{\top }W)}^{-1}$;
• ${\tilde{W}}_{\beta }:=P\mathrm{diag}({\left[\sqrt{max(0,{\lambda }_i-{\displaystyle \frac{\beta }{2}})}\right]}_i)$, where $S_{\tilde{\mathit{y}}}=:P\mathrm{diag}\left({\left[{\lambda }_i\right]}_i\right)P^{\top }$ by orthogonal diagonalization;
• $\tilde{\mathit{b}}:={\mathit{m}}_{\tilde{\mathit{y}}}$.

As a result, the value of $\beta$ does not influence the training process and can be adjusted during the evaluation phase. This is in contrast to the VIB framework, where $\beta$ affects the training setup and cannot be changed afterward.

In the following, we theoretically argue that, despite the structural advantages mentioned above, this setup and objective function can still learn the VIB objective.

Theorem 1.

Consider the model settings in Section 3.2. If ${lim}_{t\to \infty }{J}_{FVIB-R}\left({\psi }_t\right)=0$, then, for $\beta >0$, the sequence $\{{{\mathcal{L}}_{VIB}({\tilde{\mu }}_{\beta ,{\psi }_t},{\tilde{\Sigma }}_{\beta },{\tilde{W}}_{\beta },\tilde{\mathit{b}};\beta \left)\right\}}_{t\in \mathbb{N}}$ converges to ${max}_{\theta ,\varphi }{\mathcal{L}}_{VIB}(\theta ,\varphi ;\beta )$ as $t\to \infty$.

This is derived from the fact that FVIB-R with $J_{FVIB-R}\left(\psi \right)=0$ is the optimal solution to the VIB objective, and from the continuity of the VIB objective. Detailed proof is provided in Appendix A.2. This theorem shows that, for any $\beta$, when $J_{FVIB-R}\left(\psi \right)$ is sufficiently maximized, the VIB objective approaches its maximum. This characteristic indicates that FVIB-R can learn the VIB for any $\beta$, despite being a learning setup that is independent of $\beta$. Note that Theorem 1 is valid when $J_{FVIB-R}\left(\psi \right)$ converges to zero. However, in cases where gradient descent reaches a local optimum, or when the training is stopped early to prevent overfitting, $J_{FVIB-R}\left(\psi \right)$ does not strictly converge to zero. Even in such cases, the following properties are beneficial regardless of the value that $J_{FVIB-R}\left(\psi \right)$ reaches.

Theorem 2.

Consider the model settings in Section 3.2 and assume

$$\{[h_{\psi }\left({\mathit{x}}_1\right)h_{\psi }\left({\mathit{x}}_2\right)\cdots h_{\psi }\left({\mathit{x}}_N\right)]|\psi \in \Psi \}={\mathbb{R}}^{d\times N}.$$

(14)

Then, for any $\beta >0$,

$$\underset{\psi :J_{FVIB-R}\left(\psi \right)=\alpha }{min}{\mathcal{L}}_{VIB}({\tilde{\mu }}_{\beta ,\psi },{\tilde{\Sigma }}_{\beta },{\tilde{W}}_{\beta },\tilde{\mathit{b}};\beta )$$

(15)

is monotonically increasing with respect to α.

Detailed proof can be found in Appendix A.3. The assumption of Equation (14) is justified by the flexibility of neural networks, as shown in the universal approximation theorem [52]. Theorem 2 demonstrates that optimizing $J_{FVIB-R}\left(\psi \right)$ monotonically increases the worst-case value of the VIB objective for any $\beta$. In other words, this is a process similar to an increase in a lower bound of the objective function, which is commonly used in practical machine learning applications. It should be noted that this theorem does not strictly guarantee an increase in the VIB objective. However, as seen later in Figure 1 and Figure 2, in the experiments, the VIB objective function does indeed increase. These two theorems together show that the proposed setup and objective function are effective in simultaneously maximizing the VIB objective for all values of $\beta$.

4.2. Relation to FVIB

Previously, we have proposed the FVIB [32], which enables the learning of the VIB for all values of $\beta$ in a single training run for classification problems. In this study, we address a similar challenge in the context of regression problems. Here, we will summarize the relationship between FVIB-R and FVIB. Regarding the model setup during the analysis, while the feature extractor $p_{\theta }$ remains the same, the predictor $q_{\varphi }$ differs to reflect the problem setting. Specifically, in FVIB-R, $log{q}_{\varphi }$ is expressed as a quadratic function, allowing the expectation in the prediction term of the VIB objective to be computed analytically. On the other hand, in FVIB, the expectation in the prediction term cannot be solved analytically, and, instead, the Taylor approximation of $log{q}_{\varphi }$ is used. As a result, FVIB-R directly analyzes the VIB objective, whereas FVIB analyzes an approximation of the VIB objective. Additionally, FVIB assumes a deterministic scenario, where no data points exist with the same X but different Y values [34]. In contrast, FVIB-R does not require this assumption. This is because, as seen in Equation (A1) in Appendix A.1, any non-deterministic situation in regression problems can be transformed into a deterministic situation without changing the VIB objective by replacing ${\mathit{y}}_i$ with ${\tilde{\mathit{y}}}_i$. Overall, for regression tasks, it is possible to analyze the VIB with fewer assumptions and without relying on approximations.

5. Experiments

5.1. Experimental Setup

We demonstrate the optimization performance of FVIB-R using two real datasets and further investigate its compression effects with a synthetic dataset. Below, we introduce the datasets and the experimental procedure for each setting. The detailed configurations are presented in

Table 1. Experimental setup. In the description of the feature extractor, the subscripts of each module indicate the output dimension.

	California Housing Prices	Nutrition5k	Synthetic Dataset
Feature extractor	$x\in {\mathbb{R}}^8$	$x\in {\mathbb{R}}^{3\times 480\times 640}$	$x\in \mathbb{R}$
	$\to Linea{r}_{128}$	$\to ResNet{18}_{\kappa or2\kappa }$ [53]	$\to Linea{r}_{128}$
	$\to ReLU\to Linea{r}_{128}$		$\to ReLU\to Linea{r}_{128}$
	$\to ReLU\to Linea{r}_{\kappa or2\kappa }$		$\to ReLU\to Linea{r}_{128}$
			$\to ReLU\to Linea{r}_{128}$
			$\to ReLU\to Linea{r}_1$
$\kappa$ in non-FVIB-R models	32	256	−
# Epochs	{50, 100, 150, 200}	{50, 100, 150, 200}	200
Learning rate	{$1.0\times {10}^{-4}$, $1.0\times {10}^{-3}$}	{$1.0\times {10}^{-4}$, $1.0\times {10}^{-3}$}	$1.0\times {10}^{-3}$
Optimizer	Adam [54]	Adam	Adam

.

5.1.1. Real Dataset

The first dataset is the California housing prices dataset [55], which consists of 20,640 samples. This dataset, based on the 1990 California census, involves a regression problem where the goal is to predict a single variable representing house prices using eight source variables, including the latitude, longitude, and the number of rooms. We use the version of the dataset distributed in the scikit-learn package [56]. The second dataset is the Nutrition5k dataset [57]. Nutrition5k contains around 5000 images of food. The task is to predict five continuous variables that indicate calories, mass, and three major nutrients, using these image data as input. The data acquisition and normalization procedures follow [58]. For Nutrition5k, we use parameters pretrained on ImageNet1k [59] as the parameters’ initial values to compensate for the limited data.

The following is an overview of the experimental procedure. We split the data into training, validation, and test sets in a ratio of $80\%$, $10\%$, and $10\%$, respectively. First, for each model, we conduct a grid search over the candidate learning settings, such as the epochs and learning rates, as shown in Table 1. During this step, $\beta$ is fixed at $0.01$ for the California housing prices dataset and $0.05$ for the Nutrition5k dataset, and the evaluation metric is the MSE on the validation set. Next, using the best training settings obtained from the grid search, we retrain the models on the training data with $\beta \in \{{10}^{-6},\cdots ,{10}^{-2},0.1,0.2,\cdots ,1.0\}$ for the California housing prices dataset and $\beta \in \{5\times {10}^{-6},\cdots ,5\times {10}^{-2},0.5,1.0,1.5\cdots ,5.0\}$ for the Nutrition5k dataset. Finally, we evaluate these trained models on the test data to assess their performance. All experiments are conducted using PyTorch [60]. All models use a predictor composed of an affine transformation on top of the feature extractor described in Table 1.

5.1.2. Synthetic Dataset

As a synthetic dataset, we sample x and y from the data generation process represented below:

$$x\sim \mathcal{N}\left(0,1\right)$$

(16)

$$y=sin5x+\epsilon$$

(17)

where $\epsilon \sim \mathcal{N}\left(0,1\right)$. For both training and evaluation, 100 samples are used. In experiments that are repeated multiple times, different data samples are used for each trial. To examine the effects of noise, overfitting, and the behavior of compression in response to these factors, a relatively large model is employed for the small dataset. The other settings are shown in Table 1.

5.2. Results and Discussion

First, we verify the ability of FVIB-R to optimize the VIB objective. To do this, we compare the VIB objective values across the training epochs between FVIB-R and VIB. Figure 1 and Figure 2 show the VIB objective values for each training epoch on the California housing prices dataset and the Nutrition5k dataset, respectively. It is important to note that the VIB objective values presented here are computed using the training data. While the VIB is trained separately for each value of $\beta$, FVIB-R training is conducted only once. In both datasets, the VIB objective values for FVIB-R increase monotonically and converge to values similar or superior to those obtained by directly training the VIB for each $\beta$. These results align with the theoretical properties of FVIB-R discussed in Section 4.The reason that FVIB-R outperforms the VIB in optimization for $\beta =0.5$ and $2.5$ in Figure 2 can be attributed to the fact that the VIB restricts the covariance matrix $\Sigma$ of Z to be diagonal due to computational efficiency constraints, whereas FVIB-R does not impose such a restriction (see Section 3.2). This makes the model more flexible. Additionally, FVIB-R tends to converge faster compared to direct VIB training. This faster convergence might be attributed to the different handling of the expectation of the prediction term in the VIB objective between the VIB and FVIB-R. In the VIB, the expected value is computed based on the sampled Z, which introduces fluctuations in the gradient direction. In contrast, FVIB-R analytically computes the expected value and directly learns the optimal solution, avoiding the aforementioned fluctuations in the gradient direction. This difference in the training methods could be the reason for the observed faster convergence in FVIB-R.

Next, we examine FVIB-R’s optimization performance from a different perspective. Figure 3, Figure 4, Figure 5 and Figure 6 illustrate the trade-off between the compression and prediction terms of the VIB objective on the California housing prices dataset and the Nutrition5k dataset, respectively. More precisely, prediction is minus the MSE and compression is the sample mean of the KL divergence in Equation (5). Each point in these plots corresponds to a different value of $\beta$. Both the training and test datasets are shown in these figures. Once again, it is important to note that, while the VIB requires multiple training processes for different values of $\beta$, FVIB-R is trained only once per dataset. The curve for FVIB-R is shown as a plot due to the evaluation, but the model is able to continuously vary along this curve. In Figure 3 and Figure 4, FVIB-R is compared to the VIB and squared-VIB (sq-VIB) in the California housing prices dataset. The squared versions, proposed by [34], use a squared value for the compression term in the objective function for the purpose of exploring the entire trade-off. In these figures, all methods are evaluated and plotted for $\beta \in \{{10}^{-6},\cdots ,{10}^{-2},0.1,0.2,\cdots ,1.0\}$. In Figure 3, we observe that FVIB-R achieves compression–prediction performance on the test data that is comparable to that of the VIB. On the training data, FVIB-R shows better predictions at similar levels of compression compared to the VIB. This difference is likely not due to the framework’s inherent performance, but rather a result of the grid search process selecting models that were either underfitting or overfitting. Moreover, Figure 4 shows that FVIB-R is comparable to sq-VIB both in the training set and in the test set.

Figure 5 and Figure 6 show the optimization performance of FVIB-R compared to the VIB or sq-VIB in the Nutrition5k dataset. In these figures, the VIB and sq-VIB are evaluated for $\beta \in \{5\times {10}^{-6},\cdots ,5\times {10}^{-2},0.5,1.0,1.5\cdots ,5.0\}$. In Figure 5, FVIB-R demonstrates performance comparable to the VIB’s in both the training and test data. Figure 6 shows that the highest value of prediction for FVIB-R is similar to that of sq-VIB, while FVIB-R is more predictive than sq-VIB when the representation is compressed to some extent. In conclusion, despite the more efficient training process of FVIB-R, its performance is comparable (or sometimes superior) to that of the traditional VIBs on both the training and test data.

Finally, we present the quantitative evaluation of the performance. Here, we show the MSE as it reflects $I(Z;Y)$, as shown in Equations (3) and (8). In addition to the VIB, we compare FVIB-R with the Nonlinear Information Bottleneck (NIB) [6], as well as the squared-VIB (sq-VIB) and squared-NIB (sq-NIB).

Table 2. Mean squared error and the number of parameters of FVIB-R compared to other IB methods.

	California Housing Prices		Nutrition5k
Method	MSE	# Parameters	MSE	# Parameters
FVIB-R	0.1910	$1.8\times {10}^4\times 1$	0.0881	$1.1\times {10}^7\times 1$
VIB	0.1979	$2.6\times {10}^4\times 15$	0.0805	$1.1\times {10}^7\times 15$
sq-VIB	0.1786	$2.6\times {10}^4\times 15$	0.0893	$1.1\times {10}^7\times 15$
NIB	0.1816	$2.2\times {10}^4\times 15$	0.1549	$1.1\times {10}^7\times 15$
sq-NIB	0.1790	$2.2\times {10}^4\times 15$	0.1431	$1.1\times {10}^7\times 15$

shows the MSE on the test data and the number of parameters required for each IB method. For each method, the best model is selected based on the MSE in the validation set across $\beta \in \{{10}^{-6},\cdots ,{10}^{-2},0.1,0.2,\cdots ,1.0\}$ for the California housing prices dataset and $\beta \in \{5\times {10}^{-6},\cdots ,5\times {10}^{-2},0.5,1.0,1.5\cdots ,5.0\}$ for the Nutrition5k dataset and then evaluated on the test set. In the parameter count notation, the last multiplication indicates the number of models required for the $\beta$ search. When comparing FVIB-R to the VIB, FVIB-R outperforms the VIB on the California housing prices dataset, while the VIB outperforms FVIB-R on the Nutrition5k dataset. Similarly, when compared to other IB methods, FVIB-R outperforms them on one dataset, while the other methods perform better on the other dataset. Overall, despite reducing the number of parameters, FVIB-R demonstrates competitive performance with the other IB methods. This performance highlights that FVIB-R can achieve comparable results with fewer parameters, making it an efficient alternative to traditional IB methods.

With the synthetic data, we investigate the distributional difference between the true $p\left(y\right|x)$ and its prediction changes for each $\beta$. For this purpose, the expected KL divergence, represented below, is utilized.

$${\mathbb{E}}_x[D_{KL}(p\left(y\right|x)\left|\right|p\left(\widehat{y}\right|x)\left)\right]$$

(18)

where $\widehat{y}$ is the predicted random variable obtained from $q_{\varphi }\left(y\right|z)$. For the estimation, the sample mean of the analytically derived KL divergence is calculated. The results are shown in Figure 7. It is observed that the predicted distribution moves closer to the true distribution with suitable compression.

The distributions of the true $p\left(y\right|x)$ and the predicted $p\left(\widehat{y}\right|x)$ for each $\beta$ are shown in Figure 8. These results demonstrate that compression leads to a smoother mean line and wider variance. When $\beta =1.0$, the predicted distribution is the closest match to the true distribution, which is consistent with the observations in Figure 7. The detailed behavior of compression and its limitations are discussed in the Limitations section.

6. Limitations and Future Work

Lemma 1 highlights the potential shortcomings of the VIB in regression problems, while also showing that FVIB-R inherits these shortcomings. The key points are summarized below. The first limitation lies in the simplicity of the compression procedure. Lemma 1 demonstrates the optimal solution of the VIB for regression problems. In particular, $\mu \left(x\right)$ can be viewed as a combination of $\left({\tilde{\mathit{y}}}_i-{\mathit{m}}_{\tilde{\mathit{y}}}\right)$ and a linear transformation controlled by $\beta$. Here, the former is deterministically determined by the data, and thus includes noise, while the latter performs compression through its transformation. However, since this is a simple linear transformation, the noise that can be eliminated is limited. FVIB-R, which is designed based on this, shares the same limitation. Future research could explore comparisons with other IB methods in this regard and investigate more flexible compression approaches. Another limitation lies in the certain divergence from the properties of the IB. It is theoretically known that the IB undergoes phase transitions, and, to the best of our knowledge, the phase transitions of the VIB have been experimentally confirmed in classification problems [33]. However, since Lemma 1 represents the solution of the VIB in regression problems, it suggests that phase transitions may not occur in such cases. This is because the VIB (and FVIB-R) aims to increase the IB’s lower bound for the empirical distribution, and there is no guarantee that they will reach the IB optimal solution for the true distribution. Improvements to the VIB and FVIB-R in this regard will be left for future work.

7. Conclusions

In this study, we explore the VIB in the context of regression problems and analytically derive its optimal solution. We anticipate that this analysis will contribute to the theoretical understanding and improvement of both the VIB and the broader IB framework in the future. Using this analysis, we propose a framework that allows the VIB objective to be learned for all values of $\beta$ in a single training run for regression problems. Additionally, this framework explicitly demonstrates the compression method, enabling a more intuitive understanding. Finally, we theoretically justify the optimization performance of this framework and validate its properties through experiments. This method not only offers direct practical applications but also serves as a valuable tool for an understanding of the behavior and impact of IBs in regression problems. We believe that this framework will facilitate further research and practical advancements in the field of IBs.