Training large-scale vision transformer neural networks
Methods, systems, and apparatus, including computer programs encoded on computer storage media, for training Vision Transformer (ViT) neural networks.
1 . A method of training a vision Transformer neural network, the vision Transformer neural network configured to:
obtain a plurality of image patches of an image, wherein each image patch comprises a different subset of the pixels of the image;
process the plurality of image patches to generate an input sequence comprising a respective input element at each of a plurality of positions, wherein the input sequence includes a respective input element corresponding to each of the plurality of image patches;
process the input sequence through a plurality of self-attention neural network blocks to generate an output sequence comprising a respective output element at each of the positions; and
process one or more of the output elements using one or more output layers to generate a classification output for the image, and the method comprising:
obtaining first training data, the first training data comprising a plurality of training images and a respective target classification output for each training image; and
training the vision Transformer neural network on the first training data, the training comprising:
during the training,
updating parameters of the one or more output layers using a first weight decay value, and
updating parameters of the plurality of self-attention neural network blocks using a second weight decay value,
wherein the first weight decay value is higher than the second weight decay value, and wherein updating the parameters of the one or more output layers and the parameters of the one or more self-attention blocks comprises repeatedly performing the following:
computing, using one or more of the training examples, gradients of an objective function; and
applying an optimizer to the gradients to generate a gradient-based update to the parameters, wherein the optimizer makes use of a respective momentum value for each of the parameters, and wherein the respective momentum values are stored with a reduced precision relative to the parameters.
2 . The method of claim 1 , wherein the first weight decay value is greater than or equal to .3 while the second weight decay value is less than .3.
3 . The method of claim 2 , wherein the first weight decay value is greater than or equal to 3.0.
4 . The method of claim 3 , wherein the second weight decay value is less than or equal to . 1.
5 . The method of claim 1 , wherein each input element in the input sequence corresponds to a respective one of the image patches, and wherein the one or more output layers comprise:
an aggregation layer block that is configured to aggregate all of the output elements to generate an aggregated output element; and
one or more final output layers that are configured to generate the classification output from the aggregated output element.
6 . The method of claim 5 , wherein the one or more final output layers are a single linear layer that is configured to map the aggregated output element to the classification output.
7 . The method of claim 5 , wherein the aggregation block is configured to apply multihead attention pooling to the output elements to generate the aggregated output element.
8 . The method of claim 5 , wherein the aggregation block is configured to apply global average pooling to the output elements to generate the aggregated output element.
9 . The method of claim 1 , wherein the training further comprises:
during an initial phase of the training, linearly annealing a learning-rate for the parameters of the output layers and the self-attention blocks away from zero;
during a final phase of the training, linearly annealing the learning-rate toward zero; and
during a main phase of the training that is between the initial phase and the final phase, applying a schedule that prevents the learning-rate from reaching zero.
10 . The method of claim 1 , wherein the respective momentum values are stored with half-precision relative to the parameters.
11 . The method of claim 1 , further comprising:
after training the vision Transformer on the first training data, training the plurality of self-attention neural network blocks jointly with a different set of one or more output layers on second training data to perform a different, downstream task.
12 . One or more non-transitory computer storage media storing instructions that when executed by one or more computers cause the one more computers to perform operations for training a vision Transformer neural network, the vision Transformer neural network configured to:
obtain a plurality of image patches of an image, wherein each image patch comprises a different subset of the pixels of the image;
process the plurality of image patches to generate an input sequence comprising a respective input element at each of a plurality of positions, wherein the input sequence includes a respective input element corresponding to each of the plurality of image patches;
process the input sequence through a plurality of self-attention neural network blocks to generate an output sequence comprising a respective output element at each of the positions; and
process one or more of the output elements using one or more output layers to generate a classification output for the image, and the method comprising:
obtaining first training data, the first training data comprising a plurality of training images and a respective target classification output for each training image; and
training the vision Transformer neural network on the first training data, the training comprising:
during the training,
updating parameters of the one or more output layers using a first weight decay value, and
updating parameters of the plurality of self-attention neural network blocks using a second weight decay value,
wherein the first weight decay value is higher than the second weight decay value, and wherein updating the parameters of the one or more output layers and the parameters of the one or more self-attention blocks comprises repeatedly performing the following:
computing, using one or more of the training examples, gradients of an objective function; and
applying an optimizer to the gradients to generate a gradient-based update to the parameters, wherein the optimizer makes use of a respective momentum value for each of the parameters, and wherein the respective momentum values are stored with a reduced precision relative to the parameters.
13 . A system comprising one or more computers and one or more storage devices storing instructions that when executed cause the one more computers to perform operations for training a vision Transformer neural network, the vision Transformer neural network configured to:
obtain a plurality of image patches of an image, wherein each image patch comprises a different subset of the pixels of the image;
process the plurality of image patches to generate an input sequence comprising a respective input element at each of a plurality of positions, wherein the input sequence includes a respective input element corresponding to each of the plurality of image patches;
process the input sequence through a plurality of self-attention neural network blocks to generate an output sequence comprising a respective output element at each of the positions; and
process one or more of the output elements using one or more output layers to generate a classification output for the image, and the method comprising:
obtaining first training data, the first training data comprising a plurality of training images and a respective target classification output for each training image; and
training the vision Transformer neural network on the first training data, the training comprising:
during the training,
updating parameters of the one or more output layers using a first weight decay value, and
updating parameters of the plurality of self-attention neural network blocks using a second weight decay value,
wherein the first weight decay value is higher than the second weight decay value, and where in updating the parameters of the one or more output layers and the parameters of the one or more self-attention blocks comprises repeatedly performing the following:
computing, using one or more of the training examples, gradients of an objective function; and
applying an optimizer to the gradients to generate a gradient-based update to the parameters, wherein the optimizer makes use of a respective momentum value for each of the parameters, and where in the respective momentum values are stored with a reduced precision relative to the parameters.
14 . The system of claim 13 , wherein the first weight decay value is greater than or equal to .3 while the second weight decay value is less than .3.
15 . The system of claim 14 , wherein the first weight decay value is greater than or equal to 3.0.
16 . The system of claim 15 , wherein the second weight decay value is less than or equal to .1.
17 . The system of claim 13 , wherein each input element in the input sequence corresponds to a respective one of the image patches, and wherein the one or more output layers comprise:
an aggregation layer block that is configured to aggregate all of the output elements to generate an aggregated output element; and
one or more final output layers that are configured to generate the classification output from the aggregated output element.
18 . The system of claim 13 , the operations further comprising:
after training the vision Transformer on the first training data, training the plurality of self-attention neural network blocks jointly with a different set of one or more output layers on second training data to perform a different, downstream task.
19 . The system of claim 13 , wherein the respective momentum values are stored with half-precision relative to the parameters.
20 . The system of claim 13 , wherein the training further comprises:
during an initial phase of the training, linearly annealing a learning-rate for the parameters of the output layers and the self-attention blocks away from zero;
during a final phase of the training, linearly annealing the learning-rate toward zero; and
during a main phase of the training that is between the initial phase and the final phase, applying a schedule that prevents the learning-rate from reaching zero.