MLP-based architecture with variable length input for automatic speech recognition

1 INTRODUCTION . Self-attention , the well-known building block of Transformer ( Vaswani et al. , 2017 ) , appeared in natural language processing ( NLP ) where it caused a breakthrough . Soon enough , it propagated to the fields of computer vision and automatic speech recognition ( ASR ) . In particular , recent end-toend ASR systems based on self-attention architecture , e.g. , Transformer ( Karita et al. , 2019a ) and Conformer ( Gulati et al. , 2020 ) , provide state-of-the-art performance . Recently , several architectures entirely based on MLP have been proposed in the area of computer vision . MLP-based architectures have a simple structure , and achieve performance competitive with that of Transformer-based architectures , despite having fewer parameters and lower computational complexity . They split an image into patches and reshape it into a ( channels× patches ) matrix used as input . An illustration of the process , and its analog for variable length sequences , is shown in Figure 1 . MLP-based architecture such as MLP-Mixer ( Tolstikhin et al. , 2021 ) and gMLP ( Liu et al. , 2021 ) consist of MLP across the channel dimension and MLP across the patch dimension ( also referred to as spatial dimension in computer vision tasks ) . MLP across the channel dimension mixes information between channels like a typical feed-forward network ( FFN ) . MLP across the patches dimension mixes information between patches . All these different works demonstrate that this process capture sufficient information and that self-attention is not always necessary . The small model size and low computational cost of the MLP-based architecture are also useful for ASR . For acoustic data , the input sequence is typically first split into contiguous or overlapping blocks and transformed to some kind of frequency domain representation , e.g . a Mel-spectrogram . Then , the time and frequency dimensions are analogous to patch and channel , respectively , for images . We adopt the terminology of ” token ” to describe the vector of channel values at a given time . It is more apt for sequences and consistent with that used in the MLP-Mixer and gMLP works . Now , in contrast to images , sequences will produce inputs with a variable number of tokens , making it impractical to apply an MLP directly to this dimension . In this paper , we propose three approaches that can work with variable length inputs . First , the input sequences are broken down into contiguous chunks that we call tokens , as explained earlier and shown in Figure 1 . Building on previous approaches , we propose three new token mixing units . Two of them rely on gating , where the input is split into two parts , one is transformed and then multiplied with the other . In the first kind of unit , Convolutional MLP ( C-MLP ) , we use a depthwise convolution to transform the input . This operation is locally MLP-like along the token dimension and makes the network shift invariant . The second approach , Temporal-Shift MLP ( TSMLP ) , concatenates shifted parts of the input . The third approach is the non-gating one . In this case , we apply a simple depthwise convolution , thus mixing tokens with their local neighbors . However , we make use of the fast Fourier transform ( FFT ) to do so efficiently . This approach is termed Fourier MLP ( F-MLP ) . We applied these MLP-based methods to the connectionist temporal classification ( CTC ) ( Graves et al. , 2006 ) ASR model and evaluated them on two datasets , Librispeech ( Panayotov et al. , 2015 ) and Tedlium2 ( Rousseau et al. , 2014 ) . We found different trade-offs between accuracy in terms of word error rate ( WER ) , the number of parameters , and inference speed . However , when matching the number of parameters , the proposed MLP architectures consistently outperform selfattention based models , decreasing the WER by as much as 3 % . Related works . Recently , end-to-end ( E2E ) ASR models , where a single neural network produces a transcription directly from the input acoustic features , have come to dominate benchmarks . The E2E ASR model consists of an encoder that converts audio input into a latent representation and a decoder that converts the latent representation into text output ( token ) . To build the encoder-decoder architecture , attention based encoder-decoder models such as Listen , Attend and Spell ( LAS ) ( Chan et al. , 2015 ) and recurrent neural network transducer ( RNN-T ) ( Graves , 2012 ) have shown high performance , but they suffer from slow inference speed due to the way they auto-regressively output tokens . Another choice to build E2E ASR is non-autoregressive models with the audio encoder and a decoder using CTC ( Graves et al. , 2006 ) . In recent years , many studies have been conducted on CTC-based models , and reported results demonstrate high inference speed for a performance comparable to that of attention-based encoder-decoder models . For the architecture of the audio encoder , recurrent neural networks ( RNNs ) and Convolution neural networks ( CNNs ) have been employed ( Graves & Jaitly , 2014 ; Li et al. , 2019 ; Kriman et al. , 2020 ) . More recently , models applying the self-attention based architecture ( Karita et al. , 2019a ; Zhang et al. , 2020 ; Gulati et al. , 2020 ) have shown promising performance and become the de facto architecture . However , self-attention is expensive in terms of memory and computation , and reduction of both is an important endeavor . Next , we summarize the recently proposed MLP-based architectures . MLP-Mixer ( Tolstikhin et al. , 2021 ) consists of channel-mixing MLPs and token-mixing MLPs . It alternately applies these MLPs and mixes information between channels and information between tokens . gMLP ( Liu et al. , 2021 ) also consists of MLPs across the channel and token dimensions . To more effectively mix tokens , it proposes a spatial gating unit ( SGU ) , a module containing MLP across the token dimension and used to modulate the output of the unit . S2-MLP ( Yu et al. , 2021a ) and S2-MLPv2 ( Yu et al. , 2021b ) mix information across the token dimension by shift operation along the width dimension and height dimension called spatial-shift operation . GFNet ( Rao et al. , 2021 ) efficiently applies learnable filters in the frequency domain . It performs a 2D Fourier transform of the input image , multiplies by the filters , and transforms back to the image domain . Then , MLP is applied across the channels . In the area of NLP , FNet ( Lee-Thorp et al. , 2021 ) has been proposed . It applies a 2D Fourier transform to the input data , retains the real part , and follows with MLP across channels . Contributions . Our contributions in this work are as follows . ( i ) We propose three MLP-based architectures suitable for sequences of arbitrary lengths . ( ii ) We evaluate our architecture for ASR on two different datasets . We show that the proposed C-MLP architecture outperforms the celebrated self-attention , both in accuracy and inference speed . To the best of our knowledge , this is the first time that MLP-based architecture is applied to ASR . 2 CONVENTIONAL MLP-BASED ARCHITECTURES . The so-called MLP-based architectures are built using two modules , each based on simple MLPs : the channel mixing and token mixing modules . The channel mixing module typically consists of linear projections on the channel dimension for each token , similar to the point-wise feed-forward network in the Transformer . The token mixing module mixes the information in the token dimension as a replacement for self-attention , and how to design it is one of the keys to the MLP-based models . We explain conventional architectures dividing them into two types according to how they stack the modules : MLP-mixer and gMLP types , corresponding to the architectures proposed in ( Tolstikhin et al. , 2021 ) and ( Liu et al. , 2021 ) , respectively . We start by describing the outer construction , i.e. , how the channel and token mixing modules are combined . Then , we focus on the inners of a few kinds of token mixing module . Throughout , we omit layer normalization , residual connections , and bias of linear projections for the sake of brevity . Figure 2 shows diagrams of all the parts . 2.1 OUTER CONSTRUCTION . MLP-Mixer type . Figure 2 ( a ) shows the MLP-Mixer type which first applies the token-mixing module and then applies the channel-mixing module . It is adopted by its namesake ( Tolstikhin et al. , 2021 ) , and GFNet ( Rao et al. , 2021 ) . Let Xin ∈ RD×N denote the input data matrix with D channels and N tokens . It first goes through the token mixing module , X = TokenMixingModule ( Xin ) ∈ RD×N , ( 1 ) where TokenMixingModule ( ) is an MLP-based module that usually consists of multiple linear projections on the token dimension . How to design this module is one of the key considerations for MLP-based models . Conventional approaches are introduced in Section 2.2 . Next , the output of the token-mixing module , X , is fed to the channel-mixing module . The channelmixing module is the same as the point-wise feed forward network in Transformer , which performs two linear projection across the channel dimension , Xout = W2σ ( W1X ) ∈ RD×N , ( 2 ) where W1 ∈ RD ′×D and W1 ∈ RD×D ′ are the weights of the linear projections . We remark that the construction of the MLP-Mixer type is similar to that of a Transformer , with the token-mixing replacing self-attention . This highlights the critical importance of token mixing . gMLP type . Figure 2 ( b ) shows the gMLP type where the two layers of the channel-mixing module are split and the token-mixing module placed in-between . This is the structure adopted by gMLP , S2-MLP and S2-MLPv2 . Let Xin ∈ RD×N denote the input data with D channels and N tokens . First , linear projection on channel dimension D to D′ is applied to Xin , X = σ ( W1Xin ) ∈ RD ′×N , ( 3 ) where W1 ∈ RD ′×D refers to weights of linear projection . Next , the token-mixing module follows , X′ = TokenMixingModule ( X ) ∈ RD ′′×N . ( 4 ) Finally , linear projection across the channel dimension D′′ to D is applied as Xout = W3X ′ ∈ RD×N , ( 5 ) where W3 ∈ RD×D ′′ is the weight matrix of the linear projection . Here , additional linear projections can be utilized for mixing the channel information . For example , S2-MLP stacks two additional linear projections after the token-mixing module . 2.2 CONVENTIONAL TOKEN-MIXING MODULES AND CHALLENGES FOR ASR . There have been several studies on how to design the token-mixing module , such as token-mixing MLP in MLP-mixer , spatial gating unit ( SGU ) in gMLP , Fourier Domain Mixing in GFNet ( Rao et al. , 2021 ) and FNet ( Lee-Thorp et al. , 2021 ) . In this paper , we introduce these modules and discuss the issues in applying them to variable sequence length input data . In the following , the input is always assumed to be a data matrix with D channels and N tokens , X ∈ RD×N . Token-mixing MLP . The token-mixing MLP first transposes the input X and performs linear projection across the token dimension to mix the token information , X′ = ( W2σ ( W1X T ) ) T ∈ RD×N , ( 6 ) where W1 ∈ RN ′×N and W2 ∈ RN×N ′ refer to weights of linear projections , and σ ( · ) is an element-wise activation function . Albeit the components are the same as the channel-mixing module of Equation 2 , the linear projections are done on the token dimension by transposing the input . This module is shown in Figure 2 ( c ) . Spatial Gating Unit . The structure of the SGU is illustrated in Figure 2 ( d ) . The SGU first splits the input X into Xr ∈ R D 2 ×N and Xg ∈ R D 2 ×N . In gMLP , Xr is set as the first half of X and Xg for the second half of X . Then , it performs a gating operation . The output of SGU X′ is , X′ = Xr HSGU ∈ R D 2 ×N , with HSGU = ( W2XTg ) T ∈ RD2 ×N , ( 7 ) where denotes element-wise multiplication , and W2 ∈ RN×N refers to the weights of the linear projection . This gating operation mixes the information of the token effectively . The effectiveness of the SGU has been experimentally shown with image and language experiments compared to non-gating linear projection and other variants . Fourier Domain Mixing . GFNet and FNet use the discrete Fourier transform ( DFT ) in the tokenmixing module . GFNet applies a fixed length learnable filter by direct multiplication in the frequency domain . FNet mixes tokens by the application of a single forward 2D DFT , retaining only the real part of the output .

Summary: This paper proposes three architectural modifications on recently introduced MLP-based neural networks, such as MLP-Mixer or gMLP. These networks could not handle the variable-length input because the token-level mixing was performed by MLP. The proposed three modifications are: C-MLP (w/ depth-wise convolution), TS-MLP (w/ shift operator), and F-MLP (w/ Fourier transform and w/o gating). The authors evaluated their model on two ASR datasets, LibriSpeech-100h and TedLium-2, and achieved better WER compared to Transformer-based models with a similar number of parameters.