Training neural networks using adaptive matrix preconditioning with transformation invariance

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for training a neural network using an optimizer that uses a transformation-invariant matrix preconditioning method.

Description

Attorney Docket No. 56113-0823WO1TRAINING NEURAL NETWORKS USING ADAPTIVE MATRIX PRECONDITIONING WITH TRANSFORMATION INVARIANCECROSS-REFERENCE TO RELATED APPLICATION

[0001] This application claims priority’ to U.S. Provisional Application No. 63 / 702,640, filed on October 2, 2024. The disclosure of the prior application is considered part of and is incorporated by reference in the disclosure of this application.BACKGROUND

[0002] This specification relates to training neural networks.

[0003] Neural networks are machine learning models that employ one or more layers of nonlinear units to predict an output for a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as input to the next layer in the network, i.e.. the next hidden layer or the output layer. Each layer of the network generates an output from a received input in accordance with current value inputs of a respective set of parameters.SUMMARY

[0004] This specification describes a training system implemented as computer programs on one or more computers in one or more locations that trains a neural network to perform one or more machine learning tasks on a network input.

[0005] In general, one innovative aspect of the subject matter described in this specification can be embodied in a method performed by one or more computers for training a neural network that includes a plurality of layers that include a particular layer, wherein the particular layer is associated with a set of base parameters and a set of adapter parameters, wherein the set of adapter parameters is represented by a first parameter matrix and a second parameter matrix, and wherein the method may comprise: maintaining, for the particular layer of the neural network, a first unmagnified moment of gradients of a loss function with respect to the first parameter matrix and a second unmagnified moment of the gradients of the loss function with respect to the first parameter matrix; repeatedly performing training operations comprising: performing, using a plurality of training examples, a training step to obtain initial gradients of the loss function with respect to the first parameter matrix; determining unmagnified gradients of the loss function with respect to the first parameter matrix based on the initial gradientsAtorney Docket No. 56113-0823WO1 and the second parameter matrix; determining a first preconditioning matrix based on the second parameter matrix; updating the second unmagnified moment based on the first preconditioning matrix and the unmagnified gradients of the loss function with respect to the first parameter matrix; determining an unmagnified precondition step based on the unmagnified gradients of the loss function with respect to the first parameter matrix and the updated second unmagnified moment; updating the first unmagnified moment based on the first preconditioning matrix and the unmagnified precondition step; and updating the first parameter matrix based on the updated first unmagnified moment; and providing data specifying the adapted neural network for performing inference.

[0006] The training operations may comprise performing a polar decomposition of the second parameter matrix to generate an orthogonal matrix of the second parameter matrix and a symmetric positive semi-definite matrix of the second parameter matrix.

[0007] Determining the unmagnified gradients of the loss function may comprise: determining the unmagnified gradients of the loss function based on the initial gradients and an inverse of the upper triangular matrix of the second parameter matrix.

[0008] Determining the first preconditioning matrix may comprise: determining the first preconditioning matrix based on a transpose of the orthogonal matrix of the second parameter matrix.

[0009] The method may comprise maintaining, for the particular layer of the neural network, an escaped mass, wherein the training operations comprise updating the escaped mass based on the first preconditioning matrix.

[0010] Determining the unmagnified precondition step may comprise: determining the unmagnified precondition step based on the unmagnified gradients of the loss function, the updated second unmagnified moment, and the updated escaped mass.

[0011] Updating the first parameter matrix may comprise: updating the first parameter matrix based on the updated first unmagnified moment and the symmetric positive semi-defmite matrix.

[0012] The method may further comprise: maintaining, for the particular layer of the neural network, a first unmagnified moment of gradients of a loss function with respect to the second parameter matrix and a second unmagnified moment of the gradients of the loss function with respect to the second parameter matrix; and wherein performing the training step further comprises performing, using the plurality of training examples, the training step to obtain initial gradients of the loss function with respect to the second parameter matrix, and wherein the training operations may further comprise: determiningAttorney Docket No. 56113-0823WO1 unmagnified gradients of the loss function with respect to the second parameter matrix based on the initial gradients and the second parameter matrix; determining a second preconditioning matrix based on the first parameter matrix; updating the second unmagnified moment based on the second preconditioning matrix and the unmagnified gradients of the loss function with respect to the second parameter matrix; determining an unmagnified precondition step based on the unmagnified gradients of the loss function with respect to the second parameter matrix and the updated second unmagnified moment; updating the first unmagnified moment based on the second preconditioning matrix and the unmagnified precondition step; and updating the second parameter matrix based on the updated first unmagnified moment.

[0013] The training operations may comprise performing a polar decomposition of the first parameter matrix to generate an orthogonal matrix of the first parameter matrix and a symmetric positive semi-definite matrix of the first parameter matrix.

[0014] Determining the unmagnified gradients of the loss function may comprise: determining the unmagnified gradients of the loss function based on the initial gradients and an inverse of the upper triangular matrix of the first parameter matrix.

[0015] Determining the second preconditioning matrix may comprise: determining the second preconditioning matrix based on a transpose of the orthogonal matrix of the first parameter matrix.

[0016] The method may comprise maintaining, for the particular layer of the neural network, an escaped mass, wherein the training operations comprise updating the escaped mass based on the second preconditioning matrix.

[0017] Determining the unmagnified precondition step may comprise: determining the unmagnified precondition step based on the unmagnified gradients of the loss function, the updated second unmagnified moment, and the updated escaped mass.

[0018] Updating the second parameter matrix may comprise: updating the second parameter matrix based on the updated first unmagnified moment and the symmetric positive semi-definite of the first parameter matrix.

[0019] The neural network may be a generative neural network and the particular layer is an attention layer.

[0020] Other embodiments of this aspect include corresponding computer systems, apparatus, and computer programs recorded on one or more computer storage devices, each configured to perform the actions of the methods. A system of one or more computers can be configured to perform particular operations or actions by virtue ofAtorney Docket No. 56113-0823WO1 software, firmware, hardware, or any combination thereof installed on the system that in operation may cause the system to perform the actions. One or more computer programs can be configured to perform particular operations or actions by virtue of including instructions that, when executed by data processing apparatus, cause the apparatus to perform the actions.

[0021] Particular embodiments of the subject matter described in this specification can be implemented so as to realize one or more of the following advantages.

[0022] Low-rank adaption (LoRA) is a technique for adapting pre-trained neural networks to specific downstream tasks or domains. LoRA freezes the pre-trained parameter values of a neural network and injects a pair of trainable rank decomposition matrices into each of one or more layers of the neural network, thereby reducing the number of trainable parameters for downstream tasks or domains.

[0023] Nevertheless, training the rank decomposition matrices as part of the adaptation process may be problematic because applying traditional optimizers (e.g., an Adam optimizer or a RMSProp optimizer) or some existing second-order methods that use a matrix preconditioner (e.g., Shampoo or CASPR) to train the rank decomposition matrices may lead to updates that are not transformation invariant. In other words, two different pairs of trainable rank decomposition matrices that equivalently represent the same set of adapter parameters may receive different updates, even though both pairs represent the same set of adapter parameters. This inconsistency negatively impacts computational resource efficiency during training - for example when only one of the two rank decomposition matrices in each pair may be sufficiently updated, while the other matrix remains nearly unchanged.

[0024] The adaptive matrix preconditioning techniques described in this specification can be used to improve computational resource efficiency during the LoRA optimization. By achieving transformation invariance during the training of the rank decomposition matrices, the techniques ensure that any two equivalent pairs of trainable rank decomposition matrices will always have equivalent updates. This ensures that the same updates are applied to different pairs of rank decomposition matrices representing the same set of adapter parameters regardless of the specific factorization of the two rank decomposition matrices and avoids making imbalanced updates to either one of them, which would negatively impact computational resource efficiency.

[0025] Moreover, because the described techniques compute preconditioning matrices in the low rank space and rather than in the full rank space (as is the case of ShampooAttorney Docket No. 56113-0823WO1 preconditioning, for example), transformation invariance is achieved without substantially increasing time and memory complexity. Therefore, using the described techniques, the neural network can be adapted to a downstream task more quickly and using fewer computational resources, e.g., fewer processing cycles, than conventional approaches.

[0026] The details of one or more embodiments of the subj ect matter of this specification are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims.BRIEF DESCRIPTION OF THE DRAWINGS

[0027] FIG. 1A shows an example training system.

[0028] FIG. IB shows an example inference system.

[0029] FIG. 2 is a flow diagram of an example process for updating a first parameter matrix of an adaptation layer of a neural network.

[0030] FIG. 3 is a flow diagram of an example process for updating a second parameter matrix of an adaptation layer of a neural network.

[0031] FIGS. 4A-B show quantitative examples of the performance gains that can be achieved during training by using the optimizer described in this specification compared to existing optimizer.

[0032] Like reference numbers and designations in the vanous drawings indicate like elements.DETAILED DESCRIPTION

[0033] FIG. 1 A shows an example training system 100. The training system 100 is an example of a system implemented as computer programs on one or more computers in one or more locations that trains an adapted neural network 140 to perform one or more machine learning tasks on training data 120.

[0034] The adapted neural network 140 can be trained to perform any kind of machine learning task, i.e., can be configured to receive any kind of digital data input and to generate any kind of digital data output based on the input.

[0035] In some implementations, the adapted neural network 140 can be a generative neural network that can be configured through training to perform a generative task to generate, as output, data that includes, for example, text data, image data, video data, audio data, or multimodal data that includes data in two or more different modalities.Attorney Docket No. 56113-0823WO1

[0036] In some of these implementations, the adapted neural network 140 can be configured as, or include, a generative (large) language model, a foundation model, or a multi-modal model, e.g., a visual and language model.

[0037] In some implementations, the adapted neural network 140 can be an autoregressive neural network when the neural network auto-regressively generates an output sequence of tokens. More specifically, the auto-regressively generated output is created by generating each particular token in the output sequence conditioned on a current input sequence that includes any tokens that precede the particular token in the output sequence, i.e., the tokens that have already been generated for any previous positions in the output sequence that precede the particular position of the particular token.

[0038] For example, the adapted neural network 140 can be an auto-regressive Transformer-based neural network that includes (i) a plurality of attention blocks that each apply a self-attention operation and (ii) an output subnetwork that processes an output of the last attention block to generate the score distribution.

[0039] A neural network block refers to a group of one or more neural network layers in a neural network. For example, an attention block can include an attention layer and, optionally, a feed-forward layer and possibly other layers, e.g., residue connection layers, normalization layers, and so forth.

[0040] Generally, to apply the self-attention operation, each attention block uses one or more attention heads. Each attention head generates a set of queries, a set of keys, and a set of values, and then applies any of a variety of variants of query-key -value (QKV) attention, e.g., a dot product attention function or a scaled dot product attention function, using the queries, keys, and values to generate an output. Each query, key, value can be a vector that includes one or more vector elements. When there are multiple attention heads, the attention block then combines the outputs of the multiple attention heads, e.g., by concatenating the outputs and, optionally, processing the concatenated outputs through a linear layer.

[0041] The output block can generate an output of the adapted neural network 140. For example, the output can include a score distribution, e.g., a probability distribution, over tokens in a vocabulary of tokens. The score distribution assigns a respective score, e.g., a respective probability, to each token in the vocabulary of tokens.

[0042] The vocabulary of tokens can include any of a variety of tokens that represent text symbols or other symbols. For example, the vocabulary of tokens can include one orAttorney Docket No. 56113-0823WO1 more of: characters, sub-words, words, punctuation marks, numbers, or other symbols that appear in a corpus of natural language text and / or computer code.

[0043] Additionally or alternatively, the vocabulary of tokens can include tokens that can represent data other than text. For example, the vocabulary of tokens can include image tokens that represent a discrete set of image patch embeddings of an image that can be generated by an image encoder neural network based on processing the image patches of the image. As another example, the vocabulary of tokens can include audio tokens that represent code vectors in a codebook of a quantizer, e g., a residual vector quantizer.

[0044] In this example, the adapted neural network 140 can have any of a variety of Transformer-based neural network architectures. Examples of such Transformer-based neural network architectures include those described in PaLM: Scaling Language Modeling with Pathways, arXiv preprint arXiv:2204.02311; Rohan Anil, et al. Palm 2 technical report. arXiv preprint arXiv: 2305.10403, 2023; the Gemini Team, Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312. 11805 (2023);Gemini Team, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.05530 (2024); and Comanici. Gheorghe, et al. Gemini 2.5: Pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capabilities. arXiv preprint arXiv:2507.06261 (2025).

[0045] In some implementations, the adapted neural network 140 can be a non-auto- regressive Transformer neural network that similarly includes a plurality of attention blocks that each apply a self-attention operation, but generates an output sequence of tokens in a non-auto-regressive manner, i.e., simultaneously generates multiple tokens during each single forward pass.

[0046] Examples of such architectures in these cases include those described in Gu, J., Bradbury, J., Xiong, C., Li, V. O. K., and Socher, R. Non-autoregressive neural machine translation. In 6th International Conference on Learning Representations, ICLR 2018; Gu, J., Wang, C., and Zhao, J. Levenshtein transformer. In Wallach, H. M., Larochelle. H.. Beygelzimer, A.. d’Alche-Buc, F.. Fox, E. B.. and Garnett, R. (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019; and Gu, J. and Kong, X. Fully non- autoregressive neural machine translation: Tricks of the trade. In Zong, C., Xia, F., Li, W.. and Navigli, R. (eds.), Findings of the Association for Computational Linguistics: ACL / IJCNLP 2021.Attorney Docket No. 56113-0823WO1

[0047] In some implementations, the adapted neural network 140 can be a diffusion neural network that similarly includes a plurality of attention blocks that each apply a self-attention operation, and that generates an output data item, e.g., conditioned on a conditioning input, across multiple updating iterations by performing a reverse diffusion process. For example, the output data item can include text data, audio data, image data, or video data.

[0048] Examples of such architectures in these cases include those described in Saharia, Chitwan, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems 35 (2022): 36479- 36494; Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. arXiv:2006. 11239, 2020; and Yang Song and Stefano Ermon. Generative Modeling by Estimating Gradients of the Data Distribution. NeurlPS, 2019; and Zhao, Yang, et al. Mobilediffusion: Instant text-to-image generation on mobile devices. European Conference on Computer Vision. Cham: Springer Nature Switzerland, 2024.

[0049] Some examples of machine learning tasks, including generative tasks, that the adapted neural network 140 when implemented using one of the architectures described above or other known architectures can be configured to perform follow.

[0050] As one example, the task may be a neural machine translation task. For example, if the input to the neural network is a sequence of text, e.g., a sequence of words, phrases, characters, or word pieces, in one language, the output generated by the neural network may be a translation of the sequence of text into another language, i.e., a sequence of text in the other language that is a translation of the input sequence of text. As a particular example, the task may be a multi-lingual machine translation task, where a single neural network is configured to translate between multiple different source language - target language pairs. In this example, the source language text may be augmented with an identifier that indicates the target language into which the neural network should translate the source language text.

[0051] As another example, the task may be an audio processing task. For example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural netw ork may be a score for each of a set of pieces of text, each score representing an estimated likelihood that the piece of text is the correct transcript for the utterance. As another example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network can indicate whether a particular word or phrase (“hotword”) was spoken in the utterance. As anotherAtorney Docket No. 56113-0823WO1 example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network can identify the natural language in which the utterance was spoken.

[0052] As another example, the task can be a natural language processing or understanding task, e.g., an entailment task, a paraphrase task, a textual similarity task, a sentiment task, a sentence completion task, a grammaticality task, and so on, that operates on a sequence of text in some natural language.

[0053] As another example, the task can be a text to speech task, where the input is text in a natural language or features of text in a natural language and the network output is a spectrogram, a waveform, or other data defining audio of the text being spoken in the natural language.

[0054] As another example, the task can be a health prediction task, where the input is a sequence derived from electronic health record data for a patient and the output is a prediction that is relevant to the future health of the patient, e.g., a predicted treatment that should be prescribed to the patient, the likelihood that an adverse health event will occur to the patient, or a predicted diagnosis for the patient.

[0055] As another example, the task can be a text generation task, where the input is a sequence of text, and the output is another sequence of text, e.g., a completion of the input sequence of text, a response to a question posed in the input sequence, or a sequence of text that is about a topic specified by the first sequence of text. In this example, both the input sequence of text and the output sequence of text can include tokens from a vocabulary of text tokens that includes, e.g., one or more of characters, sub-words, words, punctuation marks, numbers, or other symbols that appear in a natural language or a computer language.

[0056] As a similar example, the task can be an automatic code generation task, where the input is a sequence of words, wordpieces or characters in a first natural language and the output is a sequence of tokens that represent instructions in a computer programming or markup language, or instructions for controlling an application program to perform a task e.g. build a data item such as an image or web page.

[0057] As a particular example of this, the input can represent a context input that includes a text description of a desired piece of code or a snippet of computer code in a programming language and the output can be an output sequence that includes computer code, e.g., a snippet of code that is described by the context input or a snippet of code that follows the context input in a computer program.Attorney Docket No. 56113-0823WO1

[0058] As another example, the input to the text generation task can be an input other than text, e.g., an image or audio, and the output sequence can be text that describes the input. In this example, both the input sequence of text and the output sequence of text can include tokens from a vocabulary of tokens that includes tokens that can represent data other than text, in addition to the text tokens mentioned above.

[0059] For example, the vocabulary of tokens can additionally include image tokens that represent a discrete set of image patch embeddings of an image that can be generated by an image encoder neural network based on processing the image patches of the image. As another example, the vocabulary of tokens can additionally include audio tokens that represent code vectors in a codebook of a quantizer, e g., a residual vector quantizer.

[0060] As another example, the task can be an image generation task, where the input is a conditioning input, e.g., text, a lower-resolution image, or a partial image, and the output is a sequence of intensity value inputs for the pixels of an image.

[0061] As another example, the task can be an image processing task. For example, the input can be the intensity values of the pixels of the image or an encoded representation of the intensity values of the pixels generated by an encoder neural network, and the network output can be (i) an image classification output that classifies the input image into one of a plurality of object categories (ii) an object detection output, i.e.. a sequence that specifies the coordinates of one or more bounding boxes in the image that are predicted to encompass objects or (iii) a segmentation output that classifies each pixel in the input image into one of a plurality of categories. As another example, the input can include the intensity7values of the pixels of the image or an encoded representation of the intensity' values of the pixels generated by an encoder neural network and optionally text, and the network output can be text that characterizes the image, e.g., captions the image or answers a question posed by the text in the input about the image.

[0062] As another example, the task can be an audio generation task, where the input is a conditioning input, e.g., text, an image, or context audio, and the output is a sequence of tokens that represents audio.

[0063] As another example, the task can be an audio processing task. For example, the input can include audio or an encoded representation of the audio generated by an encoder neural network, and the network output can be text or an image that characterizes the audio.Attorney Docket No. 56113-0823WO1

[0064] As another example, the task can be an agent control task, where the input is a sequence of observations or other data characterizing states of an environment and the output defines an action to be performed by the agent in response to the most recent data in the sequence. The agent can be, e.g., a real-world or simulated robot, a control system for an industrial facility', or a control system that controls a different kind of agent.

[0065] As another example, the task can be a genomics task, where the input is a sequence representing a fragment of a DNA sequence or other molecule sequence and the output is either an embedding of the fragment for use in a downstream task, e.g., by making use of an unsupervised learning technique on a data set of DNA sequence fragments, or an output for the downstream task. Examples of downstream tasks include promoter site prediction, methylation analysis, predicting functional effects of non-coding variants, and so on.

[0066] In some cases, the machine learning task is a combination of multiple individual machine learning tasks, i.e., the neural netw ork is configured to perform multiple different individual machine learning tasks, e.g., two or more of the machine learning tasks mentioned above. For example, the neural network can be configured to perform multiple individual natural language understanding tasks, with the network input including an identifier for the individual natural language understanding task to be performed on the network input.

[0067] In some cases, the machine learning task is a multi-modal processing task that requires processing multi-modal data. In general, multi-modal data is a combination of two or more different types of data, e.g., two or more of audio data, image data, text data, or graph data. As one example the multi-modal data may comprise audio-visual data, comprising a combination of pixels of an image or of video and audio data representing values of a digitized audio waveform. As another example the multi-modal data may comprise a combination of i) text data representing text in a natural language and ii) pixels of an image or of video or audio data representing values of an audio w aveform. Optionally, but not necessarily, the different types of data may represent the same or overlapping objects using the different modalities (types), and when processing multimodal data the data may be mapped into a common embedding space.

[0068] As a particular example, the task is a multi-modal processing task that requires processing both text and image inputs, so that the neural network includes both a computer vision neural network and a text processing neural network. That is, the target output to be generated by the computer vision neural network for a given image dependsAttorney Docket No. 56113-0823WO1 on one or more outputs generated by the text processing neural network for one or more corresponding text inputs (and vice versa). Examples of such tasks include openvocabulary image classification, open-vocabulary object detection, image captioning, text-based image search, image-based retrieval, and so on.

[0069] More generally, the multi-modal processing task may correspond to any of the tasks previously described for any of the types of data making up the multi-modal combination. For example, an accuracy of the previously described tasks may be increased when the task is applied to multi-modal data combining the data for which the task has been previously described and another type of data. For example detection or classification of an object or event may be improved when data of multiple different types (modalities) is processed. As another example, the quality (e.g.. accuracy, fidelity, or intelligibility) of a generated image, video, or audio may be improved when data of multiple different types (modalities) is processed.

[0070] In some cases, once trained, the neural network can perform tasks that it was not explicitly trained to perform. For example the neural network can perform translation tasks (provided that the training corpus included words in different languages), generative tasks, and many other tasks.

[0071] In these cases, the neural network can be made to perform a particular task by providing a natural language description of the desired response as a part of the input or “prompt”. The prompt may be a few-shot prompt where a few, e.g.. 1 to 10, examples of a query and an example output are provided in the text prior to the actual query.

[0072] Additional description of generative tasks that the neural network (when configured as a “generative” neural network) can perform are discussed below.

[0073] Generally, the generative neural network is configured to process a conditioning input (“prompt”) to generate a data item. The data item can include data in any of a variety’ of modalities, e.g., text data, image data, video data, or audio data. Generally, the data item represents a response to the conditioning input which may be, e.g. a “prompt” for the generative neural network. For example, the conditioning input can characterize one or more desired properties for the generated data item.

[0074] In some implementations the generative neural network generates an output token sequence from an input token sequence including the conditioning input. The generative neural network may then be configured to process the input token sequence to generate for each position in the output token sequence, a respective score for each tokenAttorney Docket No. 56113-0823WO1 in a vocabulary of output tokens, that is used to select an output token for the output token sequence.

[0075] In some implementations the tokens can represent text, e.g., words, wordpieces or characters, in a natural or computer language. For example, text may be received, e.g., as a series of encoded characters, e.g. UTF-8 encoded characters; such “characters” can include Chinese and other similar characters, as well as logograms, syllabograms and the like. A text encoder, i.e. a tokenizer, can process a sequence of text to represent the text as a series of text tokens from a vocabulary of text tokens, e.g. that each represent words, wordpieces or characters in a natural or computer language. The computer language may be any formal language used to communicate with a computer, e.g. a markup language, or a command or configuration language, or a data exchange language such as JSON, or a programming language. The tokenizer can, e.g., implement BPE (Byte Pair Encoding) or Wordpiece tokenization. Optionally the text can be obtained from audio data representing speech; the output tokens may be converted into audio data that represent speech corresponding to the text.

[0076] Also, or instead the tokens may represent an image. For example, a set (sequence) of input or output tokens can represent an image. Each image token may comprise a block encoding of values of the pixels in a different region of an image that maps a set of values of the pixels to a respective image token. The block encoder may comprise a neural network, e.g. having one or more (self-)attention layers, such as a Transformer neural network.

[0077] Also, or instead the tokens may represent an audio waveform. For example, a set (sequence) of input or output tokens can represent audio data representing a waveform e.g. instantaneous audio amplitude values or time-frequency audio data. Each image token may comprise a block encoding of the audio waveform in a different time segment of the audio that maps a set of values representing the audio waveform to a respective image token. The block encoder may comprise a neural network, e.g. having one or more (self-)attention layers, such as a Transformer neural network.

[0078] In a multimodal system audio data or an image may be flagged by a start-of- audio token or start-of-image token.

[0079] In some implementations the generative neural network can be a multimodal network that is configured to process a conditioning input comprising one or more of text data, audio data defining an audio signal (e.g. as amplitude values of the audio signal or as a time-frequency representation of the audio signal), or a still or moving image (e.g. asAttorney Docket No. 56113-0823WO1 image pixel values), to generate a data item that can similarly comprise text data, audio data, or a still or moving image.

[0080] For example, the conditioning input may comprise text and the data item may comprise an image or an audio signal that represents speech an image generated in response to the text, e.g. described by the text. Also, or instead the conditioning input may comprise an audio signal that represents speech, or an image, and the data item may comprise text, e.g. that describes the conditioning input.

[0081] As another example the conditioning input may comprise an observation, e.g. of a real world environment, e.g. from sensor such as a camera or other image sensor; and optionally additional information such as information defining a particular task to be deformed. The output data item may comprise agent control data that defines one or more actions to be performed by an agent, e.g. by a mechanical agent such as a robot or autonomous vehicle, to perform a task. The reward model(s) may, e.g., define a preferred trajectory of motion of the mechanical agent in the (real-world) environment.

[0082] Prior to training the adapted neural network 140, the training system 100 obtains data specifying a trained neural network 1 10. For example, the data specifying the trained neural network 110 can include data specifying the trained values of the base parameters of the trained neural network 110 and, optionally, the data specifying the architecture of the trained neural network 110.

[0083] The base parameters of the trained neural network 110 include a set of base parameters 1 14 of an adaptation layer 130 that represents the weights and, optionally, biases of the adaptation layer 130 of the trained 1 network 110.

[0084] The adaptation layer 130 can be any layer, e.g., an attention layer, a feedforward layer, an embedding layer, or a convolutional layer, of the trained neural network 110.

[0085] Although one adaptation layer 130 is depicted in FIG. 1A for convenience, the adapted neural network 140 can include many adaptation layers, and the techniques described in this specification as being performed with respect to the adaptation layer 130 apply to each of these adaptation layers.

[0086] Much like the adapted neural network 140, the trained neural network 110 can have any neural network architecture, e.g., one of the example generative neural network architectures mentioned above.

[0087] The trained neural network 110 can have been pre-trained, by the training system 100 or a separate training system, on training data that includes a largeAttorney Docket No. 56113-0823WO1 unsupervised dataset, and thus the training system 100 can obtain data specifying the trained neural network 110 either from a local storage device of the training system 100 or from a remote storage device that is physically remote from the training system 100.

[0088] For example, the trained neural network 110 can have been pre-trained on a next token prediction task, e.g., a task that requires predicting, given a current sequence of tokens, the next token that follows the current sequence in the training data.

[0089] As a particular example, the trained neural network 110 can have been pretrained on a maximum-likelihood objective on a large dataset of text in one or more natural languages, e.g., text that is publicly available from the Internet or another text corpus, a large dataset of computer code in one or more programming languages, e.g., Python, C++, C#, Java, Ruby, PHP, and so on, e.g., computer code that is publicly available from the Internet or another code repository, a large dataset of audio samples, e.g., audio recordings or waveforms that represent the audio recordings, a large dataset of images where each image includes an array of pixels, a large dataset of videos where each video includes a temporal sequence of frames, or a large multi-modal dataset that includes a combination of two or more of these datasets.

[0090] Optionally, in some implementations, the pre-training of the trained neural network 110 also involves supervised fine-tuning (SFT), direct preference optimization (DPO), reinforcement learning from human feedback (RLHF), or another training technique, to train the trained neural network 110 on specific machine learning tasks, e.g., one or more of the tasks mentioned above and possibly other tasks, using labeled training data on a supervised training objective.

[0091] Having obtained the data specifying the trained neural network 110, the training system 100 generates the adapted neural network 140 based on adding, to the adaptation layer 130 of the trained neural network 110, a first parameter matrix 1 16 and a second parameter matrix 118, and then training the adapted neural network 140 using the training data 120 to learn the trained values of the first and second parameter matrices 116, 118 that have been added to the adaptation layer 130.

[0092] In doing so. training system 100 adapts (customizes) the trained neural network to a specific downstream, e.g., any one of the machine learning task tasks mentioned above and possibly other tasks, adapts (customizes) the trained neural network to a specific domain, e.g., a particular entity that provides prompts specific the particular entity that make the adapted neural network perform tasks (tasks on which the trainedAttorney Docket No. 56113-0823WO1 neural network may or may not have seen during its pre-training), or adapts (customizes) the trained neural network to another downstream use case / application.

[0093] For the adaptation layer 130, the first and second parameter matrices 116, 1 18 collectively represent (approximate) a larger set of adapter parameters that correspond to newly added parameters of the adaptation layer 130 of the adapted neural network 140 that will be used in addition to the set of base parameters 114 when performing one or more machine learning tasks.

[0094] The first and second parameter matrices 116, 118 are smaller matrices that, when combined together through multiplication, can approximately recover the set of adapter parameters of the adaptation layer 130. Each smaller matrix contains much fewer parameters than an original matrix representing the set of adapter parameters that would be added to the adaptation layer 130.

[0095] For example, suppose VF G JRmxnis a weight matrix representing the set of base parameters 114 of the adaptation layer 130, the training system 100 can add the first parameter matrix 116 (a rank decomposition matrix A) and the second parameter matrix 118 (a rank decomposition matrix B) that collectively represent a larger set of adapter parameters (represented by adapter matrix Z):

[0096] Because the number of new ly added parameters that needs to be learned is therefore reduced, computing resources which would be spent on neural network training or training data collection can be conserved, thereby reducing the consumption of computational resources such as processor usage, memory' usage, and / or network bandwidth.

[0097] The training data 120 includes multiple training examples which, in turn, each include a training input and, optionally, a corresponding target output for the training input for the machine learning task on which the adapted neural netw ork 140 is being trained to perform, i.e., a target output to be generated by the adapted neural netw ork 140 by processing the training input.

[0098] Generally, the training system 100 trains the adapted neural network 140 to minimize a loss function for the machine learning task.Attorney Docket No. 56113-0823WO1

[0099] For example, the loss function can be a loss function used in supervised fine- tuning (SFT), direct preference optimization (DPO), reinforcement learning from human feedback (RLHF), or another training technique.

[0100] More generally, the loss function can be any appropriate loss function for the machine learning task. The loss function includes one or more terms that measure, for each training input, the quality of a training output for the training input, e.g.. relative to a respective target output for the training input (when it is included in the training example that includes the training input).

[0101] For each training input, the training output is generated by the adapted neural network 140 by performing one or more forward passes through the plurality of layers of the adapted neural network 140. The plurality of layers include the adaption layer 130. The adaption layer 130 is associated with the set of base parameters 114 and a set of adapter parameters. The set of adapter parameters is represented (approximated) by the first parameter matrix 116 and the second parameter matrix 118.

[0102] For example, the one or more terms can be cross entropy loss terms, mean squared error loss terms, negative log likelihood loss terms, and so on. The loss function can also include other terms, e g., regularization terms, auxiliary loss terms, unsupervised learning loss terms, and so on, that do not depend on the target outputs for the training inputs (e.g., when they are not included in the training examples).

[0103] The training system 100 performs the training over a plurality of training steps. At each training step, the training system 100 updates the first parameter matrix 116, the second parameter matrix 118, or both using a plurality of training examples (a “batch” or a “mini-batch” of training examples) sampled from the training data 120.

[0104] Thus, by repeatedly performing training steps, the training system 100 repeatedly updates the values of the first parameter matrix 116 and the second parameter matrix 118 of the adapted neural network 140 to determine trained values of the first parameter matrix 116 and the second parameter matrix 118 that will cause the adapted neural network 140 to perform well on the machine learning task.

[0105] In some implementations, training the adapted neural network 140 to leam the trained values of the first parameter matrix 116 and the second parameter matrix 118 can include holding the pre-trained values of the set of base parameters 114 fixed.

[0106] At each training step, the training system 100 computes, using the plurality of training examples, a gradient of the loss function for the machine learning task withAttorney Docket No. 56113-0823WO1 respect to each of the parameters of the first parameter matrix 116 and the second parameter matrix 118 of the adaptation layer 130.

[0107] The training system 100 then uses an optimizer 150 to determine an update to the values of the parameters of the first parameter matrix 116 and the second parameter matrix 118 from the gradients.

[0108] The optimizer 150 overcomes limitations faced by traditional optimizers and some existing second-order methods that use a matrix preconditioner, which may result in imbalanced updates to the first and second parameter matrices 11 , 1 18 due to their lack of transformation invariance, by employing an adaptive, transformation-invariant matrix preconditioning method that efficiently balances updates to both the first and second parameter matrices 116, 118.

[0109] The optimizer 150 ensures that the same updates are applied to the first and second parameter matrices 116, 118 representing the same set of adapter parameters regardless of the specific factorization of the two parameter matrices.

[0110] The optimizer 150 achieves improved computational resource efficiency during training because it addresses the unbalanced training dynamics faced by traditional optimizers, where one of the two parameter matrices receives significantly larger updates than the other, leading to inefficient learning and a suboptimal trained instance of the adapted neural network 140.

[0111] To achieve transformation invariance and thereby improve computational resource efficiency during training, the optimizer 150 uses adaptive matrix preconditioning techniques to update the first and second parameter matrices 116, 118.

[0112] To use the adaptive matrix preconditioning techniques as will be explained further below, the optimizer 150 maintains (stores) two unmagnified moments of gradients of the loss function with respect to each parameter matrix of the adaptation layer 130. These unmagnified moments are used by the optimizer 150 to utilize adaptive preconditioning, thereby maintaining transformation invariance during the training.

[0113] Prior to the first training step, the two unmagnified moments can be initialized to have predetermined initial values (e.g., zero) and then iteratively updated over the course of the plurality of training steps.

[0114] The optimizer 150 maintains a first unmagnified moment 152 of gradients of the loss function with respect to the first parameter matrix 116.

[0115] The optimizer 150 maintains a second unmagnified moment 154 of the gradients of the loss function with respect to the first parameter matrix 116.Attorney Docket No. 56113-0823WO1

[0116] The optimizer 150 maintains a first unmagnified moment 156 of gradients of a loss function with respect to the second parameter matrix 118.

[0117] The optimizer 150 maintains a second unmagnified moment 158 of the gradients of the loss function with respect to the second parameter matrix 118.

[0118] A moment is called “unmagnified’' when it is a moment of unmagnified gradients. Unmagnified gradients rely solely on the column spaces of the first parameter matrix 116 and the second parameter matrix 118, thus remaining invariant to transformations of the set of adapter parameters represented by the first and second parameter matrices 116, 118.

[0119] In contrast, gradients used by existing optimizers (e.g., stochastic gradient decent or Adam optimizer) are dependent on both the basis and the magnitude of the first and second parameter matrices 116, 118. Because of their dependence on both the magnitude and basis, when used with the existing optimizers, these gradients are not transformation invariant, leading to suboptimal model updates.

[0120] Having trained the adapted neural network 140 on the training data 120, the training system 100 can provide data specifying the adapted neural network 140. e.g., the trained values of at least a subset of the parameters of the adapted neural network and, optionally, data specifying at least a portion of the architecture of the adapted neural network to another system, e.g., an inference system, e.g.. the inference system 170 of FIG. IB, or another system in a deployment environment, for use in performing inference, i.e., performing one or more of the machine learning tasks mentioned above.

[0121] Instead of or in addition to providing the data specifying the adapted neural network, the training system 100 can use the adapted neural network to perform machine learning tasks, e.g., perform generative tasks, e.g.. process received prompts to generate, as output, data that includes, for example, text data, image data, video data, audio data, or multimodal data that includes data in two or more different modalities.

[0122] FIG. IB shows an example inference system 170. The inference system 170 is an example of a system implemented as computer programs on one or more computers in one or more locations that implements the adapted neural network 140 to perform inference.

[0123] The inference system 170 performs one or more machine learning tasks using the base parameters of the adapted neural network 140, e.g., the set of base parameters 114 of the adaptation layer 130 of the adapted neural network 140, and an approximation of the adapter parameters of the adapted neural network 140, e.g., the first parameterAttorney Docket No. 56113-0823WO1 matrix 116 (with rank = r) and the second parameter matrix 118 (with rank = r) that approximate the set of the adapter parameters of the adaptation layer 130.

[0124] In some implementations, the inference system 170 stores the base parameters together with the approximation of the adapter parameters. For example, the base parameters can be stored together with the approximation of the adapter parameters in one or more storage devices local to the inference system.

[0125] In some implementations, as indicated by the dashed line in FIG. IB, the inference system 170 stores the base parameters 1 16 separate from the approximation of the adapter parameters. For example, the approximation of the adapter parameters can be stored in one or more storage devices local to the inference system, while the base parameters can be stored in one or more storage devices remote from the inference system.

[0126] In some implementations, the inference system 170 stores the respective approximation of each of multiple sets of adapter parameters that are specific to different machine learning tasks, different domains, or different downstream use cases / appli cations. In other words, the inference system 170 stores, for the same base parameters of the adapted neural network 140, respective approximations of multiple sets of adapter parameters that can be user together with the same base parameters to perform different machine learning tasks, perform the same tasks in different domains, or in different downstream use cases / applications.

[0127] In any implementation, the first parameter matrix 1 16 and the second parameter matrix 118 that approximate the set of the adapter parameters of the adaptation layer 130 can have been trained in a computationally efficient manner by using the adaptive matrix preconditioning techniques described in this specification.

[0128] FIG. 2 is a flow- diagram of an example process 200 for updating a first parameter matrix (matrix A) of an adaptation layer of a neural network. For example, the neural netw ork can correspond to the adapted neural network 140 of FIG. 1A. The neural network includes a plurality of layers. The plurality of layers include the adaptation layer. The adaptation layer is associated with a set of base parameters and a set of adapter parameters. The set of adapter parameters is represented (approximated) by the first parameter matrix and a second parameter matrix.

[0129] For convenience, the process 200 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system, e.g., the training system 100 of FIG. 1, appropriately programmed, canAtorney Docket No. 56113-0823WO1 perform the process 200. An iteration of the process 200 can be performed at each training step during the training of the adapted neural network.

[0130] The system maintains, for the adaptation layer of the neural network, (i) a first unmagnified moment of gradients of a loss function with respect to the first parameter matrix, (ii) a second unmagnified moment of the gradients of the loss function with respect to the first parameter matrix, and (iii) an escaped mass.

[0131] The system performs, using a plurality of training examples, a training step to obtain initial gradients of the loss function with respect to the first parameter matrix (step 202).

[0132] Each training example includes a training input and, optionally, a target output for each training input. The system will generally obtain different training examples at different training steps, e.g., by sampling a fixed number of examples from a larger set of training data at each iteration.

[0133] For example, the system can perform a forward pass through the neural network using the training examples and then perform a backward pass through the neural network to compute the initial gradients through backpropagation.

[0134] As described above, the loss function can be any appropriate loss function for the machine learning task. Generally, the loss function includes one or more terms that measure, for each training input, the quality of a training output for the training input generated by performing a forward pass through the neural network, e.g.. relative to a respective target output for the training input. For example, the one or more terms can be cross entropy loss terms, mean squared error loss terms, negative log likelihood loss terms, and so on.

[0135] The loss function can also include other terms, e.g., regularization terms, auxiliary loss terms, unsupervised learning loss terms, and so on, that do not depend on the target outputs for the training inputs.

[0136] The system determines unmagnified gradients of the loss function with respect to the first parameter matrix based on the initial gradients and the second parameter matrix (step 204). These unmagnified gradients rely solely on the column spaces of the first and second parameter matrices, maintaining invariance to transformations.

[0137] To do this, the system performs a matrix decomposition of the second parameter matrix to generate a first sub-matrix of the second parameter matrix and a second sub-matrix of the second parameter matrix, and then determines the unmagnified gradients of the loss function based on computing a product of (i) the initial gradients andAttorney Docket No. 56113-0823WO1(ii) an inverse or a pseudo-inverse of the second sub-matrix of the second parameter matrix. For example, the system determines:

[0138] In some implementations, the matrix decomposition is a polar decomposition. For example, polar decomposition can be computed by singular value decomposition (SVD). In these implementations, the first sub-matrix of the second parameter matrix is an orthogonal matrix of the second parameter matrix, representing a rotation or reflection, and the second sub-matrix of the second parameter matrix is a positive semi-definite symmetric matrix of the second parameter matrix, representing a scaling or stretching.For example, the system determines:

[0139] In some implementations, the matrix decomposition is a QR decomposition. For example, QR decomposition can be computed by a Gram-Schmidt process.Householder transformations, or Givens rotations. In these implementations, the first submatrix of the second parameter matrix is an orthogonal matrix of the second parameter matrix, representing a rotation or reflection, and the second sub-matrix of the second parameter matrix is an upper triangular matrix of the second parameter matrix, representing a scaling or shearing.

[0140] The system determines a first preconditioning matnx based on computing a product of (i) a transpose of the first sub-matrix of the second parameter matrix and (ii) a historic first sub-matrix of the second parameter matrix determined in a previous iteration of the process 200 (step 206). The product is then used as the first preconditioning matrix.For example, the system computes:j.

[0141] The system updates the second unmagnified moment based on the first preconditioning matrix and the unmagnified gradients of the loss function with respect to the first parameter matrix (step 208).

[0142] To do this, the system computes a first product of (i) the first preconditioning matrix, (ii) a historic second unmagnified moment determined in the previous iteration of the process 200, and (iii) a transpose of the first preconditioning matrix.

[0143] The system computes a second product of (i) a transpose of the unmagnified gradients of the loss function with respect to the first parameter matrix, (ii) the unmagnified gradients of the loss function with respect to the first parameter matrix, and(iii) a constant (e.g., one) over a scaling factor m, where m represents the number of rows in the adapter matrix representing the set of adapter parameters.Attorney Docket No. 56113-0823WO1

[0144] The system computes a sum of the first product and the second product. The sum is then used as the updated second unmagnified moment. For example, the system computes:

[0145] The system updates the escaped mass based on the first preconditioning matrix(step 210). The escaped mass compensates for the potential loss of information when the updated second unmagnified moment is transformed to a new basis, i.e., when the orthogonal representations (bases) of the second parameter matrix undergoes a change, as iterations of the process 200 continue.

[0146] To do this, the system computes an information loss function that is in the form of:where El is the historic second unmagnified moment determined in the previous iteration of the process 200, E2 is a product of (i) the first preconditioning matrix, (ii) the historic second unmagnified moment determined in the previous iteration of the process 200, and (iii) the transpose of the first preconditioning matrix, and U represents the first sub-matrix of the second parameter matrix.

[0147] The system adds an output of the information loss function to a historic escaped mass determined in the previous iteration of the process 200. The sum is then used as the updated escaped mass. For example, the system computes:

[0148] The system determines an unmagnified precondition step based on the unmagnified gradients of the loss function with respect to the first parameter matrix, the updated second unmagnified moment, and the updated escaped mass (step 212). This unmagnified precondition step incorporates adaptive preconditioning using the updated second unmagnified moment and the updated escaped mass.

[0149] To do this, the system computes a sum of (i) the updated second unmagnified moment and (ii) a product of the updated escaped mass and an identity matrix of appropriate dimensions.

[0150] The system computes an inverse square root of the sum.

[0151] The system computes a product of (i) the inverse square root of the sum and (ii) the unmagnified gradients of the loss function with respect to the first parameterAttorney Docket No. 56113-0823WO1 matrix. The product is then used as the unmagnified precondition step. For example, the system

[0152] The system updates the first unmagnified moment based on the first preconditioning matrix and the unmagnified precondition step (step 214). That is, the system applies an exponential moving average (EMA) update to the first unmagnified moment to account for basis changes.

[0153] To do this, the system computes a first product of (i) a momentum hyperparameter for the first unmagnified moment, (ii) a historic first unmagnified moment determined in the previous iteration of the process 200, and (iii) the transpose of the first preconditioning matrix.

[0154] The system computes a second product of (i) a constant (e.g., one) minus the momentum hyperparameter for the first unmagnified moment and (ii) the unmagnified precondition step. The momentum hyperparameter for the first unmagnified moment is a tunable hyperparameter which can take any value within a predetermined range, e.g., between 0 and 1.

[0155] The system computes a sum of the first product and the second product. The sum is then used as the updated first unmagnified moment. For example, the system computes

[0156] The system updates the first parameter matrix based on the updated first unmagnified moment and the second sub-matrix of the second parameter matrix (step 216).

[0157] To do this, the system computes a product of (i) a learning rate, (ii) the updated first unmagnified moment, and (iii) an inverse or a pseudo-inverse of a transpose of the second sub-matrix of the second parameter matrix. The inverse or a pseudo-inverse of the transpose of the second sub-matrix of the second parameter matrix adjusts the magnitude for transformation invariance. In some implementations, the system uses different learning rates in different iterations of the process 200.

[0158] The system computes a difference of (i) a historic first parameter matrix determined in the previous iteration of the process 200 and (ii) the product. The difference is then used as the updated first parameter matrix. For example, the system computes:Attorney Docket No. 56113-0823WO1

[0159] An example algorithm for updating a first parameter matrix of an adaptation layer of a neural network is shown below.Algorithm1: Initialize: unmagnified first and second moment AT-- 02: for t .1 . . . T do3: Compute the gradient V A< :11 1 : end for

[0160] In the example algorithm shown above, lines 3-5 correspond to step 202 of FIG. 2, lines 4-5 correspond to step 204 of FIG. 2, line 5 corresponds to step 206 of FIG.2, line 6 corresponds to step 208 of FIG. 2, line 7 corresponds to step 210 of FIG. 2, line 8 corresponds to step 212 of FIG. 2, line 9 corresponds to step 214 of FIG. 2, and line 10 corresponds to step 216 of FIG. 2. The first parameter matrix is denoted as matrix A. and the second parameter matrix is denoted as matrix B. Collectively, matrices A and B are referred to as LoRA factors.

[0161] FIG. 3 is a flow diagram of an example process 300 for updating a second parameter matrix (matrix B) of an adaptation layer of a neural network. For example, the neural network can correspond to the adapted neural network 140 of FIG. 1 A. The neural network includes a plurality of layers. The plurality of layers include the adaptation layer. The adaptation layer is associated with a set of base parameters and a set of adapter parameters. The set of adapter parameters is represented (approximated) by a first parameter matrix and the second parameter matrix.

[0162] For convenience, the process 300 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system, e.g., the training system 100 of FIG. 1, appropriately programmed, can perform the process 300.

[0163] The system maintains, for the adaptation layer of the neural network, (i) a first unmagnified moment of gradients of a loss function with respect to the second parameter matrix, (ii) a second unmagnified moment of the gradients of the loss function with respect to the second parameter matrix, and (iii) an escaped mass.Attorney Docket No. 56113-0823WO1

[0164] The system performs, using a plurality of training examples, a training step to obtain initial gradients of the loss function with respect to the second parameter matrix (step 302).

[0165] The system determines unmagnified gradients of the loss function with respect to the second parameter matrix based on the initial gradients and the second parameter matrix (step 304).

[0166] To do this, the system performs a matrix decomposition of the first parameter matrix to generate a first sub-matrix of the first parameter matrix and a second sub-matrix of the first parameter matrix, and then determines the unmagnified gradients of the loss function based on computing a product of (i) the initial gradients and (ii) an inverse or a pseudo-inverse of the second sub-matrix of the first parameter matrix.

[0167] In some implementations, the matrix decomposition is a polar decomposition. For example, polar decomposition can be computed by singular value decomposition (SVD). In these implementations, the first sub-matrix of the first parameter matrix is an orthogonal matrix of the first parameter matrix, representing a rotation or reflection, and the second sub-matrix of the first parameter matrix is a positive semi-definite symmetric matrix of the first parameter matrix, representing a scaling or stretching.

[0168] In some implementations, the matrix decomposition is a QR decomposition. For example, QR decomposition can be computed by a Gram-Schmidt process.Householder transformations, or Givens rotations. In these implementations, the first submatrix of the first parameter matrix is an orthogonal matrix of the first parameter matrix, representing a rotation or reflection, and the second sub-matrix of the first parameter matrix is an upper triangular matrix of the first parameter matrix, representing a scaling or shearing.

[0169] The system determines a second preconditioning matrix based on computing a product of (i) a transpose of the first sub-matrix of the first parameter matrix and (ii) a historic first sub-matrix of the first parameter matrix determined in a previous iteration of the process 300 (step 306). The product is then used as the second preconditioning matrix.

[0170] The system updates the second unmagnified moment based on the second preconditioning matrix and the unmagnified gradients of the loss function with respect to the second parameter matrix (step 308).

[0171] To do this, the system computes a first product of (i) the second preconditioning matrix, (ii) a historic second unmagnified moment determined in theAttorney Docket No. 56113-0823WO1 previous iteration of the process 300, and (iii) a transpose of the second preconditioning matrix.

[0172] The system computes a second product of (i) a transpose of the unmagnified gradients of the loss function with respect to the second parameter matrix, (ii) the unmagnified gradients of the loss function with respect to the second parameter matrix, and (iii) a constant (e.g., one) over a scaling factor n, where n represents the number of columns in the adapter matrix representing the set of adapter parameters.

[0173] The system computes a sum of the first product and the second product. The sum is then used as the updated second unmagnified moment.

[0174] The system updates the escaped mass based on the second preconditioning matrix (step 310). To do this, the system computes an information loss function that is in the form o dx(Eiwhere El is the historic second unmagnified moment determined in the previous iteration of the process 300, E2 is a product of (i) the second preconditioning matrix, (ii) the historic second unmagnified moment determined in the previous iteration of the process 300, and (iii) the transpose of the second preconditioning matrix, and U represents the first sub-matrix of the first parameter matrix.

[0175] The system adds an output of the information loss function to a historic escaped mass determined in the previous iteration of the process 300. The sum is then used as the updated escaped mass.

[0176] The system determines an unmagnified precondition step based on the unmagnified gradients of the loss function with respect to the second parameter matrix, the updated second unmagnified moment, and the updated escaped mass (step 312).

[0177] To do this, the system computes a sum of (i) the updated second unmagnified moment and (ii) a product of the updated escaped mass and an identity matrix of appropriate dimensions.

[0178] The system computes an inverse square root of the sum.

[0179] The system computes a product of (i) the inverse square root of the sum and (ii) the unmagnified gradients of the loss function with respect to the second parameter matrix. The product is then used as the unmagnified precondition step.

[0180] The system updates the first unmagnified moment based on the second preconditioning matrix and the unmagnified precondition step (step 314).Atorney Docket No. 56113-0823WO1

[0181] To do this, the system computes a first product of (i) a momentum hyperparameter for the first unmagnified moment, (ii) a historic first unmagnified moment determined in the previous iteration of the process 300, and (iii) the transpose of the second preconditioning matrix.

[0182] The system computes a second product of (i) a constant (e.g., one) minus the momentum hyperparameter for the first unmagnified moment and (ii) the unmagnified precondition step. The momentum hyperparameter for the first unmagnified moment is a tunable hyperparameter which can take any value within a predetermined range, e.g., between 0 and 1.

[0183] The system computes a sum of the first product and the second product. The sum is then used as the updated first unmagnified moment.

[0184] The system updates the second parameter matrix based on the updated first unmagnified moment and the second sub-matrix of the second parameter matrix (step 316).

[0185] To do this, the system computes a product of (i) a learning rate, (ii) the updated first unmagnified moment, and (iii) an inverse or a pseudo-inverse of a transpose of the second sub-matrix of the first parameter matrix. In some implementations, the system uses different learning rates in different iterations of the process 300.

[0186] The sy stem computes a difference of (i) a historic second parameter matrix determined in the previous iteration of the process 300 and (ii) the product. The difference is then used as the updated second parameter matrix.

[0187] To train the neural network, the system can repeatedly perform iterations of the process 200 to repeatedly update the first parameter matrix and repeatedly perform iterations of the process 300 to repeatedly update the second parameter matrix of each of the one or more adaptation layers of the neural network until a termination criterion has been satisfied, e.g., until a threshold number of iterations of the processes 200 and 300 have been performed, until a threshold amount of wall clock time has elapsed, or until the values of the first and second parameter matrices have converged.

[0188] To train the neural network, iterations of the process 200 can be performed alongside iterations of the process 300. In some implementations, iterations of the process 300 can be executed in parallel or interleaved with iterations of the process 200. For example, an iteration of the process 300 can be executed in parallel with or immediately following an iteration of process 200 together as part of the same training step for training the neural network, based on the same plurality of training examples.Atorney Docket No. 56113-0823WO1

[0189] FIG. 4A shows a quantitative example of the performance gains that can be achieved during training by using the optimizer described in this specification compared to existing optimizer.

[0190] In particular, FIG. 4A shows training loss curves of different instances of a neural network having the same architecture when trained using different optimizers on the Super-Natural instruction dataset (described in Yizhong Wang, et al. Super- naturalinstructions: Generalization via declarative instructions on 1600+ nip tasks. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 5085-5109, 2022). The Super-Natural instruction dataset includes a collection of 1600+ NLP tasks, including both text-based classification and generative tasks. It will be appreciated that the neural network trained using the described optimizer (LoRA-RITE) has the lowest training loss.

[0191] FIG. 4B shows another quantitative example of the performance gains that can be achieved during training by using the optimizer described in this specification compared to existing optimizer.

[0192] In particular, FIG. 4A shows training loss curves of different instances of a neural network having the same architecture when trained using different optimizers on the OpenBookQA dataset (described in Todor Mihaylov, et al. Can a suit of armor conduct electricity? a new dataset for open book question answering. In EMNLP, 2018.). It will be appreciated that the neural network trained using the described optimizer (LoRA-RITE) has the lowest training loss.

[0193] In this specification, the term “configured” is used in relation to computing systems and environments, as well as computer program components. A computing system or environment is considered “configured” to perform specific operations or actions when it possesses the necessary software, firmware, hardware, or a combination thereof, enabling it to cany' out those operations or actions during operation. For instance, configuring a system might involve installing a software library with specific algorithms, updating firmware with new instructions for handling data, or adding a hardware component for enhanced processing capabilities. Similarly, one or more computer programs are “configured” to perform particular operations or actions when they contain instructions that, upon execution by a computing device or hardware, cause the device to perform those intended operations or actions.

[0194] The embodiments and functional operations described in this specification can be implemented in various forms, including digital electronic circuitry, software,Attorney Docket No. 56113-0823WO1 firmware, computer hardware (encompassing the disclosed structures and their structural equivalents), or any combination thereof. The subject matter can be realized as one or more computer programs, essentially modules of computer program instructions encoded on a tangible non-transitory storage medium for execution by or to control the operation of a computing device or hardware. The storage medium can be a storage device such as a hard drive or solid-state drive (SSD), a storage medium, a random or serial access memory device, or a combination of these. Additionally or alternatively, the program instructions can be encoded on a transmitted signal, such as a machine-generated electrical, optical, or electromagnetic signal, designed to carry' information for transmission to a receiving device or system for execution by a computing device or hardware. Furthermore, implementations may leverage emerging technologies like quantum computing or neuromorphic computing for specific applications, and may be deployed in distributed or cloud-based environments where components reside on different machines or within a cloud infrastructure.

[0195] The term “computing device or hardware’" refers to the physical components involved in data processing and encompasses all types of devices and machines used for this purpose. Examples include processors or processing units, computers, multiple processors or computers working together, graphics processing units (GPUs), tensor processing units (TPUs), and specialized processing hardware such as field- programmable gate arrays (FPGAs) or application-specific integrated circuits (ASICs). In addition to hardware, a computing device or hardware may also include code that creates an execution environment for computer programs. This code can take the form of processor firmware, a protocol stack, a database management system, an operating system, or a combination of these elements. Embodiments may particularly benefit from utilizing the parallel processing capabilities of GPUs, in a General-Purpose computing on Graphics Processing Units (GPGPU) context, where code specifically designed for GPU execution, often called kernels or shaders, is employed. Similarly, TPUs excel at running optimized tensor operations crucial for many machine learning algorithms. By leveraging these accelerators and their specialized programming models, the system can achieve significant speedups and efficiency gains for tasks involving artificial intelligence and machine learning, particularly in areas such as computer vision, natural language processing, and robotics.

[0196] A computer program, also referred to as software, an application, a module, a script, code, or simply a program, can be written in any programming language, includingAtorney Docket No. 56113-0823WO1 compiled or interpreted languages, and declarative or procedural languages. It can be deployed in various forms, such as a standalone program, a module, a component, a subroutine, or any other unit suitable for use within a computing environment. A program may or may not correspond to a single file in a file system and can be stored in various ways. This includes being embedded within a file containing other programs or data (e.g., scripts within a markup language document), residing in a dedicated file, or distributed across multiple coordinated files (e.g.. files storing modules, subprograms, or code segments). A computer program can be executed on a single computer or across multiple computers, whether located at a single site or distributed across multiple sites and interconnected through a data communication network. The specific implementation of the computer programs may involve a combination of traditional programming languages and specialized languages or libraries designed for GPGPU programming or TPU utilization, depending on the chosen hardware platform and desired performance characteristics.

[0197] In this specification, the term "engine" broadly refers to a software-based system, subsystem, or process designed to perform one or more specific functions. An engine is typically implemented as one or more software modules or components installed on one or more computers, which can be located at a single site or distributed across multiple locations. In some instances, one or more dedicated computers may be used for a particular engine, while in other cases, multiple engines may operate concurrently on the same one or more computers. Examples of engine functions within the context of Al and machine learning could include data pre-processing and cleaning, feature engineering and extraction, model training and optimization, inference and prediction generation, and post-processing of results. The specific design and implementation of engines will depend on the overall architecture and the distribution of computational tasks across various hardware components, including CPUs, GPUs, TPUs, and other specialized processors.

[0198] The processes and logic flows described in this specification can be executed by one or more programmable computers running one or more computer programs to perform functions by operating on input data and generating output. Additionally, graphics processing units (GPUs) and tensor processing units (TPUs) can be utilized to enable concurrent execution of aspects of these processes and logic flows, significantly accelerating performance. This approach offers significant advantages for computationally intensive tasks often found in Al and machine learning applications, suchAttorney Docket No. 56113-0823WO1 as matrix multiplications, convolutions, and other operations that exhibit a high degree of parallelism. By leveraging the parallel processing capabilities of GPUs and TPUs, significant speedups and efficiency gains compared to relying solely on CPUs can be achieved. Alternatively or in combination with programmable computers and specialized processors, these processes and logic flows can also be implemented using specialized processing hardware, such as field-programmable gate arrays (FPGAs) or applicationspecific integrated circuits (ASICs), for even greater performance or energy efficiency in specific use cases.

[0199] Computers capable of executing a computer program can be based on general- purpose microprocessors, special-purpose microprocessors, or a combination of both. They can also utilize any other type of central processing unit (CPU). Additionally, graphics processing units (GPUs), tensor processing units (TPUs), and other machine learning accelerators can be employed to enhance performance, particularly for tasks involving artificial intelligence and machine learning. These accelerators often work in conjunction with CPUs, handling specialized computations while the CPU manages overall system operations and other tasks. Typically, a CPU receives instructions and data from read-only memory (ROM), random access memory (RAM), or both. The elements of a computer include a CPU for executing instructions and one or more memory devices for storing instructions and data. The specific configuration of processing units and memory will depend on factors like the complexity of the Al model, the volume of data being processed, and the desired performance and latency requirements. Embodiments can be implemented on a wide range of computing platforms, from small embedded devices with limited resources to large-scale data center systems with high-performance computing capabilities. The system may include storage devices like hard drives, SSDs, or flash memory for persistent data storage.

[0200] Computer-readable media suitable for storing computer program instructions and data encompass all forms of non-volatile memory, media, and memory devices. Examples include semiconductor memory devices such as read-only memory (ROM), solid-state drives (SSDs), and flash memory devices; hard disk drives (HDDs); optical media; and optical discs such as CDs, DVDs, and Blu-ray discs. The specific type of computer-readable media used will depend on factors such as the size of the data, access speed requirements, cost considerations, and the desired level of portability or permanence.Atorney Docket No. 56113-0823WO1

[0201] To facilitate user interaction, embodiments of the subject matter described in this specification can be implemented on a computing device equipped with a display device, such as a liquid crystal display (LCD) or an organic light-emitting diode (OLED) display, for presenting information to the user. Input can be provided by the user through various means, including a keyboard), touchscreens, voice commands, gesture recognition, or other input modalities depending on the specific device and application. Additional input methods can include acoustic, speech, or tactile input, while feedback to the user can take the form of visual, auditor , or tactile feedback. Furthermore, computers can interact with users by exchanging documents with a user's device or application. This can involve sending web content or data in response to requests or sending and receiving text messages or other forms of messages through mobile devices or messaging platforms. The selection of input and output modalities will depend on the specific application and the desired form of user interaction.

[0202] Machine learning models can be implemented and deployed using machine learning frameworks, such as TensorFlow or JAX. These frameworks offer comprehensive tools and libraries that facilitate the development, training, and deployment of machine learning models.

[0203] Embodiments of the subject matter described in this specification can be implemented within a computing system comprising one or more components, depending on the specific application and requirements. These may include a back-end component, such as a back-end server or cloud-based infrastructure; an optional middleware component, such as a middleware server or application programming interface (API), to facilitate communication and data exchange; and a front-end component, such as a client device with a user interface, a web browser, or an app, through which a user can interact with the implemented subject matter. For instance, the described functionality could be implemented solely on a client device (e.g., for on-device machine learning) or deployed as a combination of front-end and back-end components for more complex applications. These components, when present, can be interconnected using any form or medium of digital data communication, such as a communication network like a local area network (LAN) or a wide area network (WAN) including the Internet. The specific system architecture and choice of components will depend on factors such as the scale of the application, the need for real-time processing, data security requirements, and the desired user experience.Attorney Docket No. 56113-0823WO1

[0204] The computing system can include clients and servers that may be geographically separated and interact through a communication network. The specific type of network, such as a local area network (LAN), a wide area network (WAN), or the Internet, will depend on the reach and scale of the application. The client-server relationship is established through computer programs running on the respective computers and designed to communicate with each other using appropriate protocols. These protocols may include HTTP. TCP / IP, or other specialized protocols depending on the nature of the data being exchanged and the security requirements of the system. In certain embodiments, a server transmits data or instructions to a user's device, such as a computer, smartphone, or tablet, acting as a client. The client device can then process the received information, display results to the user, and potentially send data or feedback back to the server for further processing or storage. This allows for dynamic interactions between the user and the system, enabling a wide range of applications and functionalities.

[0205] While this specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or on the scope of what may be claimed, but rather as descriptions of features that may be specific to particular embodiments of particular inventions. Certain features that are described in this specification in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Moreover, although features may be described above as acting in certain combinations and even initially be claimed as such, one or more features from a claimed combination can in some cases be excised from the combination, and the claimed combination may be directed to a subcombination or variation of a subcombination.

[0206] Similarly, while operations are depicted in the drawings and recited in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order shown or in sequential order, or that all illustrated operations be performed, to achieve desirable results. In certain circumstances, multitasking and parallel processing may be advantageous. Moreover, the separation of various system modules and components in the embodiments described above should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally beAttorney Docket No. 56113-0823WO1 integrated together in a single software product or packaged into multiple software products.

[0207] Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in a different order and still achieve desirable results. As one example, the processes depicted in the accompanying figures do not necessarily require the particular order shown, or sequential order, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous.

[0208] What is claimed is:

Claims

Attorney Docket No. 56113-0823WO1CLAIMS1. A method performed by one or more computers for training a neural network that includes a plurality of layers that include a particular layer, wherein the particular layer is associated with a set of base parameters and a set of adapter parameters, wherein the set of adapter parameters is represented by a first parameter matrix and a second parameter matrix, and wherein the method comprises: maintaining, for the particular layer of the neural network, a first unmagnified moment of gradients of a loss function with respect to the first parameter matrix and a second unmagnified moment of the gradients of the loss function with respect to the first parameter matrix; repeatedly performing training operations comprising: performing, using a plurality of training examples, a training step to obtain initial gradients of the loss function with respect to the first parameter matrix; determining unmagnified gradients of the loss function with respect to the first parameter matrix based on the initial gradients and the second parameter matrix; determining a first preconditioning matrix based on the second parameter matrix; updating the second unmagnified moment based on the first preconditioning matrix and the unmagnified gradients of the loss function with respect to the first parameter matrix; determining an unmagnified precondition step based on the unmagnified gradients of the loss function with respect to the first parameter matrix and the updated second unmagnified moment; updating the first unmagnified moment based on the first preconditioning matrix and the unmagnified precondition step: and updating the first parameter matrix based on the updated first unmagnified moment; and providing data specifying the adapted neural network for performing inference.

2. The method of claim 1, wherein the training operations comprise performing a polar decomposition of the second parameter matrix to generate an orthogonal matrix of the second parameter matrix and a symmetric positive semi-definite matrix of the second parameter matrix.Attorney Docket No. 56113-0823WO13. The method of claim 2. wherein determining the unmagnified gradients of the loss function comprises: determining the unmagnified gradients of the loss function based on the initial gradients and an inverse of the upper triangular matrix of the second parameter matrix.

4. The method of any one of claims 2-3, wherein determining the first preconditioning matrix comprises: determining the first preconditioning matrix based on a transpose of the orthogonal matrix of the second parameter matrix.

5. The method of claim 1, wherein the method comprises maintaining, for the particular layer of the neural network, an escaped mass, wherein the training operations comprise updating the escaped mass based on the first preconditioning matrix.

6. The method of claim 5. wherein determining the unmagnified precondition step comprise: determining the unmagnified precondition step based on the unmagnified gradients of the loss function, the updated second unmagnified moment, and the updated escaped mass.

7. The method of any one of claims 2-6, wherein updating the first parameter matrix comprises: updating the first parameter matrix based on the updated first unmagnified moment and the symmetric positive semi-definite matrix.Atorney Docket No. 56113-0823WO18. The method of any one of claims 1-7, wherein the method further comprises: maintaining, for the particular layer of the neural network, a first unmagnified moment of gradients of a loss function with respect to the second parameter matrix and a second unmagnified moment of the gradients of the loss function with respect to the second parameter matrix; and wherein performing the training step further comprises performing, using the plurality of training examples, the training step to obtain initial gradients of the loss function with respect to the second parameter matrix, and wherein the training operations further comprise: determining unmagnified gradients of the loss function with respect to the second parameter matrix based on the initial gradients and the second parameter matrix; determining a second preconditioning matrix based on the first parameter matrix; updating the second unmagnified moment based on the second preconditioning matrix and the unmagnified gradients of the loss function with respect to the second parameter matrix; determining an unmagnified precondition step based on the unmagnified gradients of the loss function with respect to the second parameter matrix and the updated second unmagnified moment; updating the first unmagnified moment based on the second preconditioning matrix and the unmagnified precondition step; and updating the second parameter matrix based on the updated first unmagnified moment.

9. The method of claim 8, wherein the training operations comprise performing a polar decomposition of the first parameter matrix to generate an orthogonal matrix of the first parameter matrix and a symmetric positive semi-defmite matrix of the first parameter matrix.

10. The method of claim 9, wherein determining the unmagnified gradients of the loss function comprises: determining the unmagnified gradients of the loss function based on the initial gradients and an inverse of the upper tnangular matrix of the first parameter matrix.Attorney Docket No. 56113-0823WO111. The method of any one of claims 8-9, wherein determining the second preconditioning matrix comprises: determining the second preconditioning matrix based on a transpose of the orthogonal matrix of the first parameter matrix.

12. The method of claim 8, wherein the method comprises maintaining, for the particular layer of the neural network, an escaped mass, wherein the training operations comprise updating the escaped mass based on the second preconditioning matrix.

13. The method of claim 12, wherein determining the unmagnified precondition step comprise: determining the unmagnified precondition step based on the unmagnified gradients of the loss function, the updated second unmagnified moment, and the updated escaped mass.

14. The method of any one of claims 9-13, wherein updating the second parameter matrix comprises: updating the second parameter matrix based on the updated first unmagnified moment and the symmetric positive semi-definite of the first parameter matrix.

15. The method of any one of claims 1-14, wherein the neural network is a generative neural network and the particular layer is an attention layer.

16. A system comprising one or more computers and one or more storage devices storing instructions that when executed by the one or more computers cause the one more computers to perform the operations of the respective method of any one of claims 1-15.

17. One or more computer storage media storing instructions that when executed by one or more computers cause the one more computers to perform the operations of the respective method of any one of claims 1-15.

Images