Associated Learning: an Alternative to End-to-End Backpropagation that Works on CNN, RNN, and Transformer

1 INTRODUCTION . 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 . 2 ASSOCIATED LEARNING . 2.1 AN OVERVIEW OF AL : AL FORM , TRAINING , AND INFERENCE . 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 2.1.1 APPLYING AL TO DIFFERENT NEURAL NETWORK STRUCTURES . 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 . 2.2 PROPERTIES OF AL . This section presents three properties of AL that do not exist in BP : forward shortcuts , dynamic layer accumulation , and pipelines . 2.2.1 FORWARD SHORTCUTS . 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 .

Summary: 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.