Initialization of parameters for machine-learned transformer neural network architectures
An online system trains a transformer architecture by an initialization method which allows the transformer architecture to be trained without normalization layers of learning rate warmup, resulting in significant improvements in computational efficiency 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. The initialization method may be 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.
1. A system comprising:
a processor configured to execute instructions;
a computer-readable medium containing instructions for execution on the processor, the instructions causing the processor to perform steps of:
accessing a machine-learned model including:
a set of blocks that each perform a respective operation on one or more inputs to the block to generate one or more outputs for the block, wherein one or more of the blocks are configured as attention blocks and an 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, and wherein at least a subset of the blocks are configured as residual blocks in which an input to a residual block is added to an output of the residual block;
initializing parameters of the machine-learned model including parameters of the value matrix and the output matrix for each attention block;
scaling the parameters of the value matrix and the output matrix by multiplying a scaling factor that depends on a number of the subset of residual blocks;
obtaining a set of training data, each training data instance including an ordered set of training input embeddings and an ordered set of training output embeddings;
determining a loss function indicating a difference between one or more estimated output embeddings and the ordered set of training output embeddings, the estimated output embeddings generated by applying the machine-learned model to the ordered set of training input embeddings; and
updating the parameters of the machine-learned model to reduce the loss function.
2. The system of claim 1 , wherein at least one attention block does not include a normalization layer coupled to receive a set of inputs and normalize the set of inputs.
3. The system of claim 1 , wherein the set of blocks are arranged as a set of decoders, and at least one decoder includes a first attention block, a second attention block, and a multi-layer perceptron (MLP) block placed after the first and second attention blocks,
wherein the first attention block is a self-attention block coupled to receive the input key, the input query, and the input value that each corresponds to the sequence of output embeddings or an output of a previous decoder, and wherein the second attention block is a cross-attention block coupled to receive the input query from the sequence of output embeddings or an output of a previous decoder, and the input key and the input value from an encoder.
4. The system of claim 3 , wherein the scaling factor is a number inverse to 3N+1, wherein N is the number of decoders in the set.
5. The system of claim 1 , wherein the set of blocks are arranged as a set of encoders, and at least one encoder includes an attention block and a multi-layer perceptron (MLP) block placed after the attention block,
wherein the attention block is a self-attention block coupled to receive the input key, the input query, and the input value that each corresponds to the sequence of output embeddings or an output of a previous decoder.
6. The system of claim 5 , wherein the scaling factor is a number inverse to 2N+1, where N is the number of encoders in the set.
7. The system of claim 1 , wherein the scaling factor is inverse to the number of residual blocks.
8. The system of claim 1 , wherein initializing the parameters of the machine-learned model comprises initializing values of at least a portion of the parameters of the machine-learned model using a Xavier initialization method.
9. The system of claim 8 , wherein initializing the parameters of the machine-learned model further comprises sampling the values of the at least the portion of the parameters from a uniform distribution with a range of [−1/sqrt (n), 1/sqrt (n)], where n is a size of a previous neural network layer.
10. The system of claim 1 , wherein the subset of blocks configured as residual blocks include the one or more blocks configured as attention blocks.
11. A method, comprising:
accessing, by a processor when executing instructions, a machine-learned model including:
a set of blocks that each perform a respective operation on one or more inputs to the block to generate one or more outputs for the block, wherein one or more of the blocks are configured as attention blocks and an 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, and wherein at least a subset of the blocks are configured as residual blocks in which an input to a residual block is added to an output of the residual block;
initializing, by the processor, parameters of the machine-learned model including parameters of the value matrix and the output matrix for each attention block;
scaling, by the processor, the parameters of the value matrix and the output matrix by multiplying a scaling factor that depends on a number of the subset of residual blocks;
obtaining, by the processor, a set of training data, each training data instance including an ordered set of training input embeddings and an ordered set of training output embeddings;
determining, by the processor, a loss function indicating a difference between the one or more estimated output embeddings and the ordered set of training output embeddings; and
updating, by the processor, the parameters of the machine-learned model to reduce the loss function for the training data in the set.
12. The method of claim 11 , wherein at least one attention block does not include a normalization layer coupled to receive a set of inputs and normalize the set of inputs.
13. The method of claim 11 , wherein the set of blocks are arranged as a set of decoders, and at least one decoder includes a first attention block, a second attention block, and a multi-layer perceptron (MLP) block placed after the first and second attention blocks,
wherein the first attention block is a self-attention block coupled to receive the input key, the input query, and the input value that each corresponds to the sequence of output embeddings or an output of a previous decoder, and wherein the second attention block is a cross-attention block coupled to receive the input query from the sequence of output embeddings or an output of a previous decoder, and the input key and the input value from an encoder.
14. The method of claim 13 , wherein the scaling factor is a number inverse to 3N+1, wherein N is the number of decoders in the set.
15. The method of claim 11 , wherein the set of blocks are arranged as a set of encoders, and at least one encoder includes an attention block and a multi-layer perceptron (MLP) block placed after the attention block,
wherein the attention block is a self-attention block coupled to receive the input key, the input query, and the input value that each corresponds to the sequence of output embeddings or an output of a previous decoder.
16. The method of claim 15 , wherein the scaling factor is a number inverse to 2N+1, where N is the number of encoders in the set.
17. The method of claim 11 , wherein the scaling factor is inverse to the number of residual blocks.
18. The method of claim 11 , wherein initializing the parameters of the machine-learned model comprises initializing values of at least a portion of the parameters of the machine-learned model using a Xavier initialization method.
19. The method of claim 18 , wherein initializing the parameters of the machine-learned model further comprises sampling the values of the at least the portion of the parameters from a uniform distribution with a range of [−1/sqrt (n), 1/sqrt (n)], where n is a size of a previous neural network layer.
20. The method of claim 11 , wherein the subset of blocks configured as residual blocks include the one or more blocks configured as attention blocks.