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 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).