Year: 2022

Abstract:
We study associated learning (AL), an alternative methodology to end-to-end backpropagation (BP). We introduce the workflow to convert a neural network into an AL-form network such that AL can be used to learn parameters for various types of neural networks. We compare AL and BP on some of the most successful neural networks: convolutional neural networks, recurrent neural networks, and Transformers. Experimental results show that AL consistently outperforms BP on open datasets. We discuss possible reasons for AL’s success, its limitations, and AL’s newly discovered properties. Our implementation is available at https://github.com/Hibb-bb/AL.

heading: 1 INTRODUCTION
text: Backpropagation (BP) is the keystone of modern deep learning. Although BP is the standard way to learn network parameters, it is far from ideal. Some of the most discussed issues of BP are optimization difficulties (e.g., vanishing gradients and exploding gradients (Hochreiter et al., 2001)) and training performance (e.g., backward locking (Jaderberg et al., 2017)).
It appears that custom network structures may be needed for different types of learning tasks. Among the various forms, convolutional neural networks (CNNs), recurrent neural networks (RNNs), and Transformer networks (along with their extensions, e.g., LSTM (Hochreiter & Schmidhuber, 1997) and VGG (Simonyan & Zisserman, 2015)) are particularly useful. These networks have been successfully applied in fields as varied as computer vision, natural language processing, signal processing, and others (Goodfellow et al., 2016; Deng & Yu, 2014).
This paper studies a new learning approach—associated learning (AL)—an alternative to end-toend backpropagation learning. AL decomposes BP’s global end-to-end training strategy into several small local optimization targets such that each layer has an isolated gradient flow. However, since most layers in AL do not interact with the final loss directly, we would expect AL-training models to be less accurate than BP. Surprisingly, the original AL paper compares AL and BP using the CNN network (and its extensions, e.g., VGG) and shows impressive results based on image classification datasets (MNIST, CIFAR-10, and CIFAR-100) (Kao & Chen, 2021). We continue this line of study in two ways. First, we discover more interesting properties of AL. Second, we show how to apply AL on different network structures, including VGG (Simonyan & Zisserman, 2015), LSTM (Hochreiter & Schmidhuber, 1997), and Transformer (Vaswani et al., 2017). Eventually, we compare the networks learned via AL and via BP on various tasks (image classification, sentiment analysis, and topic classification) based on public datasets (CIFAR-10, CIFAR-100, IMDB movie reviews, AG’s News, Stanford Sentiment Treebank (SST), and DBpedia). We find that AL consistently outperforms BP in most datasets. Additionally, AL requires fewer epochs than BP when using early stopping but still yields excellent accuracy. These results suggest that AL is a strong alternative to BP, as AL is effective for various tasks and various network structures.
The rest of the paper is organized as follows. In Section 2, we introduce associated learning and its properties. Section 3 presents the experiments, including a comparison of the model accuracies and
convergence speed of AL and BP on VGG, LSTM, and Transformer, a generalization test of AL and BP, and an ablation study. Section 4 reviews previous work on backpropagation alternatives and the network structures that may look similar to AL. Finally, we conclude our contribution in Section 5.

heading: 2 ASSOCIATED LEARNING

heading: 2.1 AN OVERVIEW OF AL: AL FORM, TRAINING, AND INFERENCE
text: To apply AL, we must transform the network into a different structure, which we call AL form below. Instead of learning a function to map a feature vector x to a target y, the AL form performs metric learning by searching for functions to transform x and y into latent representations such that the distance between these two latent representations is close.
Algorithm 1 shows the procedure to convert a neural work into an AL form. We use Figure 1 (a neural network with 3 hidden layers) as an example to show the process. We can regard functions f1, f2, f3 as the encoders to convert the input feature vector into a latent representation and clf as a classifier to transform the latent representation into a target. When converting this network to its AL form (referring to Figure 2), we only keep the encoders fi-s and extend the model architecture by adding a bridge function bi and an autoencoder for each layer i. Let si−1 and ti−1 be the inputs of layer i (assuming s0 := x and t0 := y), functions fi and gi convert si−1 and ti−1 into latent representations si and ti. The function fi can be any type of forwarding block, such as a convolution layer, an LSTM layer, or a residual block in ResNet. Each bridge function bi projects si (the latent representation of x at layer i) to be close to ti (the latent representation of y at layer i), i.e., bi(fi(si−1)) ≈ gi(ti−1). Meanwhile, gi not only extracts the latent representation of y, but also serves as the encoder for an autoencoder such that hi(gi(ti−1)) ≈ ti−1 (hi is the decoder). For all the AL-form networks in our experiments, each of the functions gi and hi is constructed by a linear transform followed by a non-linear activation function (e.g., tanh), and bi is a linear transformation
of the input vector or input matrix. We discuss the number of neurons for the input, output, and hidden layers of the autoencoders in Section 6.9.
To summarize, we minimize the following objective for each layer i during training:
LiAL = LiA + LiAE = ∥∥bi(fi(si−1))− gi(ti−1)∥∥2 +∥∥hi(gi(ti−1))− ti−1∥∥2 , (3)
where LiA denotes the associated loss of layer i, and LiAE indicates the autoencoder loss of layer i. Algorithm 2 gives the training algorithm for AL. Since each layer has its objective function (Equation 3), we can parallelly update the parameters of different layers using a pipeline to increase the training throughput (details in Section 2.2.3).
For inference, if an AL network has L layers, an input x goes through f1, . . . , fL to generate x’s latent representation (sL), which is transformed to y’s latent representation (tL) via bL. Next, tL is converted to y via hL, . . . , h1. The autoencoders’ encoding functions gi-s are not used during inference. Take Figure 2 as an example, the inference path is (f1 −→ f2 −→ f3 −→ b3 −→ h3 −→ h2 −→ h1). Although it could be unintuitive to concatenate two unconnected transformation functions as an inference path (i.e., hi+1 followed by hi, i = 1, 2, 3), it is valid because hi+1(ti+1) = t′i ≈ ti.
Algorithm 1 Converting a neural network to its AL form Input a neural network N = (f1, f2, ...fL), target y ∈ R1, feature vector x ∈ Rm Output an AL network A
1: s0 ← x; t0 ← y 2: for i = 1 to L do 3: si = fi(si−1) 4: insert a function gi, where ti = gi(ti−1) 5: insert a function bi, where s′i = bi(si) ▷ adding bridge function 6: insert a function hi, where t′i−1 = hi(ti) ▷ adding a decoder for an autoencoder 7: LiAL = MSE(s′i, ti) +MSE(ti−1, t′i−1) ▷ associated loss and autoencoder loss 8: return A = (a1, . . . , aL), ai = {fi, bi, gi, hi}Li=1
Algorithm 2 Training an AL network Input an AL network A = (a1, . . . , aL), ai = {fi, bi, gi, hi}Li=1, training features and targets (X,Y ) Output a fine-tuned AL network A
1: repeat 2: Sample x, y from X,Y . 3: s0 ← x; t0 ← y 4: for i = 1 to L do 5: si ← fi(si−1) 6: ti ← gi(ti−1) 7: s′i ← bi(si) 8: t′i−1 ← hi(ti) 9: LiA ←MSE(s′i, ti) ▷ associated loss 10: LiAE ←MSE(ti−1, t′i−1) ▷ autoencoder loss 11: Update the parameter of fi, bi, gi according to∇LA 12: Update the parameter of gi, hi according to∇LAE 13: until converges

heading: 2.1.1 APPLYING AL TO DIFFERENT NEURAL NETWORK STRUCTURES
text: This section introduces details on converting some of the most successful neural network structures into their corresponding AL forms. We discuss vanilla CNN and VGG (used as the representative models for CNN-based models), LSTM (used as the representative model for RNN-based models), and Transformer. We also discuss how to integrate word embeddings into an AL-form network.
CNN The AL form of an CNN architecture (here, vanilla CNN and VGG) uses fi-s to convert an input image x into latent representations by convolutions, just like a regular CNN or a VGG does. The associated loss at layer i is defined as the distance between bi(si) (transforming the flattened feature map at layer i into ti’s shape) and ti (the latent representation of y).
RNN The AL form of an RNN architecture (here, LSTM and Bi-LSTM) uses the internal state to iteratively process each element in the input sequence x and generates the new internal state, just like a regular RNN does. After reading the entire input sequence x, we define si the latent representation of x at layer i by the final internal state. Consequently, the associated loss is defined as the distance between bi(si) (transforming the final internal state at layer i into ti’s shape) and ti (the latent representation of y). Details are shown in Section 6.2.
Transformer The AL form of a Transformer encodes the input sequence x into a list of vectors, as a regular Transformer does. We define si (the latent representation of x) by computing the mean-pooling on the encoded vectors. The associated loss is defined as the distance between bi(si) (transforming the mean-pooling output into ti’s shape) and ti (the latent representations of y). Details are shown in Section 6.2.
Word Embedding Word embeddings are frequently used as the input of LSTM or Transformer for NLP tasks. We use mean-pooling to aggregate all token’s word embeddings as si.

heading: 2.2 PROPERTIES OF AL
text: This section presents three properties of AL that do not exist in BP: forward shortcuts, dynamic layer accumulation, and pipelines.

heading: 2.2.1 FORWARD SHORTCUTS
text: Forward shortcuts enable faster inference. However, we can also leverage “shortcut paths” for faster inference. As shown in Algorithm 3, given an integer ℓ (1 ≤ ℓ ≤ L), the bridge function bℓ can serve as a shortcut to transform sℓ to s′ℓ, which should be close to tℓ when the network is well trained. As a result, we can skip fj , bj , and hj for all j > ℓ to reduce the length of the inference function. In other words, we have multiple inference functions based on a single model. When inference time is
critical, we can select a shorter inference path, e.g., x f1−→ s1 b1−→ t1 h1−→ y. On the other hand, if inference time is unimportant, we can dynamically adjust the model complexity by modifying the number of AL layers used at the inference phase to reduce overfitting or underfitting.

heading: 2.2.2 DYNAMIC LAYER ACCUMULATION
text: An AL-form network allows the dynamic creation of new layers during training. Specifically, we initially create an AL-form network with k AL layers and train the network based on the strategy introduced in Section 2.1. If the model still underfits, we simply add the (k + 1)-th AL layer on top of the first k AL layers. We may choose to fix the parameters in the first k AL layers and train only the parameters in the (k+1)-th AL layer. In contrast, when using a standard neural network, adding layers dynamically is more complicated.

heading: 2.2.3 PIPELINE
text: In an AL-form network, the network parameters in different AL layers can be trained simultaneously via pipelines because each AL layer i has a local objective function. In contrast, it is difficult to simultaneously update the parameters of different layers for a standard BP-based neural network.
Given n training instances in a standard neural network with L hidden layers, if we apply BP for training, the time complexity is O(nL) because each instance must complete the entire forward (L+1 transformations) and backward (L+1 transformations) process. Thus, for such an architecture, the time complexity increases linearly with the total number of layers L. However, since each AL layer has its local objective, each layer can update its parameters without waiting for the gradient from preceding layers. As a result, for an AL-form network with L AL layers, using L computation units can make the parameter update process fully pipelined, as explained below.
We denote task ji as the task of computing both forward and backward propagation for an AL layer i based on the j-th instance. If an AL layer i finishes task1i at time step t, at the next time step t+ 1, AL layers i can continue performing task2i since all parameters in this layer are updated based on task1i . Meanwhile, AL layer i + 1 performs task 1 i+1. Consequently, training based on the first instance requires O(L), and each of the following n − 1 instances requires extra O(1) time units. The time complexity of one training epoch eventually reduces to O(n+ L).

heading: 3 EXPERIMENTS
text: To evaluate how AL works on different network architectures, we conducted experiments on image classification with vanilla CNN and VGG and on sentiment and topic classification with LSTM and Transformer (using pre-trained word embeddings as input). In addition, we conducted a case study on AL’s generalization ability, evaluated by fitting AL models and BP models on different portions of noisy labels. Finally, we discuss possible reasons for AL’s success based on ablation studies. We report more detailed settings in the Appendix.

heading: 3.1 DATASETS
text: We used four datasets (IMDB Review, AG’s News corpus, DBpedia Ontology, and the Stanford Sentiment Treebank) to compare AL and BP on LSTM and Transformer, and two datasets (CIFAR-10 and Fashion-MNIST) to compare AL and BP on CNN and VGG. Dataset details are in Appendix 6.4

heading: 3.2 TEXT CLASSIFICATION (LSTM AND TRANSFORMER)
text: We compared BP and AL on LSTM and Transformer based on the text classification task.
For LSTM, we used a 2-layer bidirectional LSTM with pre-trained GloVe word embeddings (Pennington et al., 2014) as input. For Transformer, we used a 2-layer Transformer encoder, again with pre-trained GloVe word embeddings. During training, we recorded the accuracies of the validation set for every epoch and saved that model with the best performance on the validation set for testing. For each model, we repeated the process five times and reported the average accuracy and standard deviation. For a fair comparison, we forced the AL parameter counts to not exceed those of the compared standard neural network by reducing AL’s hidden dimension size.
We report the accuracies of LSTM and Transformer trained via BP and AL in Table 1. We also report the accuracies based on the predictions using shorter inference paths in the same Table. Although AL is trained by isolated gradient flows, AL-form LSTMs outperform their standard BP-trained counterparts on most datasets. Transformer-AL also outperforms Transformer on every dataset. Also, we observe that standard Transformer easily overfits smaller datasets like SST and IMDB, resulting in poor test accuracies, but Transformer-AL does not seem to have this weakness.
In Table 2 and Table 3, we report the epoch that yielded the best results. We discover that LSTMAL converges much faster than BP, suggesting that training time will be improved if applying early
stopping. For the Transformer-AL, although the best validation and test accuracy come at roughly the 20th epoch on AGNews, Transformer-AL reaches high validation and test accuracy within the first 5 epochs. Since we report only the epoch that reaches the best validation and test accuracies, the required epochs look large. Particularly, Transformer-AL achieved 99% of the best validation set accuracy before the 5th epoch. More details can be found in Appendix 6.7.

heading: 3.3 IMAGE CLASSIFICATION (VANILLA CNN AND VGG)
text: We also reproduced the experiments in Kao & Chen (2021) to determine whether AL-form CNNs also converge faster. In image classification, we report performance on the CIFAR-10 and FashionMNIST benchmarks. We selected FashionMNIST instead of MNIST because MNIST is likely too simple to show the power of different models: CIFAR-10 is more challenging than FashionMNIST since it contains three channels, whereas FashionMNIST is limited to grayscale images. Table 4 shows the results: all AL models converge faster than BP. Figure 6 in Appendix 6.6 shows the relationship between epochs and test accuracies on CNN and CNN-AL. Also, AL models perform slightly better in terms of accuracy, consistent with the results demonstrated in Kao & Chen (2021).

heading: 3.4 GENERALIZATION
text: To assess the generalizability of a model, a common approach is to empirically compare the errors on the training data and test data. Garg et al. (2021) show another way to evaluate generalizability: a model is prone to overfitting if it closely fits a dataset with random labels. We use both methods to show that an AL-trained model may be more general than that trained by BP.
We took θ proportion of data, in which we changed the labels randomly. We used different proportions to train models. We report both the training and test accuracies on clean test data.
Table 5 shows the results. First, as shown in the last column (for a label noise rate of θ = 0), the training accuracy of the CNN is larger than that of the CNN-AL, but CNN-AL outperforms CNN
in terms of test accuracy. Second, with θ > 0, the vanilla CNN network with BP training fits the randomly labeled training data much better than the AL-form CNN. However, if we evaluate model accuracies based on the clean test data, the AL-form CNN performs better. Taken together, the above empirical evidence suggests that the new model may generalize better (Garg et al., 2021). The relationship between accuracy and training epochs (for θ = 1.0, 0.2, and 0.6) is given in Figure 3.

heading: 3.5 WHAT MAKES AL WORK – AN ABLATION STUDY
text: We empirically validates that AL generates great models, converges fast, and generalizes well. Below we report ablation studies using AG’s news dataset to understand why AL performs excellently.
First, we ensure the gradient flow in each AL layer does not influence the gradients in other layers via the gradient stopping mechanism.1 However, to optimize the objective function (Equation 3), either allowing or blocking the gradient flow between AL layers is acceptable. We sought to determine whether allowing gradient flow between AL layers improves model accuracy (although such a design may make pipelining infeasible). Hence, we disabled all the gradient stops (allowing gradients to flow through different AL layers). The results are denoted as “LSTM-AL w/o grad. stops” in Table 6: accuracy is nearly unchanged. Therefore, given enough computing nodes, applying a pipelined AL increases training throughput without sacrificing model accuracy.
Second, to test the effect of associated loss, we disabled the gradient stops and removed the associated loss (except the last AL layer) from the model. The second row of Table 6 shows the result: although the accuracy is similar to that with associated loss, far more epochs are needed to converge. So, the bridge functions likely facilitate model training by gradually bringing s′i and ti closer.
With the third experiment, we sought to investigate the effect of the autoencoder components in the AL model. To do this, we passed the output of a normal Bi-LSTM to a bridge function, whose output t1 was fed to the bottleneck layer of an autoencoder, which decodes t1 back to the target y. Compared with the original Bi-LSTM classifier, the new design results in higher accuracy, as seen in the third (LSTM + AE*1) in Table 6. This suggests that replacing the standard classifier with a decoder of an AE and training it with AL objectives could produce a better model. Additionally, we experimented with removing all autoencoders and using bi as the classifier to predict y for each AL layer (similar to the early exit technique in Teerapittayanon et al. (2016)). We use “LSTM-AL w/o AE (SC-1)” and “LSTM-AL w/o AE (full)” in Table 6 to denote “early exit at the first layer” and “using bL the last bridge as the classifier”, respectively. The results show that using autoencoders
1This is supported by Tensor.detach() in PyTorch and tf.stop_gradient() in TensorFlow.
gives a better predicting model. We suspect that each autoencoder perhaps performs some kind of feature extraction and regularization. Particularly, although it is easy to transform one-hot encoded y to t1 to decrease associated loss and transform t1 to t′0 to decrease the autoencoder loss, the network still needs to learn more challenging tasks: converting ti to ti+1 and ti+1 to t′i for i > 1 to reduce the associated loss and autoencoder loss of layer i. Consequently, adding autoencoders may prevent overfitting for the components beyond layer 1.
Finally, for LSTM-AL, the inference path goes through 4 transformation functions (b3 −→ h3 −→ h2 −→ h1) after the LSTM module. To ensure a fair comparison, we also tested a network with a 3-layer MLP classifier after the standard LSTM model. The result is denoted as LSTM + MLP*3 in Table 6. As shown, replacing the MLP from one layer to three layers does not produce a better model, i.e., the number of layers is not why LSTM-AL is a better model.

heading: 4 RELATED WORK

heading: 4.1 BACKPROPAGATION ALTERNATIVES
text: Backpropagation is fundamental in deep learning. However, backpropagation suffers from optimization and performance issues. Related work usually concerns methodologies that better imitate the signal transmission process of biological neural networks, since biological neural networks are highly efficient in analyzing visual, audio, textual, and other signals (Lee et al., 2015; Kao & Chen, 2021; Nøkland, 2016; Ororbia et al., 2018; Ororbia & Mali, 2019). Alternatives to end-to-end backpropagation can be categorized into three types: proxy objective, target propagation, and synthetic gradients (Duan & Principe, 2021). Below we introduce each kind and the representative works.
The first type—proxy objective—uses a local objective function for one layer i to learn θi, the parameters for the current layer i. The model then fixes θi, adds one more layer i+ 1 with another objective function, and learns θi+1, the parameters for the new layer. Consequently, proxy objective decouples the end-to-end backpropagation since the parameters in each layer are trained via a local objective function (Nøkland & Eidnes, 2019; Belilovsky et al., 2019; 2020; Löwe et al., 2019).
The second type—target propagation—approximates gi, the inverse of the downstream layers for every layer i. To update θi, the parameters of layer i, the model backpropagates the target (not gradient) through function gi. Consequently, target propagation handles the situation where the relationship between the parameters and the cost is highly non-linear (e.g., discrete), which is otherwise difficult to solve using backpropagation (Lee et al., 2015; Meulemans et al., 2020).
The third type is synthetic gradients, which learns a function to approximate the gradient for each layer and utilizes the approximated gradients as the true gradients to update the parameters. This method thus decouples the layer-wise parameters (Jaderberg et al., 2017; Czarnecki et al., 2017; Lansdell et al., 2019).
Associated learning is different from these methods in several ways. First, both proxy gradient and target propagation propagate the signal (loss or target) from the output layer directly to each layer. Biologically, this could be implausible because when in deep networks, hidden layers near the inputs are unlikely to receive direct signals from the output layer that are far away. Associated learning, on the other hand, does not require each component to receive signals directly from the output layer. Additionally, although proxy gradient allows each layer to have a local objective, most models of
this type still learn the parameters in a layer-wise fashion, so it may still be challenging to learn the parameters using a pipeline. One exception is Greedy InfoMax (GIM) (Löwe et al., 2019), which does not require a direct signal from the output layer. However, GIM uses contrastive loss as the loss function, so GIM can mainly be applied to self-supervised learning tasks. As a result, for supervised learning tasks, AL is likely a more natural choice than GIM. As for synthetic gradients, empirical experiments show that gradients are difficult to predict, so model training based on synthetic gradients usually yields much worse accuracies than those for models trained via backpropagation.

heading: 4.2 NEURAL NETWORKS WITH STRUCTURES SIMILAR TO AL
text: The structure of an AL-form model may look similar to a ResNet, as the bridge functions in AL appear analogous to the residual connections in ResNet. Therefore, some may argue that AL’s superiority is the result of its resemblance to ResNet. However, the residual connections in ResNet are implemented to mitigate accuracy saturation (He et al., 2016), and all the parameters are adjusted to fit a global objective function, whereas the bridge functions in AL are implemented to modularize the network such that the gradient flow in one AL layer does not pass to other layers. Therefore, the bridge functions and the residual connections are only structurally similar; their purposes and functionalities are entirely different.
An AL-form network may also appear similar to contrastive learning (CL). Particularly, CL transforms a pair of images i and j (along with their augmented images i′ and j′) into vectors with large distances between (1) images i and j and (2) augmented images i′ and j′, with small distances between (1) an image i and its augmented image i′ and (2) an image j and its augmented image j′. This process is related to AL because both models convert inputs to vectors and compare the distance between the transformed vectors. Following this line, we are interested in studying the relationship between AL and CL in two ways. First, we could apply AL to CL such that CL is trained more effectively. Second, we could extend the AL model such that the inputs include a pair of samples x1 and x2 with different labels and design an objective function that maximizes the distance between the vectors transformed from x1 and x2 (since their labels are different) and simultaneously minimizes the current AL loss function (Equation 3).
An AL network is also structurally similar to the models in (Lin et al., 2014; Yeh et al., 2017), which also transform x into a latent space and convert it to y via an autoencoder. However, the design motivations are very different from AL. The two works attempt to efficiently represent a multi-label target y by a low-dimensional latent representation c, so learning a function to convert x to c may be easier. The motivation of AL is to design isolated loss functions so that the new design has favorable properties that BP does not have, such as short gradient flows, forward shortcuts, dynamic layer accumulation, and layer-wise pipeline learning. Isolated objectives are not implemented in (Lin et al., 2014; Yeh et al., 2017).
Ladder network is another structurally similar network (Rasmus et al., 2015), which adds skip connections between each encoder to its corresponding decoder. While the ladder network may look similar to AL structurally, at least two manifest differences exist. First, the residual between y and ŷ in the ladder network is propagated through all layers to update the parameters on the inference path, but in AL, each layer has a local objective function, and most layers do not receive signals from the output layer. Second, the skip connection in the ladder network connects an encoder layer to a decoder layer. In contrast, the bridge function in AL connects a hidden representation of x to the bottleneck layer of an autoencoder to create the local loss.

heading: 5 CONCLUSION
text: This paper studies associated learning, an alternative methodology to end-to-end backpropagation. Since backpropagation is the cornerstone of today’s deep learning but still suffers from optimization and training performance issues, it is appropriate to study alternatives to backpropagation. We discussed the unique properties of AL that BP does not have. AL provides multiple inference paths; it allows dynamic layer accumulation during training; it supports pipeline training, and, perhaps most importantly, its prediction power is comparable to and frequently better than backpropagation on many of the most successful neural network structures. Ablation studies disclose why AL may work and demonstrate AL’s effectiveness and architecture flexibility in training and testing.

heading: 6 APPENDIX

heading: 6.1 AL INFERENCE ALGORITHM
text: Algorithm 3 shows the inference algorithm for AL networks. It supports forward shortcuts by going through only ℓ of the L AL layers.
From Algorithm 3, it is clear that the functions gi-s are not used in inference, and only one bridge function out of the L bi-s are used in inference. Therefore, the number of parameters of inference is smaller than the number of training parameters.
Algorithm 3 Using AL form network for inference Input AL network A = (a1, a2, ...aL), ai = {fi, bi, gi, hi}Li=1, input data x, inference length ℓ, 1 ≤ ℓ ≤ L Output predicted label
1: s← x 2: for i = 1 to ℓ do 3: s← fi(s) 4: t← bℓ(s) 5: t′ ← t ▷ Same as Figure 2 6: for i = ℓ to 1 do 7: t′ ← hi(t′) 8: return t′

heading: 6.2 CNN, RNN, AND TRANSFORMER IN THEIR AL FORMS
text: Here we extend 2.1.1 for further explanation. In AL, the process in bridge layers did not influence the actual forward path in AL. Therefore, the process in 2.1.1 is only used to calculate associated loss and will not effect the original representation shape. Figure 4 and Figure 5 show that when we take the last hidden state in the sequence input to calculate associated loss, the original sequence representation remains the same when propagate to the next layer.

heading: 6.3 IMPLEMENTATION ENVIRONMENT DETAILS
text: For LSTM and Transformer, we used the PyTorch package (Paszke et al., 2019). For CNN, we used the Tensorflow Abadi et al. (2015) package to build the AL models (vanilla CNN and a visual geometry group (VGG) network). We trained the models with GeForce 3090 GPUs on Ubuntu Linux 20.04. Since current hardware development does not support actual pipeline training, we simulated pipeline training by blocking gradient flows between AL layers with Tensor.detach() in PyTorch and with tf.stop_gradient() in TensorFlow.

heading: 6.4 DATASET STATISTICS
text: The IMDB Review dataset (Maas et al., 2011) is a binary sentiment classification dataset containing 25k training data and 25k test data. AG’s News corpus dataset (Zhang et al., 2015) includes four news topics: world, sports, business, and sci/tech. It has 120k training data and 7.6k test data. The DBpedia Ontology classification dataset (Zhang et al., 2016) was constructed by picking 14 non-overlapping topic classes from DBpedia 2014, containing 560k training samples and 70k test samples. The Stanford Sentiment Treebank (SST2) dataset (Socher et al., 2013) is a binary sentiment classification dataset from GLUE (Wang et al., 2018). The dataset contains 67k training samples, 872 validation samples, and 1.8k test samples. The CIFAR-10 dataset contains color images from 10 classes; 50k are training samples, and 10k are test samples. Finally, the Fashion-MNIST dataset (Xiao et al., 2017) is a grayscale image classification dataset with 10 classes. It consists of 60k training samples and 10k test samples.

heading: 6.5 TEXT CLASSIFICATION
text: During training, the most frequent 30k words in the training corpus were included in the word embedding. We replaced the remaining words with the [UNK] token. We trained all text classification experiments with 50 epochs. We used the Adam optimizer.

heading: 6.5.1 BI-LSTM
text: For Bi-LSTM, we applied gradient clipping with the value 5. For the standard version of Bi-LSTM, we experimented with learning rates of 0.0005 and 0.0001 and reported that with better accuracy. We also adjusted the learning rate for AL, but the learning rate has little influence on the accuracy and the convergence speed. We used a batch size of 64 for all Bi-LSTM experiments. For IMDB and DBpedia, we set the max sequence length to 500. For AG’s News and SST, we set the max sequence length to the longest sequence length; the exact length settings were applied to the Transformer experiments. The objective of the standard Bi-LSTM was a cross-entropy loss. Details are given in Table 7.

heading: 6.5.2 TRANSFORMER
text: For the standard Transformer model, we used negative log-likelihood as the objective with multihead attention (6 heads). We set the batch size to 128 because the standard Transformer converges only with a large batch size in our experiment. Table 8 gives the details.

heading: 6.6 IMAGE CLASSIFICATION
text: We normalized the image pixels for both FashionMNIST and CIFAR-10 datasets as a preprocessing step. Since CIFAR-10 is more challenging, we also performed data augmentation on this dataset by resizing, random cropping, random flipping, and adjusting the brightness of the images. We set the batch size to 128. We set the initial learning rate to 0.0001 and applied scheduled decay until converging. We used Adam as the optimizer and cross-entropy loss for the standard version.
Vanilla CNN Model We implemented the vanilla CNN by modifying the previous code. The CNN contains 13 hidden layers and an output layer. The first 8 hidden layers are convolution layers; the last is flattened to a fully connected layer with 1280 neurons. The subsequent layers consist of 4 feedforward layers, followed by an output layer with 10 classes. We set the initial learning rate to 0.0001 and applied scheduled decay.

heading: 6.7 TRANSFORMER-AL CONVERGENCE SPEED
text: Training an LSTM, a CNN, or a VGG based on AL requires much fewer epochs when compared with training them by BP; this trend is less clear when training a Transformer-AL model. However, we found that after training the Transformer-AL for roughly 5 epochs, the validation set accuracy (90.95) can already achieve 99% of the best validation set accuracy (91.31), and the test accuracy (90.98) is very close to the best test accuracy (91.17, i.e., training Transformer-AL for 20 epochs for AG News). So, we could reduce the number of epochs by sacrificing slight accuracy. TransformerAL on DBpedia has a similar trend: when training Transformer-AL on DBpedia for 5 epochs, its validation set accuracy achieves 96.0902, which is very close to the best validation set accuracy (96.6134). This result shows that Transformer-AL is efficient on convergence when training on DBpedia and AGnews.

heading: 6.8 ASSOCIATED LOSS VS. LAYER DEPTH
text: This section shows that, under the AL network structure, the latent representations in the lower layer appear to help generate the latent representations in an upper layer.
We experimented with the AL layer number i vs. the associated loss ∥si − ti∥22 using different datasets on different network structures. The following tables show that the associated loss at an upper layer (i.e., the layers closer to the output) is smaller than at a lower layer (i.e., the layers closer to the input). Consequently, the shortcuts are likely doing some kind of curriculum learning.

heading: 6.9 THE DESIGN OF THE OVERCOMPLETE AUTOENCODERS IN AL
text: An autoencoder’s hidden layer usually consists of fewer neurons than the input and output layers. However, we use an overcomplete autoencoder (i.e., the hidden layer has more neurons than the input/output layers) (Vincent et al., 2010) in AL. An overcomplete autoencoder can reach zero reconstruction loss (i.e., the autoencoder loss in Equation 3) by copying the neurons from the input layer to the output layer. However, in AL, gi (the encoding function of the autoencoder at layer i) needs to minimize both LiAE the autoencoder loss and LiA the associated loss. As a result, LiA can be regarded as a regularization term when minimizing LiAE . Additionally, we hypothesize that mapping the one-hot encoded y to a high dimensional vector can accelerate the convergence speed. According to Graf et al. (2021), the cross-entropy loss and supervised contrastive loss attain their minimum if and only if the features of each class collapse to the vertices of an origin-centered regular K − 1 simplex, where K is the number of classes. In AL, since the autoencoder in layer-1 simply maps y to a high dimensional space and back (i.e., letting h1(g1(y)) ≈ y), it is nearly guaranteed to reach minimum loss. Therefore, we can consider g1(y) as an ideal prototype for AL to fit in, such that, b1(f1(x)) ≈ g1(y). Therefore, the associated loss directly optimizes for a K − 1 regular simplex in the representation space, which is the best configuration to minimize the supervised contrastive learning loss (Khosla et al., 2020). Such an autoencoder setup probably helps AL find the best configuration in the early stage of training. This may explain why AL converges so fast. Besides, an overparameterized model may accelerate optimization (Arora et al., 2018; Chen & Chen, 2020). Empirical studies in Section 3.5 support our claim: using overcomplete autoencoders to replace classifiers performs better.

heading: 6.10 PYTORCH PSEUDOCODE
text: Here we provide a PyTorch pseudocode of a 5-layer AL network.
import torch import torch.nn as nn
class MLP(nn.Module): def __init__(self, d1, d2):
super().__init__() self.f = nn.Sequential(
nn.Linear(d1,d2), nn.Sigmoid())
def forward(self, x): return self.f(x)
class AL(nn.Module): def __init__(self, class_num):
super().__init__() # we choose 10 as hidden dimension size for simplicity. self.f = nn.ModuleList([MLP(10, 10) for i in range(5)]) self.g = nn.ModuleList([MLP(10, 10) for i in range(5)]) self.b = nn.ModuleList([MLP(10, 10) for i in range(5)]) self.h = nn.ModuleList([MLP(10, 10) for i in range(5)]) self.g[0], self.h[0] = MLP(class_num, 10), MLP(10, class_num)
self.mse, self.ce = nn.MSELoss(), nn.CrossEntropyLoss() self.c = class_num
def forward(self, x, y):
loss = 0.0 s_i, t_i = x, nn.functional.one_hot(y, num_classes=self.c).float() # convert y into one-hot vector for linear input for i in range(5):
last_t_i, s_i = t_i, self.f[i](s_i) # encode input t_i, s_i_prime = self.b[i](s_i), self.g[i](t_i) # bridge + encode target t_i_prime = self.h[i](t_i) # decode target
if i == 0: loss += self.mse(t_i.detach(), s_i_prime) + self.ce(last_t_i, t_i_prime)
else: loss += self.mse(t_i.detach(), s_i_prime) + self.mse(last_t_i, t_i_prime)
s_i, t_i = s_i.detach(), t_i.detach() # cut gradient
return loss
def inference(self, x): for layer in self.f:
x = layer(x) y = self.b[-1](x) for layer in reversed(self.h):
y = layer(y) return y
model = AL(class_num=5)
x, y = torch.rand(30, 10, dtype=torch.float), torch.randint(high=4, size=(30, )) loss, pred = model(x, y), model.inference(x)

Review summaries:
This paper is a continuation of an original associated learning paper by Kao&Chen 2021. It attempts to propoose new learning approach associated learning as an alternative way to back-propagation. On top of the original paper, it discovers more interesting properties and extend AL to CNN, LSTM and transformers (though lacking sufficient details).

This paper proposes associated learning (AL) for CNN, RNN, and transformer. Different from back-propagation (BP), AL decomposes BP’s global end-to-end training strategy into several small local optimization targets such that each sub-networks has an isolated gradient flow. To achieve this, the paper proposes to map input $x$ and output $y$ into intermediate AL layers and performs metric learning (e.g., $t_1=b_1(s_1)$) and auto-encoder learning ($t_1=t_1^{‘}$), as shown in Figure 2. Moreover, Each AL layer can be optimized locally. The idea is interesting. The experiments demonstrate the effectiveness on (IMDB Review, AG’s News corpus, DBpedia Ontology, the Stanford Sentiment Treebank, CIFAR10, and Fashion-MNIST.