Patent Publication Number: US-11663488-B2

Title: Initialization of parameters for machine-learned transformer neural network architectures

Description:
CROSS-REFERENCE TO RELATED APPLICATIONS 
     This application claims the benefit of and priority to U.S. Provisional Application No. 62/976,040, filed Feb. 13, 2020, which is hereby incorporated by reference in its entirety. 
    
    
     BACKGROUND 
     This invention relates generally to transformer neural network architectures, and more particularly to initializing parameters during training of transformer architectures. 
     Transformer architectures are machine-learned neural network models used for, e.g., natural language processing (NLP) applications. For example, a transformer architecture may receive a sequence of input embeddings that represents a question and generate a sequence of output embeddings that represent an answer to the question. As another example, the transformer architecture may receive a sequence of input embeddings that represents a sentence in English and generate a sequence of output embeddings that represents a translation of the sentence in French. As yet another example, the transformer architecture may receive a sequence of input embeddings that represents a paragraph of text and generate a sequence of output embeddings that represents a summarized version of the text. 
     A transformer architecture includes a set of encoders and a set of decoders coupled to the set of encoders. An encoder is coupled to receive a set of input representations and generate an encoded output, and a decoder is coupled to receive a set of input representations that include the encoded output of the encoders and generate a decoded output that can be used to generate word predictions. Typically, an encoder or a decoder in the transformer architecture includes an attention block that is coupled to receive an input key, an input query, an input value and generate a set of attention representations. The attention block allows an attention representation of an encoder or decoder to respectively encode or decode a respective input based on the associations between the respective input to other inputs in the set. 
     In practice, the transformer architecture may also include additional blocks, such as multi-layer perceptron (MLP) neural network layers and normalization layers, within an encoder or a decoder that contribute to improving the accuracy of the model. Thus, the number of parameters of the transformer architecture that need to be determined during the training process may quickly increase, especially when there are a significant number of encoders and decoders. Moreover, the training process for a transformer architecture often involves learning rate warmup, in which the learning rate for the parameters are gradually increased at the start of the training process, leading to prohibitively slow training time and computational efficiency. While there are other methods that try to improve the training process, these methods require changing the architecture of the transformer model or result in sub-optimal prediction accuracies. 
     SUMMARY 
     An online system trains a machine-learned transformer architecture by an initialization method that allows the transformer architecture to be trained without normalization layers of learning rate warmup, resulting in significant improvements in computational efficiency and training time for transformer architectures. Specifically, an attention block included in an encoder or a decoder of the transformer architecture generates the set of attention representations by applying a key matrix to the input key, a query matrix to the input query, a value matrix to the input value to generate an output, and applying an output matrix to the output to generate the set of attention representations. In one embodiment, the initialization method is performed by scaling the parameters of the value matrix and the output matrix with a factor that is inverse to a number of the set of encoders or a number of the set of decoders. 
     The initialization method results from a previously unknown insight that learning rate warmup is needed during the training process of the transformer architecture because of the presence of normalization layers within the set of encoders and the set of decoders. By deploying such an initialization method, the online system can train transformer architectures without normalization layers and thus, without learning rate warmup. Moreover, the initialization method allows the online system to train deep transformer architectures with a significant number (e.g., 200 or more) of encoders and decoders that may open up new avenues for improvement in prediction accuracies or transformer architecture. 
     Specifically, during the training process, the online system accesses a machine-learned model including a set of encoders and a set of decoders. The set of encoders are coupled to receive a sequence of input embeddings and generate an encoded output. The set of decoders are coupled to receive a sequence of output embeddings and the encoded output from the set of encoders and generate a prediction for a next word. Specifically, at least one encoder of the set of encoders or at least one decoder of the set of decoders may include an attention block. The online system initializes parameters of the machine-learned model including parameters of the value matrix and the output matrix. The online system scales parameters of the value matrix and the output matrix by multiplying a factor that is inverse to a number of residual blocks in the set of encoders or a number of the set of decoders. A residual block is any block that is coupled to receive a set of inputs and generate a set of outputs, and in which the set of inputs are combined with the set of outputs to generate a combined representation. The combined representation can be fed into another block or layer as the inputs. 
     The online system obtains a set of training text that each includes an ordered set of training input embeddings and an ordered set of training output embeddings. For each training text in the set, the online system generates one or more estimated output embeddings by applying the set of encoders and the set of decoders to the ordered set of training input embeddings. The online system determines a loss function indicating a difference between the one or more estimated output embeddings and the ordered set of training output embeddings. The online system updates the parameters of the machine-learned model to reduce the loss function for the training text in the set. This process is repeated for the next set of training text until a convergence criteria for the parameters is reached. 
    
    
     
       BRIEF DESCRIPTION OF THE DRAWINGS 
         FIG.  1    is a high level block diagram of a system environment for an online system, in accordance with an embodiment. 
         FIG.  2    illustrates an example inference process for a machine-learned model with a transformer architecture, in accordance with an embodiment. 
         FIG.  3    is a block diagram of an architecture of the online system, in accordance with an embodiment. 
         FIG.  4    illustrates an architecture of an attention block of the machine-learned model with multi-head attention, in accordance with an embodiment. 
         FIG.  5    illustrates a method for training the machine-learned model with a transformer architecture, in accordance with an embodiment. 
     
    
    
     The figures depict various embodiments of the present invention for purposes of illustration only. One skilled in the art will readily recognize from the following discussion that alternative embodiments of the structures and methods illustrated herein may be employed without departing from the principles of the invention described herein. 
     DETAILED DESCRIPTION 
     Overview 
     An online system trains a machine-learned transformer architecture by an initialization method that allows the transformer architecture to be trained without normalization layers of learning rate warmup, resulting in significant improvements in computational efficiency and training time for transformer architectures. Specifically, an attention block included in an encoder or a decoder of the transformer architecture generates the set of attention representations by applying a key matrix to the input key, a query matrix to the input query, a value matrix to the input value to generate an output, and applying an output matrix to the output to generate the set of attention representations. In one embodiment, the initialization method is performed by scaling the parameters of the value matrix and the output matrix with a factor that is inverse to a number of residual blocks in the set of encoders or a number of the set of decoders. 
     System Environment 
       FIG.  1    is a high level block diagram of a system environment  100  for an online system  130 , in accordance with an embodiment. The system environment  100  shown by  FIG.  1    includes one or more client devices  110 A,  110 B, a network  120 , and an online system  130 . In alternative configurations, different and/or additional components may be included in the system environment  100 . 
     The online system  130  receives requests from client devices  110 A,  110 B to perform one or more tasks using machine-learned models. In one embodiment, the tasks may be related to natural language processing (NLP) applications, and the online system  130  may train and deploy machine-learned models that are configured to perform or aid in these tasks. Typically, NLP tasks involve using artificial intelligence to analyze and understand language and may include a variety of tasks including but not limited to sentiment analysis, text summarization, text translation, auto-correct, chatbot solutions, and the like. Responsive to receiving a request, the online system  130  encodes input text in the form of embedding vectors in a latent space. The online system  130  applies the one or more machine-learned models to the input embeddings to generate output text that is a response to the task of interest. 
     For example, a machine-learned model servicing a chatbot application may receive a sequence of input embeddings that represents a question and generate a sequence of output embeddings that represents an answer to the question. As another example, a machine-learned model for automatic language translation may receive a sequence of input embeddings that represents a sentence in English and generate a sequence of output embeddings that represents a translation of the sentence in French. As yet another example, a machine-learned model for language summarization may receive a sequence of input embeddings that represents a paragraph of text and generate a sequence of output embeddings that represents a summarized version of the text. 
       FIG.  2    illustrates an example inference process for a machine-learned model with a transformer architecture, in accordance with an embodiment. In one embodiment, the online system  130  trains and deploys a machine-learned model with a transformer architecture. As shown in  FIG.  2   , the transformer architecture includes a set of N encoders E 1 , E 2 , . . . , EN and a set of N decoders D 1 , D 2 , . . . , DN coupled to the set of encoders. An encoder is coupled to receive a set of input representations and generate an encoded output. A decoder is coupled to receive a set of input representations that include the encoded output of the encoders and generate a decoded output that can be used to generate word predictions. 
     In the example shown in  FIG.  2   , the online system  130  may receive a question “how are you?” received at a chatbot application, and a request to generate an answer to the question. The online system  130  encodes the question as a set of input embeddings x 1 , x 2 , x 3  that each represent the words “how,” “are,” “you?” respectively. In particular, the first encoder E 1  may be coupled to receive a set of input representations that are the set of encoded input embeddings and generate an encoded output. Each subsequent encoder may be coupled to receive a set of input representations that are the encoded output of a previous encoder and generate an encoded output for itself. For example, the last encoder EN may be coupled to receive a set of input representations that are the encoded output of the previous encoder EN−1 and generate a final encoded output {tilde over (x)} for the set of encoders. 
     The first decoder D 1  may be coupled to receive a set of input representations including a set of output embeddings and the encoded output {tilde over (x)} of the final encoder EN and generate a decoded output that can be used to generate word predictions. The output embeddings represent a sequence of word predictions that were generated by the machine-learned model at previous timesteps during the deployment process. In the example shown in  FIG.  2   , the transformer architecture may have generated a sequence of output embeddings ŷ 1 , ŷ 2 , ŷ 3  at previous timesteps that each represent the words “hello,” “I,” “am,” respectively. Each subsequent decoder may be coupled to receive a set of input representations that are the decoded output of a previous decoder and the final encoded output {acute over (x)} for the set of encoders and generate a decoded output for itself. For example, the last decoder DN may be coupled to receive a set of input representations that are the decoded output of the previous decoder DN−1 and the final encoded output {tilde over (x)} for the set of encoders and generate a final decoded output {tilde over (y)}. 
     The transformer architecture also includes one or more blocks that are coupled to receive the final decoded output {tilde over (y)} and generate an output embedding representing the next word prediction. In the example shown in  FIG.  2   , the transformer architecture includes a linear neural network layer coupled to receive the final decoded output {tilde over (y)} and generate an output. The transformer architecture also includes a softmax layer coupled to receive the output of the linear neural network layer and generate an output embedding prediction ŷ 4  by transforming the output of the linear layer with a softmax function. The output embedding prediction ŷ 4  represents the next word prediction that should come after the existing predictions represented by the sequence of output embeddings ŷ 1 , ŷ 2 , ŷ 3 . In the example shown in  FIG.  2   , the output embedding prediction ŷ 4  may represent the word “fine.” 
     In the subsequent timestep during the deployment process, the first decoder D 1  may now be coupled to receive the next sequence of output embeddings ŷ 2 , ŷ 3 , ŷ 4 , and this decoding process is repeated until all predictions for the sequence of input embeddings are generated. For example, the decoding process may be repeated until the next word prediction generated by the transformer architecture is an end token (e.g., period “.” or an “&lt;end&gt;” indicator) that signals that all predictions for the given sequence of input embeddings have been generated. 
     Typically, an encoder or a decoder in the transformer architecture includes one or more attention blocks. An attention block is coupled to receive an input key a, an input query b, and an input value c and generate a set of attention representations. The attention block allows an attention representation of an encoder or decoder to respectively encode or decode a respective input based on the associations between the respective input to other inputs to the attention block. Specifically, in one embodiment, an attention block generates the set of attention representations by applying a key matrix k to the input key a, a query matrix q to the input query b, and a value matrix v to the input value c to generate an output, and applying an output matrix w to the output to generate the set of attention representations. The parameters of the key matrix k, the query matrix q, the value matrix v, and the output matrix w are learned during the training process of the transformer architecture. 
     The example shown in  FIG.  2    is described with respect to an attention block that includes one or more key matrices k, query matrices  q , and value matrices v that each have a set of trained parameters. However, it should be appreciated that in other embodiments, the attention representations can be generated by performing any type of operation(s) that involves applying a set of trained parameters to the input key a, input query b, or input value c. The operation may be represented as a matrix operation but is not limited thereto. For example, the attention block may include one or more neural network layers (e.g., fully connected layers, skip layers, etc.) that are coupled to receive the input key a, input query b, or input value c and generate an output by applying a set of parameters for the neural network layers to each of the input key a, input query b, or input value c to generate an output. The attention block may further include one or more neural network layers (e.g., fully connected layer, skip layers, etc.) coupled to receive the output and generate the set of attention representations by applying a set of parameters for the neural network layers to the output. 
     Specifically, an encoder includes an attention block that is a self-attention block coupled to receive an input key a, an input query b, and an input value c that each corresponds to a common set of inputs. For example, the self-attention block  204  in the first encoder E 1  is coupled to receive an input key a, an input query b, and an input value c that each are the set of input embeddings x 1 , x 2 , x 3  and generate a set of attention representations. An attention representation for the input embedding x 1  in the input query b may be an encoding that incorporates associations between the input embedding x 1  to other inputs x 2 , x 3  in the set. The self-attention block in a subsequent encoder (e.g., encoders E 2 , . . . , EN) is coupled to receive an input key a, an input query b, and an input value c that each corresponds to the encoded output of a previous encoder and similarly generate a set of attention representations. 
     An encoder may also include a first add and normalization layer placed after the self-attention block. In the example shown in  FIG.  2   , the first encoder E 1  includes a first add and normalization layer  208 . The first add and normalization layer combines the set of attention representations generated by the self-attention block with the set of inputs received at the self-attention block to generate a combined representation. Thus, as defined herein, the self-attention block in an encoder may be a residual block in which the set of inputs to the block, without performing any operation, are combined with the set of outputs generated by the block before being input to the next layer. The first add and normalization layer may further normalize the combined representation by, for example, dividing each element of the combined representation with the vector magnitude of the combined representation. 
     The encoder further includes a multi-layer perceptron (MLP) block that includes one or more neural network layers with a set of trained parameters. In the example shown in  FIG.  2   , the first encoder E 1  includes an MLP block  212 . The MLP block is coupled to receive a set of inputs that are the outputs of the first add and normalization layer and generate an MLP output by applying the set of parameters to the set of inputs. The encoder may further include a second add and normalization layer placed after the MLP block. In the example shown in  FIG.  2   , the first encoder E 1  includes a second add and normalization layer  216 . The second add and normalization layer combines the MLP output generated by the MLP block with the set of inputs received at the MLP block to generate a combined representation. Thus, the MLP block in an encoder may also be configured as a residual block. The second add and normalization layer may further normalize the combined representation to generate the encoded output for the encoder. 
     Similarly, a decoder also includes an attention block that is a self-attention block. For example, the self-attention block  254  in the first decoder D 1  is coupled to receive an input key a, an input query b, and an input value c that each corresponds to the set of output embeddings ŷ 1 , ŷ 2 , ŷ 3  and generate a set of attention representations. An attention representation for an output embedding ŷ 1  in the input query may be a decoding that incorporates associations between the output embedding ŷ 1  to other inputs ŷ 2 , ŷ 3  in the set. The self-attention block in a subsequent decoder (e.g., decoders D 2 , . . . , DN) is coupled to receive an input key a, an input query b, and an input value c that each corresponds to the decoded output of a previous decoder and similarly generate a set of attention representations. 
     A decoder may also include a first add and normalization layer placed after the self-attention block. In the example shown in  FIG.  2   , the first decoder D 1  includes a first add and normalization layer  258 . The first add and normalization layer combines the set of attention representations generated by the self-attention block with the set of inputs received at the self-attention block to generate a combined representation and may further normalize the combined representation. Thus, as defined herein, the self-attention block in a decoder may also be configured as a residual block. 
     Different from the encoder, the decoder may also include an encoder-decoder attention block that is placed after the self-attention block of the decoder. In the example shown in  FIG.  2   , the first decoder D 1  includes an encoder-decoder attention block  262 . The encoder-decoder attention block is coupled to receive an input key a that corresponds to the final encoder output {tilde over (x)}, an input query b that corresponds to outputs from the first add and normalization layer, and an input value c that corresponds to the final encoder output {tilde over (x)}. Different from the self-attention block, an attention representation for a representation in the input query b may be a decoding of the representation that incorporates associations between the representation in the input query to the encoded output {tilde over (x)}. 
     The decoder may also include a second add and normalization layer placed after the encoder-decoder attention block of the decoder. In the example shown in  FIG.  2   , the first decoder D 1  includes a second add and normalization layer  266 . The second add and normalization layer combines the set of attention representations generated by the encoder-decoder attention block with the set of inputs received at the encoder-decoder attention block to generate a combined representation and may further normalize the combined representation. Thus, as defined herein, the encoder-decoder attention block in a decoder may also be configured as a residual block. 
     The decoder further includes an MLP block coupled to receive a set of inputs that are the outputs of the second add and normalization layer and generate an MLP output by applying the set of parameters to the set of inputs. In the example shown in  FIG.  2   , the first decoder D 1  includes an MLP block  270 . The decoder may further include a third add and normalization layer placed after the MLP block. In the example shown in  FIG.  2   , the first decoder D 1  includes a third add and normalization layer  274 . The third add and normalization layer combines the MLP output generated by the MLP block with the set of inputs received at the MLP block to generate a combined representation and may further normalize the combined representation. Thus, the MLP block of a decoder may also be configured as a residual block. 
     The online system  130  trains the transformer architecture by obtaining a set of training text that each includes an ordered set of training input embeddings and an ordered set of training output embeddings that correspond to known word predictions for the set of training input embeddings that address the task the transformer architecture should be trained for. The online system  130  initializes parameters of the transformer architecture to begin the training process. For each training text in the set, the online system  130  generates one or more estimated output embeddings by applying the set of encoders and the set of decoders to the ordered set of training input embeddings. The online system  130  determines a loss function indicating a difference between the one or more estimated output embeddings and the ordered set of training output embeddings. The online system  130  updates the parameters of the machine-learned model to reduce the loss function for the training text in the set. This process is repeated for multiple training iterations until a convergence criteria for the parameters is reached. 
     In one embodiment, the online system  130  trains the transformer architecture by an initialization method that scales the parameters of the value matrix v and the output matrix w with a factor that is inverse to a number of residual blocks in the set of encoders or a number of the set of decoders. Typically, the training process for the transformer architecture is difficult because the number of parameters of the transformer architecture can quickly increase with the addition of encoders and decoders. Empirically, the training process for a transformer architecture involved learning rate warmup, in which the magnitude the update at each training iteration was bounded at earlier iterations of the training process and gradually increased as the training progressed. This led to prohibitively slow training time and computational efficiency. While there are other methods that try to improve the training process, these methods require changing the architecture of the transformer model or result in sub-optimal prediction accuracies. 
     The initialization method used by the online system  130  results from a previously unknown insight that learning rate warmup is needed during the training process of the transformer architecture because of the presence of normalization layers within the set of encoders and the set of decoders. By deploying such an initialization method, the online system  130  can train transformer architectures without normalization layers and thus, without learning rate warmup. Moreover, the initialization method allows the online system  130  to train deep transformer architectures with a significant number (e.g., 200 or more) of encoders and decoders that may open up new avenues for improvement in prediction accuracies or transformer architecture. 
     Thus, in one embodiment, the online system  130  may also deploy the initialization method to train parameters of a transformer architecture that has one or more encoders or decoders that do not include normalization layers. For example, returning to the example architecture shown in  FIG.  2   , the first encoder E 1  may include a first add layer instead of the first add and normalization layer  208  that combines the set of attention representations generated by the self-attention block  204  with the set of inputs received at the self-attention block  204  to generate a combined representation without performing normalization. The first encoder E 1  may also include a second add layer instead of the second add and normalization layer  216  that combines the MLP output generated by the MLP block  212  with the set of inputs received at the MLP block  212  to generate a combined representation without performing normalization. Thus, the self-attention block and the MLP block of an encoder can be configured as residual blocks but without normalization of the combined representations. This structure may be applied to other encoders in the set of encoders. 
     As another example, the first decoder D 1  may include a first add layer instead of the first add and normalization layer  258  that combines the set of attention representations generated by the self-attention block  254  with the set of inputs received at the self-attention block  254  to generate a combined representation without performing normalization, a second add layer instead of the second add and normalization layer  266  that combines the set of attention representations generated by the encoder-decoder attention block  262  with the set of inputs received at the encoder-decoder attention block  262  to generate a combined representation without performing normalization, or a third add layer instead of the third add and normalization layer  274  that combines the MLP output generated by the MLP block  270  with the set of inputs received at the MLP block  270  to generate a combined representation without performing normalization. Thus, the self-attention block, the encoder-decoder attention block, and the MLP block of a decoder can be configured as residual blocks but without normalization of the combined representations. This structure may be applied to other decoders in the set of decoders. 
     Returning to  FIG.  1   , the client device  110  is a computing device such as a smartphone with an operating system such as ANDROID® or APPLE® IOS®, a tablet computer, a laptop computer, a desktop computer, or any other type of network-enabled device. In the system environment  100  shown in  FIG.  1   , the client device  110  allows a user to submit requests to the online system  130  to perform tasks that can be processed by the transformer architecture. For example, a user of the client device  110  may be affiliated with an entity that deploys one or more applications that require natural language processing or other types of tasks that can be processed by the transformer architecture. A request submitted by the client device  110  may be, for example, a request to generate predictions for chatbot applications, a request to translate a sequence of words from one language to another language, a request to synthesize a story or narrative given a sequence of starting words, and the like. 
     A typical client device  110  includes the hardware and software needed to connect to the network  122  (e.g., via WiFi and/or 4G, 5G or other wireless telecommunication standards). The client device  110  may include an operating system and various applications that run on the operating system that enable the users to submit the requests. For example, the client device  110  may be include browser applications or standalone applications deployed by the online system  130  that allow users of an organization to interact with the online system  130  to submit the requests. 
     The network  122  provides a communication infrastructure between the client devices  110  and the online system  130 . The network  122  is typically the Internet, but may be any network, including but not limited to a Local Area Network (LAN), a Metropolitan Area Network (MAN), a Wide Area Network (WAN), a mobile wired or wireless network, a private network, or a virtual private network. 
     System Environment 
       FIG.  3    is a block diagram of an architecture of the online system  130 , in accordance with an embodiment. The online system  130  shown by  FIG.  3    includes a document management module  320 , a training module  330 , and a prediction module  335 . The online system  130  also includes a training corpus  360 . In alternative configurations, different and/or additional components may be included in the online system  130 . 
     The document management module  320  manages the training corpus  360  of training data that are used to train a set of parameters of the transformer architecture. In one embodiment, the training corpus  360  includes multiple instances of data that each include an ordered set of training input embeddings and an ordered set of training output embeddings that correspond to known word predictions for the set of training input embeddings that address the tasks the transformer architecture should be trained for. The training input embeddings represent a sequence of input words in a latent space, and the training output embeddings represent a sequence of output words in the latent space. 
     For example, when the transformer architecture is used for automatic translation from English to German, the set of training input embeddings may correspond to a first sentence of words in English and the set of training output embeddings may correspond to a second sentence of words in German that is a translation of the first sentence. As another example, when the transformer architecture is used for servicing a chatbot application, the set of training input embeddings may correspond to a questions asked by a user of the chatbot application and the set of training output embeddings may correspond to a prediction for an answer to the question. 
     The training module  330  trains parameters of the transformer architecture during a training process. The training module  330  may generate an appropriate transformer architecture that can address the task requests received by the client devices  110 . For example, the transformer architecture may adjust the depth of the transformer architecture by determining the number of encoders and decoders in the machine-learned model. Moreover, the transformer architecture may also include a first embedding block coupled to receive an encoding of a sequence of input words and generate the sequence of input embeddings, and a second embedding block coupled to receive an encoding of a sequence of output words and generate the sequence of output embeddings. As described previously, the transformer architecture may exclude layer normalization in one or more encoders or decoders, and instead may only include an add layer in those positions instead of an add and normalization layer. 
     The training module  330  may initialize the parameters of the transformer architecture including parameters of attention blocks (e.g., self-attention blocks or encoder-decoder attention blocks) and MLP blocks included in the set of encoders and the set of decoders, and other layers including linear layers placed outside the set of encoders or decoders. In one embodiment, the training module  330  initializes the parameters of the transformer architecture excluding parameters of the embedding blocks using a Xavier initialization method where, for example, the parameters of each layer are initialized by sampling the values of the parameters from a uniform distribution with a range [−1/sqrt(n), 1/sqrt(n)], where n is the size of the previous layer. For example, n may be the number or units or elements in the previous neural network layer. The parameters of the embedding blocks may be initialized using Gaussian initialization. 
     The training module  330  performs the initialization method that scales the parameters of the transformer architecture to improve the training process. In one embodiment, the parameters of the transformer architecture are scaled to satisfy the following equations:
 
 L   e (∥ v   e ∥ 2   ∥x∥   2   +∥w   e ∥ 2   ∥x∥   2   +∥v   e ∥ 2   ∥w   e ∥ 2 )=Θ(1)
 
where v e  denotes parameters of the value matrix in the encoder, w e  denotes parameters of the output matrix in the encoder, x denotes the set of input embeddings, and L e  is the number of residual blocks in the set of encoders (e.g., 2N or 2N+1 where N is the number of encoders in the set),
 
 L   d (∥ v   d ∥ 2   ∥w   d ∥ 2   +∥v   d ∥ 2   ∥y∥   2   +∥w   d ∥ 2   ∥y∥   2   +∥v   d ∥ 2   ∥w   d ∥ 2   +∥v   d ∥ 2   ∥m∥   2   +∥v   d ∥ 2   ∥m∥   2 )=Θ(1)
 
where v d  denotes parameters of the value matrix in the decoder, w d  denotes parameters of the output matrix in the decoder, y denotes the set of output embeddings, m denotes the final encoded output, and L d  is the number of residual blocks in the set of decoders (e.g., 3N or 3N+1 where N is the number of decoders in the set), and
 
 L   e (∥ v   d ∥ 2   ∥w   d ∥ 2 )=Θ(1).
 
     In one embodiment, when the set of decoders each include two attention blocks (e.g., self-attention block and encoder-decoder attention block) and a MLP block, the training module  330  may scale parameters of the value matrix v and the output matrix w in each attention block and the parameters of the MLP block by a scaling factor inverse to a number of residual blocks in the set of decoders. Specifically, when the dimensionality of the value matrices and the output matrices are the same, these parameters may be scaled by a scaling factor of L d   −1/4 , where L d  denotes the number of residual blocks in the set of decoders, and may indicate the total number of attention blocks, MLP blocks, and the embedding block for the set of decoders. For example, L d  may be given by L d =3N+1 or 3N where N is the number of decoders in the set of decoders. In one embodiment, when the set of encoders each include an attention block (e.g., self-attention block) and a MLP block, the training module  330  may scale the parameters of the value matrix v and the output matrix w in each attention block and the parameters of a MLP block by a scaling factor inverse to a number of residual blocks in the set of encoders. Specifically, when the dimensionality of the value matrices and the output matrices are the same, these parameters may be scaled by a scaling factor of 0.67N −1/4  assuming that there are L e =2N or 2N+1 residual blocks in the set of encoders, where N is the number of encoders in the set of encoders. 
     During the training process, the training module  330  selects a set of training text and trains parameters of the machine-learned model by repeatedly iterating between a forward pass step and a backpropagation step. During the forward pass step, the training module  330  generates one or more estimated output embeddings by applying parameters of the transformer architecture to the ordered set of training input embeddings in the set of training text. The training module  330  determines a loss function indicating a difference between the one or more estimated output embeddings and the ordered set of training output embeddings. During the backpropagation step, the training module  330  updates the parameters of the transformer architecture to reduce the loss function. This process is iteratively repeated for the next sets of training text until a convergence criterion for the parameters is reached. 
     As described previously, by performing the initialization method, the training module  330  may update the parameters of the transformer architecture without learning rate warmup, such that the training process can be performed faster and in a more computationally efficient manner. Conventionally, the training process of the transformer architecture often required significantly small learning rates at earlier iterations of the training process that bounded how much the parameters could be updated in one iteration, making training prohibitively slow. By deploying the initialization method, the updates to the parameters may be automatically adjusted so that learning rate warmup may no longer be needed. 
       FIG.  4    illustrates an architecture of an attention block of the machine-learned model with multi-head attention, in accordance with an embodiment. In one embodiment, the training module  330  generates a transformer architecture that includes attention blocks with a multi-headed structure. For example, any one or more of or all of the self-attention blocks in an encoder or the self-attention blocks or the encoder-decoder attention blocks in a decoder may have the multi-headed structure. 
     As shown in  FIG.  4   , an attention block with a multi-headed structure is coupled to receive an input key a, an input query b, and an input value c and generate a set of attention representations. Specifically, the multi-headed structure includes a plurality of attention heads. Each attention head is associated with its own key matrix k, query matrix q, and value matrix v. The attention head is coupled to receive the input key a, the input query b, and the input value c and generate a respective output by applying the respective key matrix k to the input key a, the respective query matrix q to the input query b, and a value matrix v to the input value c. For example, the attention block in  FIG.  4    includes attention heads 1, 2, . . . , H, where an attention head i is associated with its own key matrix k i , query matrix q i , and value matrix v i . An attention head i is coupled to receive the input key a, the input query b, and the input value c to generate a respective output z i . The attention representation for the attention block z is generated by concatenating the outputs z 1 , z 2 , . . . , z H  and applying an output matrix w to the concatenated outputs. 
     As described above in conjunction with  FIGS.  1  and  2   , each of the key matrices k i , query matrices q i , value matrices v i  may be represented by a set of parameters in any form, not limited to matrix operations. For example, each attention head may include one or more neural network layers (e.g., fully connected layer, skip layer, etc.) that are coupled to receive the input key a, input query b, input value c and generate the respective output z i  by applying a set of parameters for the neural network layers to the input key a, input query b, input value c. The attention block may further include one or more neural network layers (e.g., fully connected layer, skip layer, etc.) coupled to receive the concatenated outputs z 1 , z 2 , . . . , z H  and generate the set of attention representations by applying a set of parameters for the neural network layers to the concatenated outputs. 
     Thus, in one embodiment, when a multi-headed structure is applied to an attention block, the training module  330  may also initialize parameters of the transformer architecture including parameters of multi-attention blocks including each respective key matrix k i , query matrix q i , and value matrix v i  of an attention head. Moreover, the training module  330  may also scale the values of each respective key matrix k i , query matrix q i , and value matrix v i  of an attention head with the scaling factors described above during the initialization method. Subsequently, similar to that described above, the training module  330  repeatedly iterates between the forward pass step and the backpropagation step to train the parameters of the transformer architecture. 
     The prediction module  335  receives requests from client devices  110  to perform one or more tasks related to, for example, natural language processing (NLP) applications, that can be responded to by using the trained transformer architecture. In one embodiment, when the applications are NLP applications, the requests may include a sequence of words that the prediction module  335  may parse and map into a sequence of input embeddings that can be fed into the transformer architecture, or alternatively, can encode the sequence of words into encodings that can be fed into the transformer architecture such that the input embeddings can be generated within the model. 
     Responsive to receiving a request, the prediction module  335  may apply the trained transformer architecture to the sequence of input embeddings such that a sequence of output embeddings are generated as predictions. The prediction module  335  may be provided to the client device  110  of the request as the response to the request. In one embodiment, when the training module  330  trains multiple transformer architectures each configured to address a respective task (e.g., translation, chatbot application, etc.) the prediction module  335  may select the appropriate model and apply the model to the request. For example, the prediction module  335  may receive a request from a translation application that includes a sequence of words that represents an English phrase. The prediction module  335  may apply a transformer architecture trained to perform this task to generate the sequence of output words. The prediction module  335  may provide the sequence of output words as the response to the request. 
     Method of Training Transformer Architecture 
       FIG.  5    illustrates a method for training a machine-learned transformer architecture, in one embodiment. In one embodiment, the process of  FIG.  5    is performed by various modules of the online system  130 . Other entities may perform some or all of the steps of the process in other embodiments. Likewise, embodiments may include different and/or additional steps, or perform the steps in different orders. 
     The online system  130  accesses  502  a machine-learned model. The machine-learned model includes a set of encoders and a set of decoders, the set of encoders are coupled to receive a sequence of input embeddings and generate an encoded output, and a set of decoders coupled to receive a sequence of output embeddings and the encoded output and generate a prediction for a next word. At least one encoder of the set of encoders or at least one decoder of the set of decoders includes an attention block. The attention block is coupled to receive an input key, an input query, an input value and generate an attention by applying a key matrix to the input key, a query matrix to the input query, a value matrix to the input value to generate an output, and applying an output matrix to the output to generate the attention. 
     The online system  130  initializes  504  parameters of the machine-learned model including parameters of the value matrix and the output matrix. The online system  130  scales  506  the parameters of the value matrix and the output matrix by multiplying a scaling factor that is inverse to a number of the set of encoders or a number of the set of decoders. 
     The online system  130  obtains  508  a set of training text. Each training text includes an ordered set of training input embeddings and an ordered set of training output embeddings. For each training text in the set, the online system  130  generates  510  one or more estimated output embeddings by applying the set of encoders and the set of decoders to the ordered set of training input embeddings. The online system determines  512  a loss function indicating a difference between the one or more estimated output embeddings and the ordered set of training output embeddings. The online system  130  updates  514  the parameters of the machine-learned model to reduce the loss function for the training text in the set. 
     SUMMARY 
     The foregoing description of the embodiments of the invention has been presented for the purpose of illustration; it is not intended to be exhaustive or to limit the invention to the precise forms disclosed. Persons skilled in the relevant art can appreciate that many modifications and variations are possible in light of the above disclosure. 
     Some portions of this description describe the embodiments of the invention in terms of algorithms and symbolic representations of operations on information. These algorithmic descriptions and representations are commonly used by those skilled in the data processing arts to convey the substance of their work effectively to others skilled in the art. These operations, while described functionally, computationally, or logically, are understood to be implemented by computer programs or equivalent electrical circuits, microcode, or the like. Furthermore, it has also proven convenient at times, to refer to these arrangements of operations as modules, without loss of generality. The described operations and their associated modules may be embodied in software, firmware, hardware, or any combinations thereof. 
     Any of the steps, operations, or processes described herein may be performed or implemented with one or more hardware or software modules, alone or in combination with other devices. In one embodiment, a software module is implemented with a computer program product comprising a computer-readable medium containing computer program code, which can be executed by a computer processor for performing any or all of the steps, operations, or processes described. 
     Embodiments of the invention may also relate to an apparatus for performing the operations herein. This apparatus may be specially constructed for the required purposes, and/or it may comprise a general-purpose computing device selectively activated or reconfigured by a computer program stored in the computer. Such a computer program may be stored in a non-transitory, tangible computer readable storage medium, or any type of media suitable for storing electronic instructions, which may be coupled to a computer system bus. Furthermore, any computing systems referred to in the specification may include a single processor or may be architectures employing multiple processor designs for increased computing capability. 
     Embodiments of the invention may also relate to a product that is produced by a computing process described herein. Such a product may comprise information resulting from a computing process, where the information is stored on a non-transitory, tangible computer readable storage medium and may include any embodiment of a computer program product or other data combination described herein. 
     Finally, the language used in the specification has been principally selected for readability and instructional purposes, and it may not have been selected to delineate or circumscribe the inventive subject matter. It is therefore intended that the scope of the invention be limited not by this detailed description, but rather by any claims that issue on an application based hereon. Accordingly, the disclosure of the embodiments of the invention is intended to be illustrative, but not limiting, of the scope of the invention, which is set forth in the following claims.