Neural Networks Perform Sufficient Dimension Reduction

Neural networks have achieved significant success in a tremendous variety of applications (Lee and Abu-El-Haija 2017; Silver et al. 2018; Jumper et al. 2021; Brandes et al. 2022; Thirunavukarasu et al. 2023). As the most essential foundation, a feedforward neural network is typically constructed by a series of linear transformations and nonlinear activations. To be specific, a function $f$ implemented by a feedforward neural network with $L$ layers can be represented as

Here, $\circ$ is the functional composition operator, $\phi_{i}(z)=W_{i}^{\top}z+b_{i}$ denotes the linear transformation and $\sigma_{i}$ represents the elementwise nonlinear activation function. Despite a clear architecture, the formula (1) provides limited insight into how the information within the input data is processed by the neural networks. To comprehensively understand how neural networks retrieve task-relevant information, there have been extensive efforts towards unraveling their interpretability (Doshi-Velez and Kim 2017; Guidotti et al. 2018; Zhang et al. 2021). For instance, post-hoc interpretability (Lipton 2018) focused on the predictions generated by neural networks, while disregarding the detailed mechanism and feature importance. Ghorbani, Abid, and Zou (2019) highlighted the fragility of this type of interpretation, as indistinguishable perturbations could result in completely different interpretations.

For the sake of striking a balance between the complexity of representation and the power of prediction, Tishby and Zaslavsky (2015) and Saxe et al. (2019) pioneered the interpretability of deep neural networks via the information bottleneck theory. Ghosh (2022) further linked the information bottleneck to sufficient dimension reduction (SDR) (Li 1991; Cook 1996), which is a rapidly developing research field in regression diagnostics, data visualization, pattern recognition and etc.

In this paper, we provide a theoretical understanding of neural networks for representation learning from the SDR perspective. Let $x\in\mathbb{R}^{p}$ represent the covariates and $y\in\mathbb{R}$ represent the response. Consider the following regression model

where $B_{0}\in\mathbb{R}^{p\times d}$ is a nonrandom matrix with $d\leq p$, $f_{0}:\mathbb{R}^{d}\to\mathbb{R}$ is an unknown function, and $\epsilon$ is the noise such that $\mathbb{E}(\epsilon|x)=0$ and $\mathrm{Var}(\epsilon|x)=\nu^{2}$ for some positive constant $\nu$. Intuitively, the semiparametric and potentially multi-index model (2) asserts that the core information for regression is encoded in the low-dimensional linear representation $B_{0}^{\top}x$. In the literature of SDR, model (2) has been extensively studied. Based on model (2) , Cook and Li (2002) proposed the objective of sufficient mean dimension reduction as

for some matrix $B$, where $\perp\!\!\!\perp$ means statistical independence. Denote the column space spanned by $B$ as $\Pi_{B}$, which is commonly referred to as the mean dimension-reduction subspace. It is evident that the matrix $B$ satisfying (3) is far away from uniqueness. Hence, we focus on the intersection of all possible mean dimension-reduction subspace, which is itself a mean dimension-reduction subspace under mild conditions (e.g., when the domain of $x$ is open and convex; see Cook and Li (2002)), and term it the central mean subspace. Agreeing to condition (3), $B_{0}$ defined in model (2) yields the central mean subspace, denoted as $\Pi_{B_{0}}$, under certain assumption. Statistical estimation and inference about the central mean subspace is the primary goal of SDR. Popular statistical methods for recovering central mean subspace include ordinary least squares (Li and Duan 1989), principal Hessian directions (Li 1992), minimum average variance estimation (Xia et al. 2002), semiparametric approach (Ma and Zhu 2013), generalized kernel-based dimension reduction (Fukumizu and Leng 2014), and many others. Although numerous studies have demonstrated the ability of neural networks to approximate complex functions (Hornik, Stinchcombe, and White 1989; Barron 1993; Yarotsky 2017; Shen, Yang, and Zhang 2021) and to adapt low-dimensional structures (Bach 2017; Bauer and Kohler 2019; Schmidt-Hieber 2020; Abbe, Adsera, and Misiakiewicz 2022; Jiao et al. 2023; Troiani et al. 2024), it is of subsequent interest to investigate whether neural networks are capable of correctly identifying the intrinsic structure encapsulated in $B_{0}$, thereby deepening the interpretation of neural networks.

Our study was inspired by an observation that the weight matrix in the first layer, i.e., $W_{1}$, could accurately detect the presence of $B_{0}$ in a toy data set, with the rank regularization. Specifically, consider the toy model $y=(B_{0}^{\top}x)^{3}+\epsilon$ where $x\sim\mathrm{Uniform}([0,1]^{10})$, $B_{0}=(1,-2,0,\ldots,0)^{\top}$ and $\epsilon\sim\mathrm{Normal}(0,0.1^{2})$. We trained neural networks using the least squares loss with $W_{1}=W_{11}W_{12}$ where $W_{11}\in\mathbb{R}^{10\times q}$ and $W_{12}\in\mathbb{R}^{q\times 64}$ for $q=1,\ldots,10$. Evidently, the rank of $W_{1}$ did not exceed $q$. It was then observed that for each $q$, (i) $B_{0}$ was closely contained within the column space of $W_{11}$, as the absolute cosine similarity between $B_{0}$ and its projection on $\Pi_{W_{11}}$ was close to 1, and (ii) the leading eigenvector of $W_{1}W_{1}^{\top}$ closely aligned with $B_{0}$ (see Figure 1). This observation indicates that the first layer of the neural network may potentially discover the underlying low-dimensional intrinsic structure.

It is important to note that the application of neural networks for estimating the central mean subspace based on model (2) was previously explored by Kapla, Fertl, and Bura (2022) through a two-stage method. The first stage focused on obtaining a preliminary estimation of $\Pi_{B_{0}}$, which subsequently served as an initial point for the joint estimation of $\Pi_{B_{0}}$ and $f_{0}$ in the second stage. Given our toy example, however, it is prudent to critically evaluate the necessity of the first stage. Furthermore, their work lacks a comprehensive theoretical guarantee. Additionally, another related work conducted by Liang, Sun, and Liang (2022) concentrated on seeking nonlinear sufficient representations. In contrast, we focus more on revealing the fundamental nature of neural networks.

We in this paper show that, with suitable rank regularization, the first layer of a feedforward neural network conducts SDR in a regression task, wherein $d(W_{11},B_{0})\to 0$ in probability for certain distance metric $d(\cdot,\cdot)$. Furthermore, numerical experiments provide empirical evidences supporting this result, while demonstrating the efficiency of neural networks in addressing the issue of SDR.

Throughout this paper, we use $\|v\|_{2}$ to represent the Euclidean norm of a vector $v$. For a matrix $A$, $\|A\|_{F}$ is the Frobenius norm of $A$, $\pi_{A}=A(A^{\top}A)^{-}A^{\top}$ denotes the projection matrix of $A$ where $A^{-}$ is the generalized inverse of $A$, and $\Pi_{A}$ stands for the linear space spanned by the columns of $A$. For a measurable function $f:\mathcal{X}\to\mathbb{R}$, $\|f\|_{L^{2}(\mu)}$ represents the $L^{2}$ norm of $f$ with respect of a given probability measure $\mu$, and $\|f\|_{L^{\infty}(\mathcal{X})}=\sup_{x\in\mathcal{X}}|f(x)|$ represents the supreme norm of $f$ over a set $\mathcal{X}$. $\mathcal{B}(\mathcal{X})$ is the unit ball induced by $\mathcal{X}$, such that $\mathcal{B}(\mathcal{X})\subset\mathcal{X}$ and $\|v\|_{2}\leq 1$ for any $v\in\mathcal{X}$.

From the perspective of numerical experiments, we utilized simulated data sets to demonstrate that (i) the column space of $\mathcal{T}(\hat{f}_{n})$ approached the central mean subspace $\Pi_{B_{0}}$ as sample size $n$ increased, (ii) the performance of neural network-based method in conducting SDR was comparable to classical methods. In some cases, the neural network-based method outperformed classical methods, particularly when the latent intrinsic dimension $d$ was large. Five commonly used methods, sliced inverse regression (SIR), sliced average variance estimation (SAVE), principal Hessian directions (PHD), minimum average variance estimation (MAVE) and generalized kernel dimension reduction (GKDR), were included for comparisons.

The neural network-based method was implemented in Python using PyTorch. Specifically, we utilized a linear layer without bias term to represent the matrix $B$ (recall that we suppose $d$ is known), which was further appended by a fully-connected feedforward neural network with number of neurons being $h-h/2-1$. Here, we set $h=64$ when the sample size was less than $1000$, and set $h=128$ when sample size was between 1000 and 2000. Overall, the neural network architecture was $p-d-h-h/2-1$. Code is available at https://github.com/oaksword/DRNN.

We considered the following four scenarios, with two of them attaining small $d=1,2$ and the rest of $d\geq 3$.

Setting 1: $y=x_{1}^{4}+\epsilon$ where $x\sim\mathrm{Normal}(0,I_{p})$ and $\epsilon\sim\mathrm{t}_{5}$.

Setting 2: $y=\log(x_{1}+x_{1}x_{2})+\epsilon$ where $x\sim\mathrm{Uniform}([0,1]^{10})$ and $\epsilon\sim\mathrm{Normal}(0,\sigma^{2})$.

Setting 3: $y=(1+\beta_{1}^{\top}x)^{2}\exp(\beta_{2}^{\top}x)+5\cdot 1(\beta_{3}^{\top}x>0)+\epsilon$ where $x\sim\mathrm{Uniform}([-1,1]^{6})$, $\epsilon\sim\chi^{2}_{2}-2$, $\beta_{1}=(-2,-1,0,1,2,3)^{\top}$, $\beta_{2}=1_{6}$ and $\beta_{3}=(1,-1,1,-1,1,-1)^{\top}$.

Setting 4:

where $\epsilon\sim\mathrm{Normal}(0,0.1^{2})$,

In setting 1, we tested the combinations of $(n,p)=(100,10),(200,30)$ and $(300,30)$, where $n$ was the sample size. In setting 2, we fixed $n=200,p=10$ and set $\sigma=0.1,0.2,0.5$. In setting 3, we fixed $p=6$ and tested $n=200,500,1000$. In setting 4, we fixed $p=10$ and tested $(n,d)=(1000,4),(1500,6)$ and $(2000,8)$. Settings 1, 2, 4 were equipped with continuous regression functions, and setting 3 involved discontinuity. We set the number of iterations of our method to 1000 for settings 1 and 2, and increased the iterations to 2000 and 4000 for settings 3 and 4, due to their high complexity.

To evaluate the performance of each method, we calculated the distance metric $\|\pi_{\hat{B}}-\pi_{B_{0}}\|_{F}$ where $\hat{B}$ represented the estimate of $B_{0}$. Particularly, for our neural network-based method, $\hat{B}=\mathcal{T}(\hat{f}_{n})$. Smaller value of this metric indicated better performance. We ran 100 replicates for each setting. The results are displayed in Table 1 and Figure 2.

In simple cases (settings 1 and 2), the neural network-based method showed similar performance to classical methods, with sliced average variance estimation being the most effective method in setting 2. However, complex settings demonstrated the significant superiority of neural network-based method compared to other methods. Specifically, in setting 3, as the sample size $n$ increased, the metric $\|\pi_{\hat{B}}-\pi_{B_{0}}\|_{F}$ decreased rapidly, indicating the convergence of $\Pi_{\hat{B}}$ to $\Pi_{B_{0}}$. According to the results in setting 4, the neural network-based method was capable of handling higher-dimensional scenarios. In summary, numerical studies advocated neural networks as a powerful tool for SDR.

We further applied the neural network-based method involving rank regularization to a real regression task. In practice, the precise intrinsic dimension $d$ is unknown, and it is uncertain whether a low-dimensional structure exists. To address this issue, we used cross validation to determine an appropriate $d$ from the range of values $\{1,\ldots,p\}$. More specifically, the raw data set was divided into 80% training data and 20% testing data. The optimal value of $d$ was determined through cross validation on the training data. Subsequently, the final model was fit using the selected $d$ on the training data, and the mean squared error on the testing data was evaluated.

In order to reduce randomness, the aforementioned process was repeated 20 times and the resulting testing errors were recorded. Additionally, we conducted a comparative analysis between the neural network-based approach and alternative methods including vanilla neural network without rank regularization, SIR-based regression, SAVE-based regression, and MAVE-based regression. For the latter three techniques, we initially executed SDR to acquire the embedded data, followed by the utilization of a fully-connected neural network for predictive purpose. The optimal value for $d$ was also determined through cross validation.

We utilized a data set of weather records from Seoul, South Korea during the summer months from 2013 to 2017 (Cho et al. 2020), available at the UCI data set repository (bias correction of numerical prediction model temperature forecast). This data set contained 7750 observations with 23 predictors and 2 responses, specifically the next-day maximum air temperature and next-day minimum air temperature. After excluding the variables for station and date, the data set was reduced to 21 predictors, which were further standardized using StandardScaler in scikit-learn package.

It is evident from Table 2 that the neural network-based method with rank regularization outperformed other methods, demonstrating the effectiveness of the modification compared to the vanilla neural network, and the sound performance in reducing dimensions compared to other SDR methods. The presence of latent structures was partially supported by the averaged optimal $d$. It was possible that 19 or 20 combinations of raw predictors might be sufficient, as opposed to the original 21 predictors.

We have demonstrated that neural networks attain the capability of detecting the underlying low-dimensional structure that preserves all information in $x$ about the conditional mean function $\mathbb{E}(y|x)$. As a result, neural networks are suitable to be utilized for the estimation of central mean subspace $\Pi_{B_{0}}$. The theoretical investigations sharpen our understanding of neural networks, while broadening the scope of SDR as well.

In the context of SDR, a more general scenario than sufficient mean dimension reduction emerges when considering

And the column space spanned by $B_{0}$ corresponds to the central subspace (Cook 1998a; Li 2018). It is clear that relation (4) is equivalent to that $p(y|x)=p(y|B_{0}^{\top}x)$, where $p(\cdot|\cdot)$ represents the conditional probability density function. Following the work of Xia (2007), we can further adapt neural networks to estimate the central subspace by modifying the loss function.

Under mild conditions, Xia (2007) showed that

Here, $(x_{0},y_{0})$ is a fixed point, and $K_{h}(\cdot)$ is a suitable kernel function with a bandwidth $h$. Such finding then implies that

Based on this discovery, it is natural to employ a neural network function $f(B^{\top}x,y_{0}):\mathbb{R}^{d+1}\to\mathbb{R}$ to approximate $p(y_{0}|B_{0}^{\top}x)$, and obtain the estimate of $B_{0}$ at the population level by solving the following problem

where $\tilde{y}$ is an independent copy of $y$. Empirically, we define the loss function as