Understanding Dynamics of Nonlinear Representation Learning and Its Application

LeCun, Bengio, and Hinton (2015) described deep learning as one of hierarchical nonlinear representation learning approaches:

In applications such as computer vision and natural language processing, the success of deep learning can be attributed to its ability to learn hierarchical nonlinear representations by automatically changing nonlinear features and kernels during training based on the given data. This is in contrast to classical machine-learning methods where representations or equivalently nonlinear features and kernels are fixed during training.

Deep learning in practical regimes, which has the ability to learn nonlinear representation (Bengio, Courville, & Vincent, 2013), has had a profound impact in many areas, including object recognition in computer vision (Rifai et al., 2011; Hinton, Osindero, & Teh, 2006; Bengio, Lamblin, Popovici, & Larochelle, 2007; Ciregan, Meier, & Schmidhuber, 2012; Krizhevsky, Sutskever, & Hinton, 2012), style transfer (Gatys, Ecker, & Bethge, 2016; Luan, Paris, Shechtman, & Bala, 2017), image super-resolution (Dong, Loy, He, & Tang, 2014), speech recognition (Dahl, Ranzato, Mohamed, & Hinton, 2010; Deng et al., 2010; Seide, Li, & Yu, 2011; Mohamed, Dahl, & Hinton, 2011; Dahl, Yu, Deng, & Acero, 2011; Hinton et al., 2012), machine translation (Schwenk, Rousseau, & Attik, 2012; Le, Oparin, Allauzen, Gauvain, & Yvon, 2012), paraphrase detection (Socher, Huang, Pennington, Ng, & Manning, 2011), word sense disambiguation (Bordes, Glorot, Weston, & Bengio, 2012), and sentiment analysis (Glorot, Bordes, & Bengio, 2011; Socher, Pennington, Huang, Ng, & Manning, 2011). However, we do not yet know the precise condition that makes deep learning tractable in the practical regime of representation learning.

For example, one of the simplest models is the linear model in the form of $f(x,θ)=ϕ(x)⊤θ$⁠, where $ϕ:X→Rd$ is a fixed function and $ϕ(x)$ is a nonlinear representation of input data $x$⁠. This is a classical machine learning model where much of the effort goes into the design of the handcrafted feature map $ϕ$ via feature engineering (Turner, Fuggetta, Lavazza, & Wolf, 1999; Zheng & Casari, 2018). In this linear model, we do not learn the representation $ϕ(x)$ because the feature map $ϕ$ is fixed without dependence on the model parameter $θ$ that is optimized with the data set $((xi,yi))i=1n$⁠.

Similar to many definitions in mathematics, where an intuitive notion in a special case is formalized to a definition for a more general case, we now abstract and generalize this intuitive notion of the representation $ϕ(x)$ of the linear model to that of all differentiable models as follows:

In this letter, we initiate the study of the dynamics of learning gradient representation that are nonlinear in $Δt$⁠. That is, we focus on the regime where the gradient representation $∂f(x,θt)∂θt$ at the end of training time $t$ differs greatly from the initial representation $∂f(x,θ0)∂θ0$⁠. This regime was studied in the past for the case where the function $ϕ(x)↦f(x,θ)$ is affine for some fixed feature map $ϕ$ (Saxe, McClelland, & Ganguli, 2014; Kawaguchi, 2016, 2021; Laurent & Brecht, 2018; Bartlett, Helmbold, & Long, 2019; Zou, Long, & Gu, 2020; Xu et al., 2021). Unlike any previous studies, we focus on the problem setting where the function $ϕ(x)↦f(x,θ)$ is nonlinear and nonaffine, with the effect of nonlinear (gradient) representation learning. The results of this letter avoid the curse of dimensionality by studying the global convergence of the gradient-based dynamics instead of the dynamics of global optimization (Kawaguchi et al., 2016) and Bayesian optimization (Kawaguchi, Kaelbling, & Lozano-Pérez, 2015). Importantly, we do not require any wide layer or large input dimension throughout this letter. Our main contributions are summarized as follows:

1. In section 2, we identify a pair of a novel assumption and a new condition, called the common model structure assumption and the data-architecture alignment condition. Under the common model structure assumption, the data-architecture alignment condition is shown to be a necessary condition for the globally optimal model and a sufficient condition for the global convergence. The condition is dependent on both data and architecture. Moreover, we empirically verify and deepen this new understanding. When we apply representation learning in practice, we often have overwhelming options regarding which model structure to be used. Our results provide a practical guidance for choosing or designing model structure via the common model structure assumption, which is indeed satisfied by many representation learning models used in practice.

2. In section 3, we discard the assumption of the data-architecture alignment condition. Instead, we derive a novel training framework, called the exploration-exploitation wrapper (EE wrapper), which satisfies the data-architecture alignment condition time-independently a priori. The EE wrapper is then proved to have global convergence guarantees under the safe-exploration condition. The safe-exploration condition is what allows us to explore various gradient representations safely without getting stuck in the states where we cannot provably satisfy the data-architecture alignment condition. The safe-exploration condition is shown to hold true for ResNet-18 with standard benchmark data sets, including MNIST, CIFAR-10, CIFAR-100, Semeion, KMNIST, and SVHN time-independently.

3. In section 3.4, the EE wrapper is shown to not degrade practical performances of ResNet-18 for the standard data sets, MNIST, CIFAR-10, CIFAR-100, Semeion, KMNIST, and SVHN. To our knowledge, this letter provides the first practical algorithm with global convergence guarantees without degrading practical performances of ResNet-18 on these standard data sets, using convolutions, skip connections, and batch normalization without any extremely wide layer of the width larger than the number of data points. To the best of our knowledge, we are not aware of any similar algorithms with global convergence guarantees in the regime of learning nonlinear representations without degrading practical performances.

In this section, we identify the common model structure assumption and study the data-architecture alignment condition for the global convergence in nonlinear representation learning. We begin by presenting an overview of our results in section 2.1, deepen our understandings with experiments in section 2.2, discuss implications of our results in section 2.3, and establish mathematical theories in section 2.4.

We introduce the common model structure assumption in section 2.1.1 and define the data-architecture alignment condition in section 2.1.2. Using the assumption and the condition, we present the global convergence result in section 2.1.3.

Through examinations of representation learning models used in applications, we identified and formalized one of their common properties as follows:

Assumption 1 is satisfied by common machine learning models, such as kernel models and multilayer neural networks, with or without convolutions, batch normalization, pooling, and skip connections. For example, consider a multilayer neural network of the form $f(x,θ)=Wh(x,u)+b$⁠, where $h(x,u)$ is an output of its last hidden layer and the parameter vector $θ$ consists of the parameters $(W,b)$ at the last layer and the parameters $u$ in all other layers as $θ=vec([W,b,u])$⁠. Here, for any matrix $M∈Rm×m¯$⁠, we let $vec(M)∈Rmm¯$ be the standard vectorization of the matrix $M$ by stacking columns. Then, assumption 1 holds because $f(x,θ)=∑k=1d1{k∈S}θk(∂f(xi,θ)∂θk)$⁠, where $S$ is defined by ${θk:k∈S}={vec([W,b])k:k=1,2,…,ξ}$ with $vec([W,b])∈Rξ$⁠. Since $h$ is arbitrary in this example, the common model structure assumption holds, for example, for any multilayer neural networks with a fully connected last layer. In general, because the nonlinearity at the output layer can be treated as a part of the loss function $ℓ$ while preserving convexity of $q↦ℓ(q,y)$ (e.g., cross-entropy loss with softmax), this assumption is satisfied by many machine learning models, including ResNet-18 and all models used in the experiments in this letter (as well as all linear models). Moreover, assumption 1 is automatically satisfied in the next section by using the EE wrapper.

Given a target matrix $Y=(y1,y2,…,yn)⊤∈Rn×my$ and a loss function $ℓ$⁠, we define the modified target matrix $Yℓ=(y1ℓ,y2ℓ,…,ynℓ)⊤∈Rn×my$ by $Yℓ=Y$ for the squared loss $ℓ$⁠, and by $(Yℓ)ij=2Yij-1$ for the (binary and multiclass) cross-entropy losses $ℓ$ with $Yij∈{0,1}$⁠. Given input matrix $X=(x1,x2,…,xn)⊤∈Rn×mx$⁠, the output matrix $fX(θ)∈Rn×my$ is defined by $fX(θ)ij=f(xi,θ)j∈R$⁠. For any matrix $M∈Rm×m¯$⁠, we let $Col(M)⊆Rm$ be its column space. With these notations, we are now ready to introduce the data-architecture alignment condition:

The data-architecture alignment condition depends on both data (through the target $Y$ and the input $X$⁠) and architecture (through the model $f$⁠). It is satisfied only when the data and architecture align well to each other. For example, in the case of linear model $f(x,θ)=ϕ(x)⊤θ∈R$⁠, the condition can be written as $vec(Yℓ)∈Col(Φ(X))$ where $Φ(X)∈Rn×d$ and $Φ(X)ij=ϕ(xi)j$⁠. In definition 2, $fX(θ)$ is a matrix of the preactivation outputs of the last layer. Thus, in the case of classification tasks with a nonlinear activation at the output layer, $fX(θ)$ and $Y$ are not in the same space, which is the reason we use $Yℓ$ here instead of $Y$⁠.

Importantly, the data-architecture alignment condition does not make any requirements on the the rank of the Jacobian matrix $∂vec(fX(θ))∂θ∈Rnmy×d$⁠: the rank of $∂vec(fX(θ))∂θ$ is allowed to be smaller than $nmy$ and $d$⁠. Thus, for example, the data-architecture alignment condition can be satisfied depending on the given data and architecture even if the minimum eigenvalue of the matrix $∂vec(fX(θ))∂θ(∂vec(fX(θ))∂θ)⊤$ is zero, in both cases of overparameterization (e.g., $d≫n$⁠) and underparameterization (e.g., $d≪n$⁠). This is further illustrated in section 2.2 and discussed in section 2.3. We note that we further discard the assumption of the data-architecture alignment condition in section 3 as it is automatically satisfied by using the EE wrapper.

Under the common model structure assumption, the data-architecture alignment condition is shown to be what lets us avoid the failure of the global convergence and suboptimal local minima. More concretely, we prove a global convergence guarantee under the data-architecture alignment condition as well as the necessity of the condition for the global optimality:

Theorem 1i guarantees the global convergence without the need to satisfy the data-architecture alignment condition at every iteration or at the limit point. Instead, it shows that the bound on the global optimality gap decreases toward zero per iteration whenever the data-architecture alignment condition holds. Theorem 1ii shows that the data-architecture alignment condition is necessary for the global optimality. Intuitively, this is because the expressivity of a model class satisfying the common model structure assumption is restricted such that it is required to align the architecture to the data in order to contain the globally optimal model $fX(θ)=ηYℓ$ (for any $η∈R$⁠).

To better understand the statement of theorem 1i, consider a counterexample with a data set consisting of the single point $(x,y)=(1,0)$⁠, the model $f(x,θ)=θ4-10θ2+6θ+100$⁠, and the squared loss $ℓ(q,y)=(q-y)2$⁠. In this example, we have $L(θ)=f(x,θ)2$⁠, which has multiple suboptimal local minima of different values. Then, via gradient descent, the model converges to the closest local minimum and, in particular, does not necessarily converge to a global minimum. Indeed, this example violates the common model structure assumption (assumption 1) (although it satisfies the data-architecture alignment condition), showing the importance of the common model structure assumption along with the data-architecture alignment. This also illustrates the nontriviality of theorem 1i in that the data-architecture alignment is not sufficient, and we needed to understand what types of model structures are commonly used in practice and formalize the understanding as the common model structure assumption.

Since $f(x1),f(x2),…,f(xn)∈Rmy$⁠, this implies that any stationary point $θ$ is a global minimum if the minimum eigenvalue of the matrix $∂vec(fX(θ))∂θ(∂vec(fX(θ))∂θ)⊤$ is nonzero, without the common model structure assumption (see assumption 1). Indeed, in the above example with the model $f(x,θ)=θ4-10θ2+6θ+100$⁠, the common model structure assumption is violated, but we still have the global convergence if the minimum eigenvalue is nonzero—for example, $f(x,θ)=y=0$ at any stationary point $θ$ such that the minimum eigenvalue of the matrix $∂vec(fX(θ))∂θ(∂vec(fX(θ))∂θ)⊤$ is nonzero. In contrast, theorem 1 allows the global convergence even when the minimum eigenvalue of the matrix $∂vec(fX(θ))∂θ(∂vec(fX(θ))∂θ)⊤$ is zero by utilizing the common model structure assumption.

The formal version of theorem 1 is presented in section 2.4 and is proved in appendix A in the supplementary information that relies on the additional previous works of Clanuwat et al. (2019), Krizhevsky and Hinton (2009), Mityagin (2015), Netzer et al. (2011), Paszke et al. (2019a, 2019b), and Poggio et al. (2017). Before proving the statement, we first examine the meaning and implications of our results through illustrative examples in sections 2.2 and 2.3.

Theorem 1 suggests that data-architecture alignment condition $vec(Yℓ)∈Col(∂vec(fX(θt))∂θt)$ has the ability to distinguish the success and failure cases, even when the minimum eigenvalue of the matrix $∂vec(fX(θt))∂θt(∂vec(fX(θt))∂θt)⊤$ is zero for all $t≥0$⁠. In this section, we conduct experiments to further verify and deepen this theoretical understanding.

We employ a fully connected network having four layers with 300 neurons per hidden layer, and a convolutional network, LeNet (LeCun et al., 1998), with five layers. For the fully connected network, we use the two-moons data set (Pedregosa et al., 2011) and a sine wave data set. To create the sine wave data set, we randomly generated the input $xi$ from the uniform distribution on the interval $[-1,1]$ and set $yi=1{sin(20xi)<0}∈R$ for all $i∈[n]$ with $n=100$⁠. For the convolutional network, we use the Semeion data set (Srl & Brescia, 1994) and a random data set. The random data set was created by randomly generating each pixel of the input image $xi∈R16×16×1$ from the standard normal distribution and by sampling $yi$ uniformly from ${0,1}$ for all $i∈[n]$ with $n=1000$⁠. We set the activation functions of all layers to be softplus $σ¯(z)=ln(1+exp(ςz))/ς$ with $ς=100$⁠, which approximately behaves as the ReLU activation as shown in appendix C in the supplementary information. See appendix B in the supplementary information for more details of the experimental settings.

Figures 1b and 1c show the results for the convolutional networks with two random initial points using two different random seeds. In the figure panels, we report the training behaviors with different network sizes $mc=1,2$⁠, and 4; the number of convolutional filters per convolutional layer is $8×mc$ and the number of neurons per fully connected hidden layer is $128×mc$⁠. As can be seen, with the Semeion data set, the networks of all sizes achieved zero error with $QT=0$ for all $T$⁠. With the random data set, the deep networks yielded the zero training error whenever $QT$ is not linearly increasing over the time or, equivalently, whenever the condition of $vec(Yℓ)∈Col(∂vec(fX(θT))∂θT)$ holds sufficiently many steps $T$⁠. This is consistent with our theory.

Finally, we also confirmed that gradient representation $∂f(x,θt)∂θt$ changed significantly from the initial one $∂f(x,θ0)∂θ0$ in our experiments. That is, the values of $∥M(θT)-M(θ0)∥F2$ were significantly large and tended to increase as $T$ increases, where the matrix $M(θ)∈Rnmy×nmy$ is defined by $M(θ)=∂vec(fX(θ))∂θ(∂vec(fX(θ))∂θ)⊤$⁠. Table 1 summarizes the values of $∥M(θT)-M(θ0)∥F2$ at the end of the training.

In section 2.1.3, we showed that an uncommon model structure $f(x,θ)=θ4-10θ2+6θ+100$ does not satisfy assumption 1, and assumption 1 is not required for global convergence if the minimum eigenvalue is nonzero. However, in practice, we typically use machine learning models that satisfy assumption 1 instead of the model $f(x,θ)=θ4-10θ2+6θ+100$⁠, and the minimum eigenvalue is zero in many cases. In this context, theorem 1 provides the justification for common practice in nonlinear representation learning. Furthermore, theorem 1i contributes to the literature by identifying the common model structure assumption (assumption 1) and the data-architecture alignment condition (definition 1) as the novel and practical conditions to ensure the global convergence even when the minimum eigenvalue becomes zero. Moreover, theorem 1ii shows that this condition is not arbitrary in the sense that it is also necessary to obtain the globally optimal models. Furthermore, the data-architecture alignment condition is strictly more general than the condition of the minimum eigenvalue being nonzero, in the sense that the latter implies the former but not vice versa.

Our new theoretical understanding based on the data-architecture alignment condition can explain and deepen the previously known empirical observation that increasing network size tends to improve training behaviors. Indeed, the size of networks seems to correlate well with the training error to a certain degree in Figure 1b. However, the size and the training error do not correlate well in Figure 1c. Our new theoretical understanding explains that the training behaviors correlate more directly with the data-architecture alignment condition of $vec(Yℓ)∈Col(∂vec(fX(θt))∂θt)$ instead. The seeming correlation with the network size is indirect and caused by another correlation between the network size and the condition of $vec(Yℓ)∈Col(∂vec(fX(θt))∂θt)$⁠. That is, the condition of $vec(Yℓ)∈Col(∂vec(fX(θt))∂θt)$ more likely tends to hold when the network size is larger because the matrix $∂vec(fX(θt))∂θt$ is of size $nmy×d$ where $d$ is the number of parameters: that is, by increasing $d$⁠, we can increase the column space $Col(∂vec(fX(θt))∂θt)$ to increase the chance of satisfying the condition of $vec(Yℓ)∈Col(∂vec(fX(θt))∂θt)$⁠.

Note that the minimum eigenvalue of the matrix $M(θt)=∂vec(fX(θ))∂θ(∂vec(fX(θ))∂θ)⊤$ is zero at all iterations in Figures 1b and 1c for all cases of $mc=1$⁠. Thus, Figures 1b and 1c also illustrate the fact that while having the zero minimum eigenvalue of the matrix $M(θt)$⁠, the dynamics can achieve the global convergence under the data-architecture alignment condition. Moreover, because the multilayer neural network in the lazy training regime (Kawaguchi & Sun, 2021) achieves zero training errors for all data sets, Figure 1 additionally illustrates that our theoretical and empirical results apply to the models outside of the lazy training regime and can distinguish “good” data sets from “bad” data sets given a learning algorithm.

In sum, our new theoretical understanding has the ability to explain and distinguish the successful case and failure case based on the data-architecture alignment condition for the common machine learning models. Because the data-architecture alignment condition is dependent on data and architecture, theorem 1, along with our experimental results, shows why and when the global convergence in nonlinear representation learning is achieved based on the relationship between the data $(X,Y)$ and architecture $f$⁠. This new understanding is used in section 3 to derive a practical algorithm and is expected to be a basis for many future algorithms.

This section presents the precise mathematical statements that formalize the informal description of theorem 1. In the following sections, we devide the formal version of theorem 1 into theorem 2, theorem 3, and proposition 1.

Let $(θt)t=0∞$ be the sequence defined by $θt+1=θt-αtg¯t$ with an initial parameter vector $θ0$⁠, a learning rate $αt$⁠, and an update vector $g¯t$⁠. The analysis in this section relies on the following assumption on the update vector $g¯t$⁠:

Assumption 2 is satisfied by using $g¯t=Dt∇L(θt),$ where $Dt$ is any positive-definite symmetric matrix with eigenvalues in the interval $[c̲,c¯]$⁠. If we set $Dt=I$⁠, we have gradient descent, and assumption 2 is satisfied with $c̲=c¯=1$⁠. This section also uses the standard assumption of differentiability and Lipschitz continuity:

The assumptions on the loss function in assumption 3 are satisfied by using standard loss functions, including the squared loss, logistic loss, and cross-entropy loss. Although the objective function $L$ is nonconvex and non-invex, the function $q↦ℓ(q,yi)$ is typically convex.

The following theorem shows that every limit point $θ^$ of the sequence $(θt)t$ achieves a loss value $L(θ^)$ no worse than $infη∈RL*(ηY*)$ for any $Y*$ such that $vec(Y*)∈Col(∂vec(fX(θt))∂θt)$ for all $t∈[τ,∞)$ with some $τ≥0$⁠:

In practice, one can easily satisfy all the assumptions in theorem 2 except for the condition that $vec(Y*)∈Col(∂vec(fX(θt))∂θt)$ for all $t∈[τ,∞)$⁠. Accordingly, we now weaken this condition by analyzing optimality at each iteration so that the condition is verifiable in experiments.

The following theorem states that under standard settings, the sequence $(θt)t∈T$ converges to a loss value no worse than $infη∈RL*(ηY*)$ at the rate of $O(1/|T|)$ for any $T$ and $Y*$ such that $vec(Y*)∈Col(∂vec(fX(θt))∂θt)$ for $t∈T$⁠:

The results in the previous section show that the bound on the global optimality gap decreases per iteration whenever the data-architecture alignment condition holds. Using this theoretical understanding, in this section, we propose a new training framework with prior guarantees while learning hierarchical nonlinear representations without assuming the data-architecture alignment condition. As a result, we made significant improvements over the most closely related study on global convergence guarantees (Kawaguchi & Sun, 2021). In particular, whereas the related study requires a wide layer with a width larger than $n$⁠, our results reduce the requirement to a layer with a width larger than $n$⁠. For example, the MNIST data set has $n=60,000$ and hence previous studies require 60,000 neurons at a layer, whereas we only require $60,000≈245$ neurons at a layer. Our requirement is consistent and satisfied by the models used in practice that typically have from 256 to 1024 neurons for some layers.

We begin in section 3.1 with additional notations and then present the training framework in section 3.2 and convergence analysis in section 3.3. We conclude in section 3.4 by providing empirical evidence to support our theory.

For an arbitrary matrix $M∈Rm×m'$⁠, we let $M*j$ be its $j$th column vector in $Rm$⁠, $Mi*$ be its $i$th row vector in $Rm'$⁠, and $rank(M)$ be its matrix rank. We define $M∘M'$ to be the Hadamard product of any matrices $M$ and $M'$⁠. For any vector $v∈Rm$⁠, we let $diag(v)∈Rm×m$ be the diagonal matrix with $diag(v)ii=vi$ for $i∈[m]$⁠. We denote by $Im$ the $m×m$ identity matrix.

In this section, we introduce the exploration-exploitation (EE) wrapper $A$⁠. The EE wrapper $A$ is not a stand-alone training algorithm. Instead, it takes any training algorithm $G$ as its input and runs the algorithm $G$ in a particular way to guarantee global convergence. We note that the exploitation phase in the EE wrapper does not optimize the last layer; instead, it optimizes hidden layers, whereas the exploration phase optimizes all layers. The EE wrapper allows us to learn the representation $∂f(x,θt)∂θt$ that differs significantly from the initial representation $∂f(x,θ0)∂θ0$ without making assumptions on the minimum eigenvalue of the matrix $∂vec(fX(θ))∂θ(∂vec(fX(θ))∂θ)⊤$ by leveraging the data-architecture alignment condition. The data-architecture alignment condition is ensured by the safe-exploration condition (defined in section 3.3.1), which is time independent and holds in practical common architectures (as demonstrated in section 3.4).

Algorithm 1 outlines the EE wrapper $A$⁠. During the exploration phase in lines 3 to 7 of algorithm 1, the EE wrapper $A$ freely explores hierarchical nonlinear representations to be learned without any restrictions. Then, during the exploitation phase in lines 8 to 12, it starts exploiting the current knowledge to ensure $vec(Yℓ)∈Col(∂vec(fX(θt))∂θt)$ for all $t$ to guarantee global convergence. The value of $τ$ is the hyperparameter that controls the time when it transitions from the exploration phase to the exploitation phase.

In the exploitation phase, the wrapper $A$ only optimizes the parameter vector $θ(H-1)t$ at the $(H-1)$th hidden layer, instead of the parameter vector $θ(H)t$ at the last layer or the $H$th layer. Despite this, the EE wrapper $A$ is proved to converge to global minima of all layers in $Rd$⁠. The exploitation phase still allows us to significantly change the representations as $M(θt)¬≈M(θτ)$ for $t>τ$⁠. This is because we optimize the hidden layers instead of the last layer without any significant overparameterizations.

The exploitation phase uses an arbitrary optimizer $G˜$ with the update vector $g˜t∼G˜(H-1)(θt,t)$ with $g˜t=αtg^t∈RdH-1$⁠. During the two phases, we can use the same optimizers (e.g., SGD for both $G$ and $G˜$⁠) or different optimizers (e.g., SGD for $G$ and L-BFGS for $G˜$⁠).