Modularizing Deep Learning via Pairwise Learning With Kernels

Throughout, we use bold capital letters for matrices and tensors, bold lower-case letters for vectors, and unbold letters for scalars. $(\mathbf{v})_{i}$ denotes the $i^{\text{th}}$ component of vector $\mathbf{v}$. And $\mathbf{W}^{(j)}$ denotes the $j^{\text{th}}$ column of matrix $\mathbf{W}$. For a 3D tensor $\mathbf{X}$, $\mathbf{X}[:,:,c]$ denotes the $c^{\text{th}}$ matrix indexing along the third dimension from the left (or the $c^{\text{th}}$ channel). We use $\langle\cdot,\cdot\rangle_{H}$ to denote the inner product in an inner product space $H$. And the subscript shall be omitted if doing so causes no confusion. For functions, we use capital letters to denote vector/matrix/tensor-valued functions, and lower-case letters are reserved specifically for scalar-valued ones. In a network, we call a composition of an arbitrary number of layers as a module for convenience.

Since we shall propose an alternative view on NNs that redefine the network layers and also modules, it is helpful to introduce notations to distinguish our view and the conventional one. We use letter $G_{i}$ (or $g_{i}$, depending on if this layer or module is a scalar-valued function) with a numeric subscript $i\in\mathbb{N}\setminus\{0\}$ to refer to the $i^{\text{th}}$ network layer or module under the conventional view and letter $F_{i}$ (or $f_{i}$) the $i^{\text{th}}$ layer or module (of the same model) under our view. Example usages can be found in Fig. 1 and Sec. III.

Kernel machines can be considered as linear models in feature spaces, the mappings into which are potentially nonlinear [11]. Consider a feature map $\Phi:\mathbb{R}^{d}\to H$, where $H$ is a (real) Hilbert space, i.e., an inner product space over the real numbers that is complete in the metric induced by the inner product, a kernel machine is a linear model in this feature space $H$:

where $\mathbf{w},b$ are its trainable weights and bias, respectively. For certain feature spaces $H$ (called RKHSs), one can define a bivariate “kernel” function $k$ that is symmetric and positive semidefinite via

This definition is often used as an identity and is referred to as the “kernel trick” in some more modern texts [11].¹¹1The correspondence between $k$ and $\Phi$ holds true in both directions and is sometimes stated the other way around. Namely, per Moore-Aronszajn Theorem, for every symmetric, continuous, positive semidefinite bivariate function $k$ that maps into $\mathbb{R}$, one can find an RKHS $H$ such that $k(\mathbf{u},\mathbf{v})=\langle\Phi(\mathbf{u}),\Phi(\mathbf{v})\rangle_{H}$, where the feature map $\Phi$ is defined through $k$ as $\Phi(\mathbf{u}):=k(\mathbf{u},\cdot)$ [12].

An RKHS, different from a general Hilbert space, possesses a “canonical coordinate system” induced by the kernel [13] (in addition to the regular coordinate system induced by its basis) that enables one to concisely express distance between feature vectors using kernel values, thanks to the reproducing property [14]. Concretely, we have

One can use the “kernel trick” to implement highly nontrivial feature maps (thus highly complicated functions in the input space). For feature maps that are not implementable, e.g., when $\Phi(\mathbf{x})$ is an infinite series, one can approximate the KM with a set of “centers” $\mathbf{x}_{1},\ldots,\mathbf{x}_{n}$ as follows:

Assuming a $k$ satisfying Eq. 2 can be found and using the “kernel trick”, the right hand side is equal to

Now, the learnable parameters become the $\alpha_{i}$’s anb $b$.

This approximation changes the computational complexity of evaluating the kernel machine on a single example from $\mathcal{O}(h)$, where $h$ is the dimension of the feature space $H$ and can be infinite for certain kernels, to $\mathcal{O}(nd)$, where $d$ is the dimension of the input space. In practice, the runtime for a sample is quadratic in sample size since $n$ is typically on the same order as the sample size. There exists acceleration methods that reduce the complexity via further approximation in this case (e.g., [15]), yet the compromise in performance can be nonnegligible in practice especially when the input space dimension is large.

Note that, for common NNs, the kernels involved in the above methods have implementable feature maps with feature space dimensions being equal to the widths of the corresponding NN layers (so they are always finite). Therefore, one no longer needs to rely on the kernel trick and approximate the kernel machines through pairwise evaluations on a set of centers — one can simply use the explicit linear model representation $f(\mathbf{x})=\langle\mathbf{w},\Phi(\mathbf{x})\rangle+b$! The major advantage is that the computational complexity of running the model on a single data example becomes $\mathcal{O}(p)$, where $p$ is the size of the corresponding NN layer (or equivalently, the size of $\Phi(\mathbf{x})$). This reduces the runtime over a set of data from the usual super-quadratic to linear in sample size, making the model much more practical.

Suppose we have a deep model consisting of two modules $F=F_{2}\circ F_{1}$ and an objective function $L(F,S)$, where $S$ is a training set $S=\{(\mathbf{x}_{i},y_{i})\}_{i=1}^{n}$. $F_{1},F_{2}$ can be compositions of arbitrary layers.

A modular learning algorithm trains $F_{1}$, freezes it afterwards, then trains $F_{2}$. And the goal is that after wiring together the two trained modules, the overall model minimizes $L$.

For a given $S$, define

Clearly, to obtain an $F$ that minimizes $L$, the goal when training $F_{1}$ can be to find an $F_{1}^{\prime}$ that is in $\mathcal{F}_{1}^{\star}$. Then one can simply train $F_{2}$ to minimize $L(F_{2}\circ F_{1}^{\prime},S)$ and the resulting minimizer $F_{2}^{\prime}$ will satisfy $F_{2}^{\prime}\circ F_{1}^{\prime}\in\operatorname*{arg\,min}_{F}L(F,S)$. And if we can characterize $\mathcal{F}_{1}^{\star}$ independently of the trainable parameters of $F_{2}$, the training of $F_{1}$ can be decoupled from that of $F_{2}$.

Earlier we have shown how to redefine NN layers such that they become KMs, but we have not yet demonstrated any practical advantage of such a representation. In this section, we show that under this KM view on the output module, the optimal input module can be described independently of the output module, which, as stated above, is the core idea motivating our modular learning algorithm. For simplicity, we discuss only binary classification. The result easily extend to classification with more classes.

Concretely, we present a result stating that if $F_{2}$ is a single KM with kernel $k$, $\mathcal{F}_{1}^{\star}$ can be characterized independently of the trainable parameters of $F_{2}$ via pairwise evaluations of the kernel $k$ on training data from distinct classes. The proof of the following theorem is given in the Appendix.

Further, if the infimum of $k(\mathbf{u},\mathbf{v})$ is attained in $\mathbb{R}^{d_{1}}\times\mathbb{R}^{d_{1}}$ and equals $\beta$, then Eq. 11 is equivalent to

Interpreting a two-module classifier $F_{2}\circ F_{1}$ as $F_{1}$ learning a new representation of the given data on which $F_{2}$ will carry out the classification, the intuition behind our result can be explained as follows. Given some data, its “optimal” representation for a linear model to classify should be the one where examples from distinct classes are located as distant from each other as possible. Thus, the optimal $F_{1}$ should be the one that produces this optimal representation. And since our $F_{2}$ is a linear model in an RKHS feature space and because distance in an RKHS can be expressed via evaluations of its reproducing kernel, the optimal $F_{1}$ can then be fully described with only pairwise kernel evaluations over the training data, i.e., $k\left(F_{1}(\mathbf{x}_{i}),F_{1}(\mathbf{x}_{j})\right)$ for $\mathbf{x}_{i},\mathbf{x}_{j}$ being training examples from different classes. In other words, an RKHS provides a useful “canonical coordinate system” [13] that enables a concise definition of the optimal input module for any linear classifier in that RKHS using only kernel evaluations.

The earlier Theorem imposes assumptions on the network architecture and the classification objective function. We now show that these assumptions are satisfied in many standard classification set-ups.

To convert the theoretical result into an implementable learning algorithm for classification (with potentially more than two classes), one may proceed as follows. Let an NN $G_{2}\circ G_{1}$ be given, where $G_{2}$ is a linear layer (absorb the trailing nonlinearity into the loss function if necessary) and $G_{1}=\Phi\circ F_{1}$ is a composition of arbitrary layers followed by nonlinearity $\Phi$. Suppose the activation vector of $G_{1}$ is normalized such that it is a unit vector by (elementwise) dividing the activation vector by its norm. Also let a loss function $L$ be given and suppose it (or its two-class analog) satisfies the requirements of Theorem IV.1. Define kernel $k(F_{1}(\mathbf{u}),F_{1}(\mathbf{v}))=\langle\Phi(F_{1}(\mathbf{u})),\Phi(F_{1}(\mathbf{v}))\rangle$. Determine $\beta:=\min k$ based on $\Phi$. Some examples include: $\beta=0$ for ReLU and sigmoid; $\beta=-1$ for tanh. Denote $G_{2}\circ\Phi$ as $F_{2}$. The training consists of two stages: Training $F_{1}$ and training $F_{2}$ (without touching $F_{1}$).

Training $F_{1}$. Given a batch of training data $\{(\mathbf{x}_{i},y_{i})\}_{i=1}^{n}$ and let $\mathcal{N}$ (for negative) denote all pairs of indices $i,j$ such that $y_{i}\neq y_{j}$ and $\mathcal{P}$ (for positive) denote all pairs of indices $i\neq j$ with $y_{i}=y_{j}$. Train $F_{1}$ to maximize one of the following proxy objective functions.

• •

  Alignment (negative only) (AL-NEO):

  $$L_{1}(F_{1})=\frac{\beta\sum_{(i,j)\in\mathcal{N}}k\left(F_{1}(\mathbf{x}_{i}),F_{1}(\mathbf{x}_{j})\right)}{|\beta||\mathcal{N}|\sqrt{\sum_{(i,j)\in\mathcal{N}}\left(k\left(F_{1}(\mathbf{x}_{i}),F_{1}(\mathbf{x}_{j})\right)\right)^{2}}};$$

  (19)

• •

  Contrastive (negative only) (CTS-NEO):

  $$L_{1}(F_{1})=-\frac{1}{|\mathcal{N}|}\sum_{(i,j)\in\mathcal{N}}\exp\left(k\left(F_{1}(\mathbf{x}_{i}),F_{1}(\mathbf{x}_{j})\right)\right);$$

  (20)

• •

  Negative Mean Squared Error (negative only) (NMSE-NEO):

  $$L_{1}(F_{1})=-\frac{1}{|\mathcal{N}|}\sum_{(i,j)\in\mathcal{N}}\left(k\left(F_{1}({x}_{i}),F_{1}(\mathbf{x}_{j})\right)-\beta\right)^{2}.$$

  (21)

All of the above proxy objectives can be shown to learn an $F_{1}$ that satisfies the optimality condition required by Theorem IV.1.

Note that in cases where some of these proxies are undefined (for example, when $\beta=0$), we may train $F_{1}$ to maximize the following alternative proxy objectives instead assuming $\alpha:=\sup k$ is known, where we define $k^{\star}_{ij}$ to be $\alpha$ if $(i,j)\in\mathcal{P}$ or if $i=j$ and $\beta$ if otherwise. Compared to the previous proxies, these are applicable to all $k$’s and impose a stronger constraint on learning. Specifically, in addition to controlling the inter-class pairs, these proxies force intra-class pairs to share representations.

• •

  Alignment (AL): The kernel alignment [21] between the kernel matrix formed by $k$ and $k^{\star}$ on the given data.

• •

  Upper Triangle Alignment (UTAL): Same as AL, except only the upper triangles minus the main diagonals of the two matrices are considered. This can be considered a refined version of the raw alignment.

• •

  Contrastive (CTS):

  $$L_{1}(F_{1})=\frac{\sum_{(i,j)\in\mathcal{P}}\exp\left(k\left(F_{1}(\mathbf{x}_{i}),F_{1}(\mathbf{x}_{j})\right)\right)}{\sum_{(i,j)\in\mathcal{N}\cup\mathcal{P}}\exp\left(k\left(F_{1}(\mathbf{x}_{i}),F_{1}(\mathbf{x}_{j})\right)\right)};$$

  (22)

• •

  Negative Mean Squared Error (NMSE):

  $$L_{1}(F_{1})=-\frac{1}{n^{2}}\sum_{(i,j)}\left(k\left(F_{1}(\mathbf{x}_{i}),F_{1}(\mathbf{x}_{j})\right)-k^{\star}_{ij}\right)^{2}.$$

  (23)

Empirically, we found that alignment and squared error typically produced better results than the contrastive ones, which is potentially due to how the exponential term involved in the contrastive objectives made the range of ideal learning rates to be smaller.

Training $F_{2}$. Now suppose $F_{1}$ has been trained and frozen at $F_{1}^{\prime}$, we simply train $F_{2}$ to minimize the overall loss function $L(F_{2}\circ F_{1}^{\prime})$. This process is illustrated in Fig. 3.

While some works establish connections via exactly matching one architecture to the other, others do so from a probabilistic perspective by studying large-sample behavior in the infinite layer widths limit and/or taking the expectation over the network parameters. The former line of research, to which this work belongs, often yields more direct and practical results since the theories operate under much milder assumptions.

NNs were developed with inspiration from human brain structures and functions [28]. The human brain itself is modular in a hierarchical manner, as the learning process always occurs in a very localized subset of highly inter-connected nodes which are relatively sparsely connected to nodes in other modules [29, 30]. That is to say that the human brain is organized as functional, sparsely connected subunits [31].

Many existing works in machine learning can be analyzed from the perspective of modularization. An old example is the mixture of experts [32, 33] which uses a gating function to enforce each expert in a committee of networks to solve a distinct group of training cases. Another recent example is the generative adversarial networks (GANs) [34]. Typical GANs have two competing neural networks that are essentially decoupled in functionality and can be viewed as two modules. In [35], the authors proposed an a posteriori method that analyzes a trained network as modules in order to extract useful information. Most works in this direction, however, achieved only partial modularization due to their dependence of end-to-end optimization.

The authors of [36] pioneered the idea of greedily learn the architecture of an NN, achieving full modularization. In their work, each new node is added to maximize the correlation between its output and the residual error. Several works attempted to remove the need for end-to-end backpropagation via approximating gradient signals locally at each layer or each node [37, 38, 39, 40, 41]. In [42], a backpropagation-free deep architecture based on decision trees was proposed. In [43], the authors proposed to learn the hidden layers with unsupervised contrastive learning, decoupling their training from that of the output layer. Compared to these existing methods that enable full modularization, our method is simple to implement and provide strong optimality guarantee. A similar provably optimal modular training framework was proposed in [44]. However, their optimality result was less general than ours and their framework required that the NN be modified by substituting certain neurons with KMs using classical kernels that have quadratic runtime.

Describing the task space structure via measuring the transferability between tasks, i.e., estimating to what extent representations learned from one task can help learning another, is a central problem in transfer learning, multi-task learning, meta learning, continual/lifelong learning, etc [22, 23, 10, 24].

Information-theoretic measures have been proven useful in quantifying task transferability. Early task-relatedness measures include the $\mathcal{F}$-relatedness [45] and the $\mathcal{A}$-distance [46]. $\mathcal{H}$-divergence [47] and Wasserstein distance [48] have received increasing attention in recent years. Specifically, [49] applied $\mathcal{H}$-divergence in natural language processing for text classification, whereas [50] used Wasserstein distance to estimate the similarity of linear parameters instead of the data generation distributions. Along this line of research, some more recent methods include H-score [51] and the Bregman-correntropy conditional divergence [52], the latter of which used the correntropy functional [53] and the Bregman matrix divergence [54] to quantify divergence between mappings.

Model-based methods, i.e., methods that leverage trained models to extract knowledge about tasks, are more closely related to our proposed method. [10] estimated task transferability via the negative conditional entropy between their training label sequences, assuming the two tasks share the same input data examples. The proposed method assumed the existence of an optimal trained source model but did not actually use a trained model in the estimation. [24] removed the input-sharing assumption. Both works assumed that the source and the target tasks use the cross-entropy loss. Our method works, on the other hand, asserts mild assumptions on the choice of loss function only. Taskonomy [22] estimated the transferability of a pair of tasks via the transfer performance from the tasks to a third reference task. Task2Vec [23] tuned a single “probe” network on all target tasks and extracted task-characterizing information via the parameters of this network. Both Taskonomy and Task2Vec required training models on target tasks to be able to extract task information. In contrast, our method does not require any training on target tasks.

All of the mentioned model-based methods have their limiting assumptions. The main assumption of our method, required by our optimality guarantee, is that the output module (classifier) that will be tuned on top of the transferred component admits a particular form as described in Theorem IV.1. Essentially, it is required to be a single-layer NN. It is worth noting that this set-up is common in practice. We summarize the properties of these related methods in Table I.

In this section, we attempt to answer the important question: Do end-to-end and the proposed modular training, when restricted to the architecturally identical hidden modules, drive the underlying modules to functions that are the same in terms of minimizing the overall loss function? Clearly, a positive answer would verify empirically the optimality of the modular approach.

We now test this hypothesis with toy data. The set up is as follows. We generate 1000 32-dimensional random input and assign them random labels from 10 classes. The underlying network consists two modules. The input module is a $(32\to 512)$ fully-connected (fc) layer followed by ReLU nonlinearity and another $(512\to 2)$ fc layer. The output module is a $(2\to 10)$ fc layer. The two modules are linked by a tanh nonlinearity (output normalized to unit vector). The overall loss is the cross-entropy loss. And for modular training, the proxy objective is CTS-NEO. We visualize the activations from the tanh nonlinearity as an indicator of the behavior of the hidden layers under training.

From Fig. 4, we see that the underlying input modules are indeed driven to the same functions (when restricted to the training data and in terms of minimizing the overall loss) by both modular and end-to-end in the limit of training time going to infinity. This confirms the optimality of our proposed method (albeit in a simplified set-up) and allows us to modularize learning with confidence.

We now extend the verification of optimality from the simplified set-up in the previous section to a more practical setting. Recall that we have characterized our proxy objective as function of the input module whose maximizers constitute parts of the overall loss minimizers and we have proposed to train the input module to maximize the proxy objective.

Ideally, however, a proxy objective for input module should satisfy: The overall accuracy is a function of solely the proxy value and the output module and as a function of the proxy value, the overall accuracy is strictly increasing.³³3This type of results is reminiscent of the “ideal bounds” for representation learning sought for in, e.g., [55]. We now demonstrate empirically that the proxies we proposed enjoy this property.

The set-up is as follows. We train a LeNet-5 as two modules on MNIST with $(\text{conv1}\to\text{tanh}\to\text{max pool}\to\text{conv2}\to\text{tanh}\to\text{max pool}\to\text{fc1}\to\text{tanh}\to\text{fc2})$ as the input module, and $(\text{tanh}\to\text{fc3})$ as the output one. The input module is trained with a number of different epochs so as to achieve different proxy values. And for each epoch, we freeze the input module and train output module to minimize the overall cross-entropy until it converges. Mathematically, suppose we froze input module at $F_{1}^{\prime}$ and we denote the proxy as $L_{1}$ and overall accuracy as $A$, we are visualizing $\max_{F_{2}}A(F_{2}\circ F_{1}^{\prime})$ vs. $L_{1}(F_{1}^{\prime})$, where $F_{2}$ is the output module.

Fig. 5 shows that the overall accuracy, as a function of the proxy value, is indeed approximately increasing. Further, this positive correlation becomes near perfect in high-accuracy regime, which agrees with our theoretical results. To summarize, we have extended our theoretical guarantees and empirically verified that maximizing the proposed proxy objectives effectively learns input modules that are optimal in terms of maximizing the overall accuracy, rendering end-to-end training unnecessary.

We now present the main results on MNIST and CIFAR-10 demonstrating the effectiveness of our modular learning method. To facilitate fair comparisons, end-to-end and modular training operate on the same backbone network. For all results, we used stochastic gradient descent as the optimizer with batch size $128$. For each module in the modular method as well as the end-to-end baseline, we trained with annealing learning rates ($0.1,0.01,0.001$, each for $200$ epochs). The momentum was set to $0.9$ throughout. For data preprocessing, we used simple mean subtraction followed by division by standard deviation. On CIFAR-10, we used the standard data augmentation pipeline from [9], i.e., random flipping and clipping with the exact parameters specified therein. In modular training, the models were trained as two modules, with the output layer alone as the output module and everything else as the input module as specified in Sec. IV-C1. We normalized the activation vector right before the output layer of each model to a unit vector such that the equal-norm condition required by Theorem IV.1 is satisfied. We did not observe a significant performance difference after this normalization.

From Table II and III, we see that on three different backbone networks and both MNIST and CIFAR-10, our modular learning compares favorably against end-to-end. Since under two-module modular training, the ResNets can be viewed as KMs with adaptive kernels, we also compared performance with other NN-inspired kernel methods in the literature. In Table III, the ResNet KMs clearly outperformed other kernel methods by a considerable margin.

Our modular training needs only implicit labels (weak supervision) for training the latent modules. Full supervision is only needed for the output module. And we argue that in fact, the modular training of output module requires a smaller number of fully-labeled data than end-to-end training.

Intuitively, the input module in a deep architecture can be understood as learning a new representation of the given data with which the output module’s classification task is simplified. A simple way to quantify the difficulty of a learning task is through its sample complexity, which, when put in simple terms, refers to the number of labeled examples needed for a given model and a training algorithm to achieve a certain level of test accuracy [11].

In a modular training setting, one can decouple the training of input and output modules and then it would make sense to discuss the sample complexity of the two modules individually. The output modules should require fewer labeled data examples to train since its task has been simplified when the input one has been well-trained. We now observe that this is indeed the case.

We trained ResNet-18 on CIFAR-10 with end-to-end and our modular approach. The input module was trained with the full training set in the modular method, but again, this only requires implicit pairwise labels. We now compare the need for fully-labeled data between the two training methods.

With the full training set of $50000$ labeled examples, the modular and end-to-end model achieved $94.93\%$ and $94.91\%$ test accuracy, respectively. From Fig. 6, we see that while the end-to-end model struggled in the label-scarce regime, achieving barely over $20\%$ accuracy with $50$ fully-labeled examples, the modular model consistently achieved strong performance, achieving $94.88\%$ accuracy with $30$ randomly chosen fully-labeled examples for training. In fact, if we ensure that there is at least one example from each class in the training data, the modular approach needed as few as $10$ randomly chosen examples to achieve $94.88\%$ accuracy, that is, a single randomly chosen example per class.⁴⁴4The fact that the modular model underperformed at $10$ and $20$ labels is likely caused by that there were some classes missing in the training data that we randomly selected. Specifically, $4$ classes were missing in the $10$-example training data, and $2$ in the $20$-example data.

These observations suggest that our modular training method for classification can almost completely rely on implicit pairwise labels and still produce highly performant models, which suggests new paradigms for obtaining labeled data that can potentially be less costly than the existing ones. This also indicates that the current form of strong supervision, i.e., full labels on individual data examples, relied on by existing end-to-end training, is not efficient enough. In particular, we have shown that by leveraging pairwise information in the RKHS, implicit labels in the form of whether pairs of examples are from the same class are just as informative. This may have further implications for un/semi-supervised learning.

To empirically verify the effectiveness of our transferability estimation method, we created a set of tasks by selecting $6$ classes from CIFAR-10 and grouping each pair into a single binary classification task. The classes selected are: cat, automobile, dog, horse, truck, deer. Each new task is named by concatenating the first two letters from the classes involved, e.g., cado refers to the task of classifying dogs from cats.

We trained one ResNet-18 using the proposed modular method for each source task, using AL as the proxy objective. For each model, the output linear layer plus the nonlinearity preceding it is the output module, and everything else belongs to the input module. All networks achieved on average $98.7\%$ test accuracy on the task they were trained on⁶⁶6Highest accuracy was $99.95\%$, achieved on audo. Lowest was $93.05\%$, on cado., suggesting that each frozen input module possesses information essential to the source task. Therefore, the question of quantifying input module reusability is essentially the same as describing the transferability between each pair of tasks.

The true transferability between a source and a target task in the task space can be quantified by the test performance of the optimal output module trained for the target task on top of a frozen input module from a source task [10]. Our estimation of transferability is based purely on proxy objective: A frozen input module achieving higher proxy objective value on a target task means that the source and target are more transferable. This estimation is performed using a randomly selected subset of the target task training data.

In Fig. 7, we visualize the target space structure with trdo being the source and the target, respectively. We see from the figures that our method accurately described the target space structure despite its simplicity, using only $10\%$ of target task training data. The discovered transferability relationships also align well with common sense. For example, trdo is more transferable to, e.g., detr, since they both contain the truck class. trdo also transfers well with auca because features useful for distinguishing trucks from dogs are likely also helpful for classifying automobiles from cats.