Auxiliary neural network for fast adaptation of pre-trained neural networks

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for training a customized neural network to perform a customized generative task. One of the methods include: obtaining data specifying a neural network that comprises an original set of parameters having pre-trained values; receiving adaptation data; generating a customized neural network that comprises (i) the original set of parameters having the pre-trained values and (ii) one or more additional sets of parameters; and adapting the customized neural network to the customized generative task.

Description

AUXILIARY NEURAL NETWORK FOR FAST ADAPTATION OFPRE-TRAINED NEURAL NETWORKSCROSS-REFERENCE TO RELATED APPLICATION

[0001] This application claims priority to U.S. Provisional Application No. 63 / 524,542, filed on June 30, 2023. 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 machine learning models.

[0003] Machine learning models receive an input and generate an output, e.g., a predicted output, based on the received input. Some machine learning models are parametric models and generate the output based on the received input and on values of the parameters of the model.

[0004] Some machine learning models are deep models that employ multiple layers of models to generate an output for a received input. For example, a deep neural network is a deep machine learning model that includes an output layer and one or more hidden layers that each apply a non-linear transformation to a received input to generate an output.SUMMARY

[0005] This specification describes a training system implemented as computer programs on one or more computers in one or more locations that can effectively and efficiently customize a pre-trained generative neural network, i.e., adapt a neural network that has been pre-trained on one or more generative tasks to generate a customized neural network that can generate customized output data items when performing the one or more generative tasks.

[0006] Particular embodiments of the subject matter described in this specification can be implemented so as to realize one or more of the following advantages. An auxiliary neural network is used to rapidly adapt a pre-trained neural network to generate customized output data items having one or more attributes when performing a generative task by generating an initial prediction of the values of one or more sets of additional parameters of the neural network that are specific to the one or more attributes, which will then be further fine-tuned using a relatively small amount of adaptation data.

[0007] Since the amount of additional parameters is orders of magnitude smaller than that of the original parameters and, in particular, that the fine-tuning begins from the predicted (e.g.,rather than randomly initialized) values of the additional parameters that have been generated by an trained auxiliary neural network, the adaptation process is significantly faster, e.g.. in terms of wall clock time, as well as consumes reduced computing and memory resources, than existing adaptation techniques including model fine-tuning, adapters, and prompt tuning. This enables the adaptation process to be performed by consumer hardware of end users, e g., mobile system-on-chip (SOC) devices or other embedded systems, rather than being performed in a datacenter.

[0008] The adaptation process is flexible and can be applied to adapt any of a variety of pretrained neural netw ork to any of a variety of different customized generative tasks that require generation of data items having specific attributes. Once adapted, the adapted neural network will be able to achieve performance competitive with or even exceeding the state-of- the-art on a wide range of customized generative tasks, even including subject-driven image synthesis tasks, despite having consumed significantly fewer computing and memory resources during the adaptation.

[0009] The details of one or more embodiments of the subject 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

[0010] FIG. 1 is a diagram of an example training system.

[0011] FIG. 2 is an example illustration of an auxiliary neural netw ork.

[0012] FIG. 3 is an example illustration of training a customized neural netw ork.

[0013] FIG. 4 is an example illustration of operations performed to train and use an auxiliary neural netw ork.

[0014] FIG. 5 is a flow diagram of an example process for adapting a customized neural network to a customized generative task.

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

[0016] FIG. 1 is a diagram of 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 ormore locations, in which the systems, components, and techniques described below can be implemented.

[0017] The training system 100 obtains data specifying a neural network 110 that has been pre-trained on one or more generative tasks and adaptation data 120, generates a customized neural network 140 based on the neural network 110, and use the adaptation data 120 and an auxiliary neural network 130 to train the customized neural network 140 to perform a customized generative task to generate customized output data items.

[0018] The neural network 110 can be any appropriate generative neural network that has been pre-trained, e.g., by the training system 100 or another training system, on the one or more generative tasks. The neural network 110 has an original set of parameters 114 that have pre-trained values determined as a result of the pre-training. Thus the data specifying the neural network 110 can include parameter data that defines pre-trained values of the original set of parameters 114 of the neural network 110 and, optionally, architecture data that defines the architecture of the neural network 110.

[0019] In some implementations, the neural network 110 can be configured as. or include, a diffusion neural network, or another image generation neural network, that has been pretrained to perform an image or video generation task, either unconditioned or conditioned on text prompts, e.g., text descriptions that describe the content of the image or video, where the output data items are images or videos that each have multiple frames (where each frame is an image).

[0020] In these implementations, the neural network 110 can be implemented with any appropriate diffusion neural network architecture that can be configured through the pretraining to, at any given updating iteration in an iterative update process, process a diffusion input for the updating iteration that includes a current data item or a current latent representation of an output data item (as of the updating iteration) to generate a denoising output for the updating iteration. From the denoising output, an updated data item or an updated latent representation of the output data item can be inferred.

[0021] For example, the denoising output can be an estimate of the noise that needs to be, e.g., added to a final data item, i.e., to the output data item being generated by the diffusion neural network, or to a final latent representation of the output data item, to arrive at the current data item or the current latent representation of the output data item. For example, the diffusion neural network can have been pre-trained on a set of training data items using a denoising score-matching objective to determine the pre-trained values of the original set of parameters 1 14 of the neural network 110.

[0022] For example, the neural network 110 can have a convolutional neural network architecture, e.g., a U-Net architecture, that has multiple convolutional blocks (see, for example, Ho, Jonathan, et al. "Denoising diffusion probabilistic models." Advances in neural information processing systems 33 (2020): 6840-6851).

[0023] As another example, the neural network 110 can have a Transformer neural network architecture that has a set of self-attention blocks and / or a set of cross-attention blocks and that processes the denoising input through the set of self-attention blocks and / or the set of cross-attention blocks to generate the denoising output (see, for example, Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE / CVF conference on computer vision and pattern recognition. 2022).

[0024] As another example, the neural network 110 can have a diffusion Transformer (DiT) architecture (see, for example, Peebles, William, et al. "Scalable diffusion models with transformers." Proceedings of the IEEE / CVF International Conference on Computer Vision. 2023).

[0025] As another example, the neural network 110 can have a universal vision Transformer (UViT) architecture (see, for example, Chen, Wuyang, et al. "A simple single-scale vision transformer for object localization and instance segmentation." arXiv preprint arXiv:2112.09747 (2021)).

[0026] In some implementations, the neural network 110 can be configured as, or include, a generative (large) language model neural network or a multi-modal neural network, e.g.. a visual and language model that that has been pre-trained to perform a text, image, or video generation task by generating tokens that make up the output data items in an auto-regressive manner, either unconditioned or conditioned on text prompts, where the output data items are textual data items, image data items, or audio data items, respectively.

[0027] In these implementations, the customized neural network 140 can have any of a variety of Transformer-based neural network architectures, e.g., encoder-only Transformer architectures, encoder-decoder Transformer architectures, decoder-only Transformer architectures, other attention-based architectures, and so on.

[0028] Examples of such architectures include those described in Colin Raffel. Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning w ith a unified text-to-text transformer. arXiv preprint arXiv: 1910. 10683, 2019; Daniel Adiwardana, Minh-Thang Luong. David R. So. Jamie Hall. Noah Fiedel. Romal Thoppilan, Zi Yang, Apoorv Kulshreshtha, Gaurav Nemade, Yifeng Lu, and Quoc V. Le. Towards a human-like open-domain chatbot. CoRR, abs / 2001.09977, 2020; Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020; Aakanksha Chowdhery, et al. PaLM: Scaling Language Modeling with Pathways, arXiv preprint arXiv:2204.02311; and Rohan Anil, et al. Palm 2 technical report. arXiv preprint arXiv:2305. 10403, 2023.

[0029] In these implementations, the neural network 110 can have been pre-trained on a large training dataset based on optimizing a maximum-likelihood objective to determine the pretrained values of the original set of parameters 114 of the neural netw ork 110. In various cases, such a training dataset can be 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., computer code that is publicly available from the Internet or another code repository, a large dataset of audio samples, e.g., audio recordings or w aveforms 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.

[0030] The adaptation data 120 includes one or more reference input data items. Each reference input data item has one or more attributes. An output data item to be generated by the customized neural network 140 is referred to as “customized” when the output data item includes attributes that are similar to those included in the one or more reference input data items.

[0031] Because of the way the training system 100 trains the customized neural network 140 as discussed throughout this specification, the adaptation data 120 need only include a small number of, e.g., fewer than ten, fewer than five, or even fewer than two (i.e., as few as one), reference input data items that have the one or more attributes.

[0032] The reference input data item can be any kind of input data item, e.g., textual data item, image data item, audio data item, and so on, and correspondingly, the training system 100 can train the customized neural network 140 to generate any kind of customized output data items, e.g., textual data items, image data items, audio data items.

[0033] In some implementations, the customized neural network 140 can be an image or video generation neural netw ork that generates images or videos that each have multiple frames (where each frame is an image) as sequences of pixels, i.e., each personalized outputdata item generated by the customized neural network 140 includes a sequence of color values for pixels in an output image arranged according to a specified order.

[0034] In these implementations, the images generated by the customized neural network 140 have one or more specific visual attributes that are similar to the specific visual attributes included in the one or more reference images included in the adaptation data 120. For example, a reference image may include a specific color, a specific texture, a specific brightness, and so on, and the customized images generated by the customized neural network 140 can also have the specific color, the specific texture, the specific brightness, and so on. As another example, a reference image may include a depiction of a specific object or a specific subject instance, a depiction of a specific watermark, a depiction of a specific logo, a depiction of a specific symbol, and so on, and the customized images generated by the customized neural network 140 can also have a depiction of the specific object or the specific subject instance, a depiction of the specific watermark, a depiction of the specific logo, a depiction of the specific symbol, and so on.

[0035] As a particular example, a reference image may include a depiction of a particular subject instance of an object class, and the customized images generated by the customized neural network 140 can also have a depiction of the particular subject instance (instead of variable subject instances that belong to the object class). For example, the subject instances can include landmarks, landscape or location features, vehicles, tools, food, clothing, devices, animals, and human faces, to name just a few examples.

[0036] In some implementations, the customized neural network 140 can be a text generation neural network that generates text sequences, i.e., each customized output data item generated by the customized neural network 140 is an output sequence of text that includes a sequence of text 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.

[0037] In these implementations, the text sequences generated by the customized neural network 140 have one or more specific textual attributes that are similar to the specific textual attributes included in the one or more reference text sequences included in the adaptation data 120. For example, a reference text sequence may include specific slangs, specific idioms, specific abbreviations, specific punctuations, and so on, and the customized text sequences generated by the customized neural network 140 can also have the specific slangs, the specific idioms, the specific abbreviations, the specific punctuations, and so on. As another example, a reference text sequence may represent a piece of code or a snippet ofcomputer code in a programming language that includes specific indentation structure, specific variable names, and so on, and the customized text sequences generated by the customized neural network 140 can represent pieces of code or snippets of computer code that also have the specific indentation structure, specific variable names, and so on.

[0038] In some implementations, the customized neural network 140 can be an audio generation neural network, e.g., each customized output data item generated by the customized neural network 140 is an audio signal that includes a sample of an audio wave at each of a sequence of time steps that span a specified time window. For example, the time steps can be arranged at regular intervals within the specified time window. The audio sample at a given time step can be an amplitude value of the audio wave or an amplitude value that has been compressed, companded, or both. For example, the audio sample can be a raw amplitude value or a mu-law companded representation of the amplitude value.

[0039] In these implementations, the audio signals generated by the customized neural network 140 have one or more specific audio attributes that are similar to the specific audio attributes included in the one or more reference audio signals included in the adaptation data 120. For example, a reference audio signal may include a specific a specific acoustic property, e.g., a specific speaker identity, a specific recording condition (such as a specific level of reverberation, distortion, or background noise), and so on, and the customized audio signals generated by the customized neural network 140 can also have the specific acoustic property, the specific recording condition, and so on. As another example, a reference audio signal may include a specific semantic property, e.g., a specific linguistic content (when the reference audio signal represents speech), or a specific melody or a specific rhythm (when the reference audio signal represents music), and so on, and the customized audio signals generated by the customized neural network 140 can also have the specific semantic property, and so on.

[0040] In each of the above, the specific attributes / properties / subject-instances can be unseen. That is, the attributes / properties / subject-instances do not need to be part of the training dataset used in pre-training the neural network 110. In some implementations, a text prompt can be used to specify the specific attributes / properties / subject-instances to be used from the adaptation data.

[0041] Having received the data specifying the pre-trained neural network 110, the training system 100 generates the customized neural network 140 based on the neural network 110 byexpanding the architecture of the neural network 110 to include more parameters.

[0042] FIG. 1 thus illustrates that the customized neural network 140 includes one or more additional sets of parameters 118 in addition to the original set of parameters 114 of the neural network 110. That is, the customized neural network 140 has some parameters (that is, the original set of parameters 114) that are the same as those included in the neural network 110, and some other parameters (that is, the one or more additional sets of parameters 118) that were not included in the neural network 110.

[0043] In practice, the training system 100 can expand any one or more of the layers, e.g., linear layers, convolutional layers, self-attention layers, cross-attention layers, and so forth, of the neural network 110 to add additional parameters to the one or more of the layers of the neural network 110.

[0044] After expansion, the customized neural network 140 is configured to generate the output data items that have attributes similar to those included in the one or more reference input data items in accordance with both (i) the original set of parameters 114 and (ii) the one or more additional sets of parameters 118. The original set of parameters 114 are shared across different attributes. On the other hand, the one or more additional sets of parameters 118 are specific to the attributes similar to those included in the one or more reference input data items.

[0045] To generate output data items that have different attributes, the customized neural network 140 uses the same original set of parameters 114 but different ones of the one or more additional sets of parameters. Thus, output data items having one or more first attributes will be generated by a different combination of parameters of the customized neural network 140 than output data items having one or more second attributes.

[0046] Put another w ay, some of the parameters of the customized neural network 140 (that is, the original set of parameters 114) will be used to generate output data items that have multiple different attributes, while some others of the parameters of the customized neural network 140 (for example, the one or more additional sets of parameters 118) will only be used to generate output data items that have one or more particular attributes that correspond to those parameters (for example, output data items having attributes that correspond to the one or more additional sets of parameters 118).

[0047] After having generated the customized neural network 140, the training system 100 trains the customized neural netw ork 140 to determine updated values of the one or more additional sets of parameters 118.

[0048] The one or more additional sets of parameters 118 are specific to the attributes that are similar to those included in the one or more reference input data items included in theadaptation data 120 received by the training system 100, thus the training system 100 can generally leam different sets of parameters that correspond respectively to different attributes included in different adaptation data. The amount of the one or more sets of additional parameters is smaller, usually several orders of magnitude smaller, than the amount of the original set of parameters 114 of the customized neural network 140.

[0049] The training system 100 determines the updated values of one or more additional sets of parameters 118 by using a two-step process which includes an initial prediction step followed by a fine-tuning step.

[0050] During the initial prediction step, the training system 100 sets the new values of the one or more additional sets of parameters 118. Unlike some existing training systems that set the new values in accordance with a parameter value initialization method that does not depend on the adaptation data 120, e.g., by initializing the one or more additional sets of parameters 118 to zeros, ones, or some other predetermined, e.g., randomly initialized, values, the training system 100 sets the new values of the one or more additional sets of parameters 118 in accordance with a parameter value initialization method that is dependent on the adaptation data 120.

[0051] In particular, the training system 100 makes use of an auxiliary neural network 130 to predict, based on processing the one or more reference input data items included in the adaptation data 120, the new values of the one or more additional sets of parameters 118 of the customized neural network 140. In this way. different reference input data items that have different attributes will generally result in different new values of the one or more additional sets of parameters 118 (although they will all be predicted by using the same auxiliary neural network 130).

[0052] The auxiliary neural network 130 can have any appropriate neural network architecture that allows the auxiliary neural network 130 to map the one or more reference input data items included in the adaptation data 120 to a prediction of the new values of the one or more additional sets of parameters 118 of the customized neural network 140.

[0053] A particular example architecture will be described in more detail below with reference to FIG. 2, but more generally, the auxiliary neural network 130 can include any appropriate types of neural network layers (e.g., linear layers, convolutional layers, attention layers, and so forth) in any appropriate number (e.g., 1 layer, or 5 layers, or 10 layers) and connected in any appropriate configuration (e.g., as a directed graph of layers).

[0054] During the fine-tuning step, the training system 100 then determines the updated values of the one or more additional sets of parameters 118 starting from their new values,e.g., based on optimizing a fine-tuning loss function that is suitable for the generative task, the attributes of the output data items, or both.

[0055] After the training system 100 has determined the updated values of one or more additional sets of parameters 118, i.e., once the two-step process is complete, the training system 100 can provide data specifying the customized neural network 140, e.g., data specifying the updated values of the one or more additional sets of parameters 118 and, in some cases, the pre-trained values of the original set of parameters 114 and / or the architecture of the customized neural network 140, for deployment for performing the one or more generative tasks on another system.

[0056] For example, a data item generation system that implements the customized neural network 140 can generate an arbitrary number of new customized output data items, e.g., textual data items, image data items, audio data items, and so on, in response to requests submitted by the user, where each customized output data item has one or more attributes that are similar to those included in the one or more reference input data items included in the adaptation data 120.

[0057] Alternatively or in addition, the training system 100 can deploy the customized neural network 140 and use the customized neural network 140 to generate new customized output data items, e.g., textual data items, image data items, audio data items, and so on, in response to requests submitted by the user. Furthermore, the training system 100 can output the data specifying the updated values of the one or more additional sets of parameters 118 for storage in a storage device for some future purpose.

[0058] As a particular example, the training system 100 can receive, e.g., from a user or another system, a reference image that includes a depiction of a particular subject instance of interest (e.g.. a particular person, a particular animal, a particular car. a particular boat, etc.) and generate and train a customized neural network 140 based on the reference image. Then, the training system 100 can provide the customized neural network 140 to an image synthesis system — and the image synthesis system can use the customized neural network 140 to generate an arbitrary number of high-fidelity synthetic images (e.g., photorealistic images) that each show the particular subject instance of interest with a diverse range of various contexts, poses, views, and / or lighting conditions.

[0059] For example, the customized neural network 140 in this example can be configured as a diffusion neural network, e.g., a text-to-image diffusion neural network that is configured to generate output images conditioned on text prompts, and the text prompt can include a unique identifier defining the particular subject instance. Such a unique identifier can be generatedbased on the reference image, for example, by using the techniques described in Ruiz, Nataniel. et al. "Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation." Proceedings of the IEEE / CVF Conference on Computer Vision and Pattern Recognition. 2023.

[0060] Optionally, after the training system 100 has determined the updated values of one or more additional sets of parameters 118, some implementations of the training system 100 can perform an evaluation to determine performance of the customized neural network 140 on the customized generative task, e.g., how similar the new customized output data items generated by the customized neural network 140 will be to the reference input data items included in the adaptation data 120.

[0061] For example, if the performance fails to meet a performance threshold because the new customized output data items are not similar to the reference input data items, then the training system 100 can refrain from outputting data specifying the customized neural network 140. Instead, the training system 100 can generate and train another instance of the customized neural network 140, e.g., by adding further additional parameters, using additional reference input data items, or both.

[0062] FIG. 2 is an example illustration 200 of the auxiliary' neural network 130.

[0063] The auxiliary' neural network 130 includes an encoder neural network 210, followed by one or more initial linear layers 220, followed by a decoder neural network 230, and followed by one or more subsequent linear layers 240.

[0064] For a reference input data item 123 included in the adaptation data 120, the encoder neural network 210 processes the reference input data item 123 to generate an encoded representation of the reference input data item 123. For example, the encoded representation of the reference input data item 123 can be in the form of a feature vector, e.g., a high dimensional feature vector.

[0065] The encoder neural network 210 can have any' appropriate neural network architecture that allows it to map the reference input data item 123 to the feature vector. For example when the reference input data item 123 is an image, the encoder neural network 210 can use a stack of convolutional layers, a stack of attention layers, and so on. As a particular example, the encoder neural network 210 can have a Visual Transformer (ViT) architecture (see, for example, Dosovitskiy, Alexey', et al. "An image is 'orth 16x16 w'ords: Transformers for image recognition at scale." arXiv preprint arXiv:2010. 11929 (2020)) that includes a stack of attention layers that operate on images patches of the image. As another example when thereference input data item 123 is a text sequence, the encoder neural network 210 can use a stack of fully connected layers, a stack of attention layers, and so on.

[0066] The one or more initial linear layers 220 processes the encoded representation of the reference input data item 123 to generate a projected representation of the reference input data item 123. For example, the projected representation of the reference input data item 123 can be in the form of a reduced dimensionality feature vector that has few er numeric values than the high dimensional feature vector.

[0067] The inclusion of the one or more initial linear layers 220 in the auxiliary neural network 130 is optional, as the encoder neural network 210 may directly generate the reduced dimensionality feature vector as the encoded representation of the reference input data item 123.

[0068] The decoder neural network 230 processes a decoder input that includes (i) the encoded representation of the reference input data item 123 and, optionally, in some implementations, (ii) an initial latent representation of initial values of the one or more additional sets of parameters 118, to generate an updated latent representation of the initial values of the one or more additional sets of parameters 118.

[0069] In these implementations, the auxiliary neural network 130 can combine, e.g., sum or concatenate, the encoded representation of the reference input data item 123 to the initial latent representation of the initial values of the one or more additional sets of parameters 118 to generate a combination, and then provide the combination as the decoder input for processing by the decoder neural network 230.

[0070] For example, the initial values of the one or more additional sets of parameters 118 can be set to all zeros, all ones, or some other predetermined numeric values, and the initial latent representation can be in the form of one or more feature vectors, e.g., a feature vector for each additional set of parameters, that each include all zeros, all ones, or some other predetermined numeric values.

[0071] The decoder neural network 230 can have any appropriate neural netw ork architecture that allows it to update the decoder input to generate the updated latent representation of the initial values of the one or more additional sets of parameters 118. For example, the decoder neural network 230 can a stack of linear layers, a stack of convolutional layers, a stack of attention layers, a stack of recurrent layers, and so forth.

[0072] As a particular example, the decoder neural network 230 can have a Transformer architecture, e.g., an encoder-decoder Transformer architecture or a decoder-only Transformer architecture, that includes a stack of one or more attention layers that areconfigured to generate the updated latent representation by incrementally updating the decoder input by applying an attention mechanism at each attention layer.

[0073] The one or more subsequent linear layers 240 processes the updated latent representation of the values of the one or more additional sets of parameters 118 to convert the updated latent representation to the new values of the one or more additional sets of parameters 118.

[0074] In some implementations, the outputs of the one or more subsequent linear layers 240 directly define the new values of the one or more additional sets of parameters 118. In other implementations, the outputs of the one or more subsequent linear layers 240 indirectly define the new values of the one or more additional sets of parameters 118, For example, the outputs can define a change (‘"delta”) that should be applied to the initial values to obtain the new values. In this example, the new values can be generated by adding or subtracting the delta to the initial values. In other examples, the deltas can specify a change to some or all of the pre-trained weights.

[0075] When multiple reference input data items are included in the adaptation data, the training system 100 can, for example, either (i) select one of the multiple reference input data items for processing by the auxiliary neural network 130 or (ii) process different reference input data items by the auxiliary' neural network 130 to generate different new values of the one or more additional sets of parameters 118 and then combine, e.g., average or sum, the different new values of to generate the new values of the one or more additional sets of parameters 1 18.

[0076] The inclusion of the one or more subsequent linear layers 240 in the auxiliary' neural network 130 is optional, as the decoder neural network 230 may directly or indirectly generate the new values of the one or more additional sets of parameters 118 as output.

[0077] In some implementations, the auxiliary neural network 130 can be configured to generate the new values of the one or more additional sets of parameters 118 in an iterative manner by iteratively updating the decoder input using the decoder neural network 230 over multiple update steps.

[0078] For example. FIG. 2 illustrates that, for each of one or more layers 1-L. the auxiliary neural network 130 generates the values of an additional set of parameters for the layer over K iterations.

[0079] At the first iteration, the encoder neural network 210 processes the reference input data item 123 to generate an encoded representation of the reference input data item 123, which is combined with an initial latent representation of the initial values of the additionalset of parameters of the layer and processed by the decoder neural network 230 to generate an updated latent representation of the initial values of the additional set of parameters of the layer.

[0080] At each subsequent iteration, the updated latent representation that has been generated in the immediately preceding iteration is combined with the encoded representation of the reference input data item 123 and processed by the decoder neural network 230 to generate a further updated latent representation. Thus the updated latent representation that will be converted to generate the new values of the additional set of parameters of the layer is the updated latent representation generated by the encoder neural network 210 in the last iteration (the Kth iteration).

[0081] FIG. 3 is an example illustration of training the customized neural network 140. In particular, the training system 100 uses a lightweight low-rank approximation technique 350 to train the customized neural network 140.

[0082] FIG. 3 illustrates the lightweight low-rank approximation technique 350 in comparison with a low-rank approximation technique 300 that would traditionally be used to train a neural network similar to the customized neural network 140.

[0083] Suppose that a given layer, e.g., a linear layer, a self-attention layer, or a crossattention layer, of the customized neural network 140 has an original set of parameters VF0(that are shared across different attributes) and an additional set of parameters AW (that are specific to the attributes similar to those included in the one or more reference input data items), the customized neural network 140 is configured to generate the output data items that have attributes similar to those included in the one or more reference input data items by performing one or more forward passes through layers of the customized neural network 140, including the given layer, in accordance with the parameters of the given layer W' = Wo+ AW.

[0084] As illustrated on the top of FIG. 3, for the given layer which has an additional set of parameters that can be represented by a matrix A IV e Knxni, the low-rank approximation technique 300 learns the values of the additional set of parameters AW by decomposing the matrix into a low -rank down matrix A G Mnxrand a low-rank up matrix B G Mrxm. where AW = AB and r « n, and then learns the elements of the matrices A and B. The learnable elements in matrices A and B collectively represent (or approximate) the values of the additional set of parameters API / , but the combined number of learnable elements in A and Bis smaller than the number of parameters originally included in the additional set of parameters AW / .

[0085] As illustrated on the bottom of FIG. 3, for the given layer which has an additional set of parameters that can be represented by a matrix AW G Mnxm, the lightweight low-rank approximation technique 350 leams the values of the additional set of parameters AW / by further decomposing each of the down matrix A and the up matrix B used in the low-rank approximation technique 300 into two matrices.

[0086] The lightweight low-rank approximation technique 350 involves generating, from a matrix having values that represent parameters of the given layer and based on rank decomposition, a first pair of decomposed matrices and a second pair of decomposed matrices. Each pair of decomposed matrices includes two decomposed matrices. Each decomposed matrix has values that, when multiplied with values included in other decomposed matrices, collectively represent the additional set of parameters AW / of the given layer.

[0087] More specifically, the lightweight low-rank approximation technique 350 decomposes the down matrix A = AauxAtraininto WaiiXG Uknxaand AtrainG Ukaxrand decomposes the up matrix B ^train^aux into B]ra nG

[0088] By applying the lightweight low-rank approximation technique 350, the additional set of parameters AW / G Bnxmcan thus be represented as:where r << min(n, m), a < n and b < m. Aauxand Bauxare randomly initialized with row-wise orthogonal vectors with constant magnitude and the initialized values are frozen. B trainand Strain have the learnable elements that collectively represent (or approximate) the values of the additional set of parameters AW / . The combined number of learnable elements is smaller than the number of parameters originally included in the additional set of parameters AW / , and in particular is further smaller than the combined number of learnable elements in matrices A and B used in the low-rank approximation technique 300.

[0089] In the equation above, a and b are hyperparameters. The values of the hyperparameters a and b. which control the total number of the learnable elements to represent the additional set of parameters AW / of the given layer, can be determined based on a total number of the parameters of the given layer. For example, the values of the hyperparameters a and b can be dependent on, e.g., proportional to, the total number of the parameters of the given layer.

[0090] For example, suppose the given layer is one of the attention layers included in the architecture mentioned in Rombach. Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE / CVF conference on computer vision and pattern recognition. 2022, when applying the low-rank approximation technique 300 with r = 1, the total number of learnable elements is about 386K and the size of the learnable elements about 1.6MB. In contrast, when applying the lightweight low-rank approximation technique 350 with a = 100 and b = 50, the total number of learnable elements is about 30K and the size of the learnable elements is about 120 KB.

[0091] This is more than an order of magnitude smaller than what would have to be learned by using the low-rank approximation technique 300 to train the customized neural network 140. Applying the lightweight low-rank approximation technique 350 thus results in a reduced amount of computational and memory resources consumed by the training process by learning and storing a much smaller number of values representing the additional set of parameters for the given layer. The savings in computational and memory' resources can be particularly important for training processes where the lightweight low-rank approximation technique 350 is applied to many layers of the customized neural network 140. In addition, it has been found that applying this lightweight low-rank approximation technique 350 is not detrimental to task performance.

[0092] FIG. 4 is an example illustration of operations performed to train and use an auxiliary’ neural network 130.

[0093] The top half of FIG. 4 illustrates hoyv a training system implemented as computer programs on one or more computers in one or more locations can train the auxiliary' neural network 130 together yvith the customized neural network 140 on a training dataset to determine trained values of the parameters of the auxiliary neural network 130. For example, the training system can be the same as the training system 100 of FIG. 1.

[0094] The training dataset includes a plurality of training data items, e.g., textual data items, image data items, audio data items, and so on. In general, the number of training data items included in the training dataset can be greater than the number of reference input data items included in the adaptation data.

[0095] The training system obtains a batch of training data items 402 from the training dataset. For example the training system can do this by sampling a fixed number of training data items 402 from the training dataset.

[0096] For each training data item 402 in the batch, the auxiliary neural network 130 processes, in accordance with current values of the parameters of the auxiliary neural network 130, the training data item 402 to generate current values of the one or more additional sets of parameters 118 of the customized neural network 140.

[0097] Then, for each training data item 402 in the batch, the customized neural network 140 processes, in accordance with (i) the current values of the one or more additional sets of parameters 118 and (ii) the pre-trained values of the original set of parameters 114. a customized neural network input to generate a predicted training data item 404. The predicted training data item 404 is a predicted reconstruction of the training data item 402 in the batch.

[0098] In some implementations, the customized neural network input includes a text prompt 408. For example, the text prompt can be a text description that describes the content of the training data item 402. In some implementations, e.g., in implementations where the customized neural network 140 is configured as a diffusion neural network, the customized neural network input also includes noise 406. For example, the customized neural network input can include a noisy version of the training data item 402 which can be generated by adding noise 406, e.g., a noise that is determined according to a noise schedule, to the training data item 402.

[0099] The training system computes an auxiliary' loss function for the batch. The auxiliary' loss function includes a reconstruction loss term that measures, for each training data item 402 in the batch, a difference between (i) the predicted training data item 404 which is a predicted reconstruction of the training data item 402 and (ii) the training data item 402. [000100] The auxiliary' loss function also includes a parameter space loss term that measures, for each training data item 402 in the batch, a difference between (i) initial values of the one or more additional sets of parameters 118 and (ii) the current values of the one or more additional sets of parameters 1 18. For example, the initial values of the one or more additional sets of parameters 118 can be set to all zeros, all ones, or some other predetermined numeric values.[000101] For example, the reconstruction loss term and the parameter space loss term can each be an LI loss, an L2 loss, a root-mean-square error, a Huber loss, or another difference measure.[000102] For example, the auxiliary loss function can be in the form of:where 9 represents the current values of the one or more additional sets of parameters 118 (e.g., the values included in the Btrainand Atrainmatrices), 6 represents the initial values of the one or more additional sets of parameters 118, x is a training data item 402,is the predicted training data item 404 generated by the customized neural network 140 conditioned on the noisy version of the training data item (% + e) and the text prompt c, and a and / ? are hyperparameters that control the weights of the reconstruction loss term and the parameter space loss term.[000103] Then, the training system determines, based on computing gradients of the auxiliary loss function with respect to the parameters of the customized neural network 140 and then backpropagating the gradients through appropriate parameters of the customized neural network 140 into the auxiliary neural network 130, one or more updates to the current values of the parameters of the auxiliary' neural network 130. In implementations, the system can then determine the updates by applying an update rule, e.g., an Adam update rule, an Rmsprop update rule, or a stochastic gradient descent (SGD) update rule, to the gradients.[000104] In some implementations, when updating the parameters of the auxiliary neural network 130, the training system holds the parameters of the customized neural network 140 fixed, such that the parameters of the customized neural network 140 are not updated during the training of the auxiliary neural network 130.[000105] In general, the training system can continue performing iterations of the operations described above with reference to the top half of FIG. 4 on different batches of the training data items obtained from the training dataset until termination criteria for the training of the auxiliary neural network 130 have been satisfied, e g., until the parameters have converged, until a threshold amount of wall clock time has elapsed, or until a threshold number of iterations of the operations have been performed.[000106] The bottom half of FIG. 4 illustrates how the training system can use the trained auxiliary neural network 130 to train the customized neural network 140 on the one or more reference input data items 412 included in the adaptation data to determine the updated values of the one or more additional sets of parameters 118 of the customized neural network 140 starting from their new values.[000107] At the initial prediction step, for each reference input data item 412 included in the adaptation data, the auxiliary neural network 130 processes, in accordance with the trained values of the parameters of the auxiliary neural network 130, the reference input dataitem 412 to generate the new values of the one or more additional sets of parameters 118 of the customized neural network 140.[000108] The initial prediction step is visualized in the bottom right comer of FIG. 4 by the directed arrow labeled '‘(1)” which brings circle 430 (representing the customized neural network 140 that includes (i) the one or more additional sets of parameters 118 having the initial values and (ii) the original set of parameters 114 having the pre-trained values) to circle 440 (representing the customized neural network 140 that includes (i) the one or more additional sets of parameters 118 having the new values which are generated by the auxiliary' neural network 130) and (ii) the original set of parameters 114 having the pre-trained values). [000109] The customized neural network 140 processes, in accordance with (i) the new values of the one or more additional sets of parameters 118 and (ii) the pre-trained values of the original set of parameters 114, a customized neural network input to generate a predicted reference data item. The predicted reference data item 414 is a predicted reconstruction of the reference input data item 412 included in the adaptation data.[000110] In some implementations, the customized neural network input includes a text prompt 418. In some implementations, the customized neural network input also includes noise 416. For example, the customized neural network input can include a noisy version of the reference input data item 412 which can be generated by adding noise 416, e.g., a noise that is determined according to a noise schedule, to the reference input data item 412.(0001111 At the fine-tuning step, the training system determines one or more updates to the new values of the one or more additional sets of parameters 1 18 based on optimizing a fine-tuning loss function that includes a reconstruction loss term that measures a difference between (i) the predicted reference data item 414 which is a predicted reconstruction of the reference input data item 412 and (ii) the reference input data item 412. For example, the reconstruction loss term can be an LI loss, an L2 loss, a root-mean-square error, a Huber loss, or another difference measure.[000112] For example, the fine-tuning loss function can be in the form of:where d represents the new values of the one or more additional sets of parameters 118 (e.g., the values included in the Btrainand Atrainmatrices), x is a reference data item 412, £)§ is the predicted reference data item 414 generated by the customized neural network 140 conditioned on the noisy version of the reference data item (x + e) and the text prompt c.[000113] Then, the training system determines, based on computing gradients of the fine-tuning loss function with respect to the parameters of the customized neural network 140, one or more updates to the new values of the one or more additional sets of parameters 118 of the customized neural network 140. In implementations, the system can then determine the updates by applying an update rule, e.g., an Adam update rule, an Rmsprop update rule, or a stochastic gradient descent (SGD) update rule, to the gradients.[000114] In some implementations, when updating the one or more additional sets of parameters 1 18 of the customized neural network 140, the training system holds the original set of parameters 114 of the customized neural network 140 fixed to their pre-trained values, such that the original set of parameters 114 of the customized neural network 140 are not updated when learning the updated values of the one or more additional sets of parameters 118 of the customized neural network 140.[000115] In practice, the training system can continue performing iterations of the operations described above with reference to the bottom half of FIG. 4 on each reference input data item 412 included in the adaptation data.[000116] The fine-tuning step is visualized on the bottom right comer of FIG. 4 by the directed arrow labeled “(2)” which brings circle 440 (representing the customized neural network 140 that includes (i) the one or more additional sets of parameters 118 having the new values which are generated by the auxiliary neural network 130) and (ii) the original set of parameters 114 having the pre-trained values) to circle 450 (representing the customized neural network 140 that includes (i) the one or more additional sets of parameters 118 having the updated values which are determined as a result of the fine-tuning) and (ii) the original set of parameters 114 having the pre-trained values).[000117] FIG. 5 is a flow diagram of an example process 500 for adapting a customized neural network to a customized generative task. For convenience, the process 500 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 500.[000118] The system obtains data specifying a neural network (step 502). For example, the neural network can be a diffusion neural network. As another example, the neural network can be a language model neural network. As another example, the neural network can be a multi-modal neural network.[000119] The neural network has been pre-trained to perform one or more generative tasks in accordance with an original set of parameters having pre-trained values that havebeen determined as result of the pre-training. For example, the neural network can have been pre-trained to perform one or more of: a text generation task (where the output data items generated by the neural network are textual data items), an image generation task (where the output data items generated by the neural network are image data items), or a video generation task (where the output data items generated by the neural network are video data items).[000120] The system obtains adaptation data for a customized generative task (step 504). The adaptation data includes one or more reference input data items. Each reference input data item has one or more attributes, e.g., one or more visual attributes, one or more textual attributes, or one or more audio attributes. For example, the system can obtain the adaptation data as an upload from a user of the system or can obtain the adaptation data from another system that is accessible by the system.[000121] The customized generative task is a generative task (which can be any one of the generative tasks on which the neural network has been pre-trained to perform) that requires the neural network to generate output data items that have attributes similar to the attributes of the reference input data items included in the adaptation data.[000122] In practice the similarity can be computed by using any appropriate measure, e.g., a cosine similarity measure, a dot product, a Euclidian distance, a Frechet distance, a dynamic time warping distance, to name just a few examples. Being similar means having a similarity that meets or satisfies a predetermined similarity threshold.[000123] The system generates a customized neural network based on expanding the architecture of the neural network to include more parameters (step 506). Therefore, the customized neural network includes (i) the original set of parameters having the pre-trained values and (ii) one or more additional sets of parameters.[000124] Thus, as used in this specification, '‘generating” a customized neural network refers to generating the values of the one or more additional sets of parameters.[000125] In some implementations, the original set of parameters having the pre-trained values and the one or more additional sets of parameters can be stored at the same location. In other implementations, the original set of parameters having the pre-trained values and the one or more additional sets of parameters can be stored across different, e.g., geographically remote, locations and / or using different storage devices, e.g., storage devices managed by different entities.[000126] In some implementations, the system can use the lightweight low-rank approximation technique as explained above with reference to FIG. 3 to generate thecustomized neural network. In these implementations, the one or more additional sets of parameters include elements that are included in the Aaux, Atrain, Btrain. and Bauxmatrices for each of one or more layers of the customized neural network.[000127] The system then performs steps 508-510 to adapt the customized neural network to the customized generative task.[000128] The system processes, using an auxiliary neural network, each reference input data item included in the adaptation data to generate new values of the one or more additional sets of parameters (step 508). For example, the auxiliary neural network can generate new values for the elements in the Btrainand Atrainmatrices for each of one or more layers of the customized neural network. By doing so, the new values of the one or more additional sets of parameters are set dependent on the adaptation data.[000129] When multiple reference input data items are included in the adaptation data, the system can, for example, either (i) select one of the multiple reference input data items for processing by the auxiliary neural network or (ii) process different reference input data items by the auxiliary neural network to generate different new values of the one or more additional sets of parameters and then combine, e.g., average or sum, the different new values of to generate the new values of the one or more additional sets of parameters.[000130] The system trains, using the one or more reference input data items included in the adaptation data, the customized neural network to determine updated values of the one or more additional sets of parameters starting from the new values of the one or more additional sets of parameters, while holding the pre-trained values of the original set of parameters fixed (step 510).[000131] Specifically, the system can update the values of the one or more additional sets of parameters based on optimizing a fine-tuning loss function that includes a reconstruction loss term that measures, for each reference input data item, a difference between (i) a predicted reference data item which is a predicted reconstruction of the reference input data item that has been generated by the customized neural network in accordance with (a) the new values of the one or more additional sets of parameters and (b) the pre-trained values of the original set of parameters and (ii) the reference input data item. [000132] In implementations where the lightweight low -rank approximation technique is used to generate the customized neural netw ork, the one or more additional sets of parameters include elements that are included in the Aaux. Atram. Btrain. and Bauxmatrices for each of one or more layers of the customized neural network, and the elements in theAauxand Bauxmatrices can be held fixed to their randomly initialized values, while the elements in the Btrainand Atrainmatrices can be updated. In particular, for the elements in the Btrainand Atrainmatrices, the system can update their values starting from the new values of these elements that have been generated by the auxiliary neural network.[000133] This specification uses the term “configured'’ in connection with systems and computer program components. For a system of one or more computers to be configured to perform particular operations or actions means that the system has installed on it software, firmware, hardware, or a combination of them that in operation cause the system to perform the operations or actions. For one or more computer programs to be configured to perform particular operations or actions means that the one or more programs include instructions that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions.[000134] Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry, in tangibly-embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or in combinations of one or more of them. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non transitory storage medium for execution by, or to control the operation of, data processing apparatus. The computer storage medium can be a machine- readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively or in addition, the program instructions can be encoded on an artificially generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, that is generated to encode information for transmission to suitable receiver apparatus for execution by a data processing apparatus.[000135] The term “data processing apparatus” refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including by way of example a programmable processor, a computer, or multiple processors or computers. The apparatus can also be, or further include, special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). The apparatus can optionally include, in addition to hardware, code that creates an execution environment for computer programs, e.g.. code that constitutes processor firmware, aprotocol stack, a database management system, an operating system, or a combination of one or more of them.[000136] A computer program, which may also be referred to or described as a program, software, a software application, an app, a module, a software module, a script, or code, can be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages; and it can be deployed in any form, including as a stand alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program can be stored in a portion of a file that holds other programs or data, e.g., one or more scripts stored in a markup language document, in a single file dedicated to the program in question, or in multiple coordinated files, e.g., files that store one or more modules, sub programs, or portions of code. A computer program can be deployed to be executed on one computer or on multiple computers that are located at one site or distributed across multiple sites and interconnected by a data communication network.[000137] In this specification, the term "database" is used broadly to refer to any collection of data: the data does not need to be structured in any particular way, or structured at all, and it can be stored on storage devices in one or more locations. Thus, for example, the index database can include multiple collections of data, each of which may be organized and accessed differently.|000138| Similarly, in this specification the term “engine” is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine will be implemented as one or more software modules or components, installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in other cases, multiple engines can be installed and running on the same computer or computers.[000139] The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also be performed by special purpose logic circuitry, e.g., an FPGA or an ASIC, or by a combination of special purpose logic circuitry and one or more programmed computers. [000140] Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read only memory or a random access memory or both. The essential elements of a computer are acentral processing unit for performing or executing instructions and one or more memoty devices for storing instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special purpose logic circuitry. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto optical disks, or optical disks. However, a computer need not have such devices. Moreover, a computer can be embedded in another device, e.g., a mobile telephone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few. [000141] Computer readable media suitable for storing computer program instructions and data include all forms of non volatile memory, media and memory devices, including by way of example semiconductor memory devices, e.g., EPROM, EEPROM, and flash memory devices; magnetic disks, e.g., internal hard disks or removable disks; magneto optical disks; and CD ROM and DVD-ROM disks.[000142] To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user and a keyboard and a pointing device, e.g., a mouse or a trackball, by which the user can provide input to the computer. Other kinds of devices can be used to provide for interaction with a user as well; for example, feedback provided to the user can be any form of sensory feedback, e.g., visual feedback, auditory feedback, or tactile feedback; and input from the user can be received in any form, including acoustic, speech, or tactile input. In addition, a computer can interact with a user by sending documents to and receiving documents from a device that is used by the user; for example, by sending web pages to a web browser on a user’s device in response to requests received from the web browser. Also, a computer can interact with a user by sending text messages or other forms of message to a personal device, e.g., a smartphone that is running a messaging application, and receiving responsive messages from the user in return.[000143] Data processing apparatus for implementing machine learning models can also include, for example, special-purpose hardware accelerator units for processing common and compute-intensive parts of machine learning training or production, i.e., inference, workloads.[000144] Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework or a J AX framework.[000145] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back end component, e.g., as a data server, or that includes a middleware component, e.g., an application server, or that includes a front end component, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the subject matter described in this specification, or any combination of one or more such back end, middleware, or front end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include a local area network (LAN) and a wide area network (WAN), e.g., the Internet.[000146] The computing system can include clients and servers. A client and server are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server transmits data, e.g., an HTML page, to a user device, e.g.. for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. Data generated at the user device, e g , a result of the user interaction, can be received at the server from the device.[000147] 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.[000148] 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 systemmodules 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 be integrated together in a single software product or packaged into multiple software products.[000149] This specification also provides the subject-matter of the following clauses: [000150] Clause 1. A computer-implemented method comprising:[000151] obtaining data specifying a neural network that has been pre-trained to perform one or more machine learning tasks, wherein the neural network comprises an original set of parameters having pre-trained values;[000152] receiving adaptation data for a particular task, the adaptation data comprising a reference input for the particular task;[000153] generating a customized neural network that comprises (i) the original set of parameters having the pre-trained values and (ii) one or more additional sets of parameters; and[000154] adapting the customized neural network to the particular task, wherein the adapting comprises:[000155] processing, using an auxiliary neural network, at least the reference input to generate new values for the one or more additional sets of parameters; and [000156] training, using at least the reference input, the customized neural network starting from the new values for the one or more additional sets of parameters to determine updated values for the one or more additional sets of parameters, while holding the pre-trained values for the original set of parameters fixed.[000157] Clause 2. The method of clause 1, wherein generating the new values for the one or more additional sets of parameters comprises:[000158] processing the reference input using an encoder neural network of the auxiliary neural network to generate an encoded representation of the reference input;[000159] processing (i) the encoded representation of the reference input and (ii) an initial latent representation of initial values for the one or more additional sets of parameters using a decoder neural network of the auxiliary neural network to generate an updated latent representation; and[000160] processing the updated latent representation using one or more output layers of the auxiliary neural network to generate the new values for the one or more additional sets of parameters.[000161] Clause 3. The method of clause 2, wherein generating the updated latent representation comprises iteratively updating the initial latent representation using the decoder neural network over multiple update steps.[000162] Clause 4. The method of any one of clauses 1-3, wherein generating the customized neural network comprises:[000163] generating, from a matrix having values that represent parameters of a particular neural network layer and based on rank decomposition, a pair of two decomposed matrices that include a matrix having values that represent an additional set of parameters.[000164] Clause 5. The method of clause 4, wherein generating the customized neural network comprises:[000165] determining a total number of the additional set of parameters based on a total number of the parameters of the particular neural network layer.[000166] Clause 6. The method of any one of clause 4 or 5, wherein the particular neural network layer is a linear layer, a self-attention layer, or a cross-attention layer.[000167] Clause 7. The method of any one of clauses 1-6, wherein the particular task is a task to generate new images that depict a particular subject instance, and the reference input is a reference image that depict the particular subject instance.[000168] Clause 8. The method of any one of clauses 1-7, wherein the neural network is a diffusion neural network and the customized neural network is a customized diffusion neural network.[000169] Clause 9. The method of clause 8, wherein the diffusion neural network is a text-to-image diffusion neural network that has been pre-trained to generate images conditioned on text prompts.[000170] Clause 10. The method of any one of clause 8 or 9, wherein the encoder neural network is a vision Transformer (ViT) encoder neural network.[000171] Clause 11. The method of any one of clause 8 or 9, wherein the decoder neural network is a Transformer decoder neural network.[000172][000173] Clause 12. The method of any one of clauses 8-11, wherein training the customized diffusion neural network comprises:[000174] processing, using the customized diffusion neural network and in accordance with the new values for the one or more additional sets of parameters, a diffusion network input comprising a noisy version of the reference image to generate a predicted reference image; and[000175] determining one or more updates to the new values for the one or more additional sets of parameters based on minimizing a diffusion de-noising loss term that measures a difference between (i) the predicted reference image and (ii) the reference image. [000176] Clause 13. The method of any one of clauses 8-12, further comprising training the auxiliary neural network together with the customized diffusion neural network on a training image dataset to determine trained values for parameters of the auxiliary neural network, wherein the training comprises, for each training image included in the training image dataset:[000177] generating, based on the training image and in accordance with current values for the set of parameters of the auxiliary neural network, current values for the one or more additional sets of parameters of the customized diffusion neural network;[000178] processing, using the customized diffusion neural network and in accordance with the current values for the one or more additional sets of parameters, a diffusion network input comprising a noisy version of the training image to generate a predicted training image;[000179] computing a fine-tuning loss function that comprises a diffusion denoising loss term that measures a difference between (i) the predicted training image and (ii) the training image, and a parameter space loss term that measures a difference between (i) initial values for the one or more additional sets of parameters and (ii) the current values for the one or more additional sets of parameters; and[000180] determining, based on computing a gradient of the fine-tuning loss function with respect to the parameters of the auxiliary neural network, one or more updates to the current values for the parameters of the auxiliary neural network.[000181] Clause 14. A computer-implemented method comprising: [000182] receiving a reference image that depicts a subject instance;[000183] adapting, using at least the reference image and the method of any one of the preceding clauses, a pre-trained neural network to obtain a customized neural network; and[000184] generating a plurality of output images that each depict the subject instance using the customized neural network.[000185] Clause 15. The method of clause 14, wherein pre-trained neural network comprises a diffusion neural network.[000186] Clause 16. The method of clause 15, wherein the diffusion neural network comprises a text-to-image diffusion neural network.[000187] Clause 17. The method of clause 16, wherein generating the plurality of output images comprises:[000188] providing, as input to the text-to-image diffusion neural network, a text prompt that comprises a predetermined identifier of the subject instance.[000189] Clause 18. The method of any one of clauses 14-17, wherein the plurality of output images depict the subject instance under different various contexts, poses, views, and / or lighting conditions.[000190] Clause 19. A system comprising one or more computers and one or more storage devices storing instructions that are operable, when executed by the one or more computers, to cause the one or more computers to perform the operations of the respective method of any preceding clause.[000191] Clause 20. A computer storage medium encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform the operations of the respective method of any preceding clause.[000192] 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.[000193] What is claimed is:

Claims

CLAIMS1. A computer-implemented method comprising: obtaining data specifying a neural network that has been pre-trained to perform one or more generative tasks, wherein the neural network comprises an original set of parameters having pre- trained values; receiving adaptation data for a customized generative task, the adaptation data comprising a reference input for the customized generative task; generating a customized neural network that comprises (i) the original set of parameters having the pre-trained values and (ii) one or more additional sets of parameters; and adapting the customized neural network to the customized generative task, wherein the adapting comprises: processing, using an auxiliary neural network, at least the reference input to generate new values of the one or more additional sets of parameters; and training, using at least the reference input, the customized neural network starting from the new values of the one or more additional sets of parameters to determine updated values of the one or more additional sets of parameters, while holding the pre-trained values of the original set of parameters fixed.

2. The method of claim 1, wherein generating the new values of the one or more additional sets of parameters comprises: processing the reference input using an encoder neural network of the auxiliary' neural network to generate an encoded representation of the reference input; processing (i) the encoded representation of the reference input and (ii) an initial latent representation of initial values of the one or more additional sets of parameters using a decoder neural network of the auxiliary neural network to generate an updated latent representation; and processing the updated latent representation using one or more output layers of the auxiliary neural network to generate the new values of the one or more additional sets of parameters.

3. The method of claim 2, wherein generating the updated latent representation comprises iteratively updating the initial latent representation using the decoder neural network over multiple update steps.

4. The method of any one of claims 1-3, wherein generating the customized neural network comprises: generating, from a matrix having values that represent parameters of a particular neural network layer and based on rank decomposition, a pair of two decomposed matrices having values that represent an additional set of parameters.

5. The method of claim 4, wherein generating the customized neural network comprises: determining a total number of the additional set of parameters based on a total number of the parameters of the particular neural network layer.

6. The method of any one of claim 4 or 5, wherein the particular neural network layer is a linear layer, a self-attention layer, or a cross-attention layer.

7. The method of any one of claims 1-6, wherein the customized generative task is a task to generate new images that depict a particular subject instance, and the reference input is a reference image that depict the particular subject instance.

8. The method of any one of claims 1-7, wherein the neural network is a diffusion neural network and the customized neural network is a customized diffusion neural network.

9. The method of claim 8, wherein the diffusion neural network is a text-to-image diffusion neural network that has been pre-trained to generate images conditioned on text prompts.

10. The method of any one of claim 8 or 9, wherein the encoder neural network is a vision Transformer (ViT) encoder neural network.

11. The method of any one of claim 8 or 9, wherein the decoder neural network is aTransformer decoder neural network.

12. The method of any one of claims 8-11, wherein training the customized diffusion neural network comprises: processing, using the customized diffusion neural network and in accordance with the new values of the one or more additional sets of parameters, a diffusion network input comprising a noisy version of the reference image to generate a predicted reference image; and determining one or more updates to the new values of the one or more additional sets of parameters based on minimizing a reconstruction loss term that measures a difference between (i) the predicted reference image and (ii) the reference image.

13. The method of any one of claims 8-12, further comprising training the auxiliary neural network together with the customized diffusion neural network on a training image dataset to determine trained values of parameters of the auxiliary neural network, wherein the training comprises, for each training image included in the training image dataset: generating, based on the training image and in accordance with current values of the set of parameters of the auxiliary neural network, current values of the one or more additional sets of parameters of the customized diffusion neural network; processing, using the customized diffusion neural network and in accordance with the current values of the one or more additional sets of parameters, a diffusion network input comprising a noisy version of the training image to generate a predicted training image; computing an auxiliary' loss function that comprises a reconstruction loss term that measures a difference between (i) the predicted training image and (ii) the training image, and a parameter space loss term that measures a difference between (i) initial values of the one or more additional sets of parameters and (ii) the current values of the one or more additional sets of parameters; and determining, based on computing a gradient of the auxiliary loss function with respect to the parameters of the auxiliary’ neural network, one or more updates to the current values of the parameters of the auxiliary neural network.

14. A computer-implemented method comprising: receiving a reference image that depicts a subject instance; adapting, using at least the reference image and the method of any one of the preceding claims, a pre-trained neural network to obtain a customized neural network; andgenerating a plurality of output images that each depict the subject instance using the customized neural network.

15. The method of claim 14, wherein pre-trained neural network comprises a diffusion neural network.

16. The method of claim 15, wherein the diffusion neural network comprises a text-to- image diffusion neural network.

17. The method of claim 16, wherein generating the plurality of output images comprises: providing, as input to the text-to-image diffusion neural network, a text prompt that comprises a predetermined identifier of the subject instance.

18. The method of any one of claims 14-17, wherein the plurality of output images depict the subject instance under different various contexts, poses, views, and / or lighting conditions.

19. A system comprising one or more computers and one or more storage devices storing instructions that are operable, when executed by the one or more computers, to cause the one or more computers to perform the operations of the respective method of any preceding claim.

20. A computer storage medium encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform the operations of the respective method of any preceding claim.

21. A computer storage medium storing one or more additional sets of parameters of a customized neural network, the one or more additional sets of parameters having updated values that are determined by the process of any one of claims 1-18.

Images