Joint example selection for multimodal learning

Abstract

Methods, systems, and apparatus for training a machine learning model through contrastive learning using a batch that includes a subset of training examples from a training dataset. In one aspect, a method includes obtaining a training dataset including multiple training examples and generating a batch comprising a subset of training examples. The method includes generating the batch by selecting the subset of training examples based on, for each training examples, a respective conditional score that measures a benefit to the training of the machine learning model of including the training example in the batch given that at least a subset of the training examples of the training dataset are also included in the batch. The method further includes training the machine learning model on a contrastive loss using the batch.

Description

[0001] JOINT EXAMPLE SELECTION FOR MULTIMODAL LEARNING

[0002] CROSS-REFERENCE TO RELATED APPLICATION

[0003] [1] This application claims priority to U.S. Provisional Patent Application No. 63 / 650,894, filed on May 22, 2024. The disclosure of the prior application is considered part of and is incorporated by reference in its entirety in the disclosure of this application.

[0004] BACKGROUND

[0005] [2] This specification relates to processing inputs using neural networks to generate output sequences.

[0006] [3] 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.

[0007] SUMMARY

[0008] [4] This specification describes a system implemented as computer programs on one or more computers in one or more locations that trains a machine learning model through contrastive learning using batches that each include a respective subset of training examples from a training dataset.

[0009] [5] In particular, in one aspect, the system can obtain a training dataset including multiple training examples. The system can generate a batch including a subset of the multiple training examples by selecting the subset of training examples based on, for each training example, a respective conditional score that measures a benefit to the training of the machine learning model of including the training example in the batch given at least a subset of the training examples of the training dataset

[0010] [6] In some implementations, generating the batch further includes adding a respective set of training examples to the batch at each of multiple iterations.

[0011] [7] In some implementations, at one or more of the multiple iterations, the system can compute a respective conditional score for each of the training examples that are not included in the batch as of the iteration, where the respective conditional score for each of the training examples that are not included in the batch measures a benefit to the training of the machine learning model of including the training example in the batch given the training examples that are included in the batch as of the iteration; and the system can select the respective set of training examples to be added to the batch at the iteration based on the respective conditional scores for the training examples that are not included in the batch.

[0012] [8] In some implementations, selecting the respective set of training examples to be added includes determining a respective probability for each of the training examples that are not included in the batch using the respective conditional scores for the training examples that are not included in the batch; and sampling the respective set of training examples in accordance with the respective probabilities.

[0013] [9] In some implementations, the respective conditional score for each of the training examples includes a first score that measures a contrastive loss of the machine learning model computed for the training example given that at least the subset of the training examples of the training dataset are also included in the batch.

[0014]

[0010] In some implementations, each training example includes a respective first input and a respective second input, and the method further includes, for each training example, determining the first score by processing the respective first input and the respective second input in the training example using an online approximation of the machine learning model to generate respective embeddings of the respective first input and the respective second input, and computing the contrastive loss using the respective embeddings.

[0015] 1111 In some implementations, processing the respective first input and the respective second input in the training example using an online approximation of the machine learning model includes processing a lower resolution version of the respective first input in the training example using the machine learning model.

[0016]

[0012] In some implementations, processing the respective first input and the respective second input in the training example using an online approximation of the machine learning model includes processing only a randomly-selected portion of the respective first input in the training example using the machine learning model.

[0017]

[0013] In some implementations, processing the respective first input and the respective second input in the training example using an online approximation of the machine learning model includes processing the respective first input in the training example using the machine learning model with one or more layers of the machine learning model disabled.

[0018]

[0014] In some implementations, for each training example, the contrastive loss depends on (i) a similarity between the first and second inputs in the training example and (ii) a respective similarity between the first input in the training example and each second input in each training example that is included in the at least the subset of the training examples of the training dataset.

[0019]

[0015] In some implementations, computing the respective conditional scores for each of the training examples includes processing the respective first inputs using a first encoder on a first set of processing nodes, where the first encoder belongs to the machine learning model; and processing the respective second inputs using the second encoder on a second set of processing nodes, where the second encoder belongs to the machine learning model.

[0020]

[0016] In some implementations, computing the respective conditional scores for each of the training examples includes processing the respective first inputs using a first encoder on a first set of processing nodes, where the first encoder belongs to the machine learning model being trained; and processing the respective second inputs using the second encoder on a second set of processing nodes, where the second encoder belongs to the machine learning model being trained.

[0021]

[0017] In some implementations, the respective conditional score for each of the training examples includes a second score that measures a loss of a reference machine learning model computed for the training example given that at least the subset of the training examples of the training dataset are also included in the batch.

[0022]

[0018] In some implementations, for each training example in the subset, respective outputs of the reference machine learning model for the training example have been precomputed and stored in a cache, and where the method further includes computing the second score using outputs retrieved from the cache.

[0023]

[0019] In some implementations, the contrastive loss function is a softmax contrastive loss function.

[0024]

[0020] In some implementations, the contrastive loss function is a sigmoid contrastive loss function.

[0025]

[0021] In some implementations, for each training example, the respective first input is of a first modality and the respective second input is of a second, different modality.

[0026]

[0022] In some implementations, the first modality is one of an image, audio, video or text.

[0027]

[0023] In some implementations, the second modality is one of an image, audio, video, or text.

[0028]

[0024] In another aspect, the system can perform a method by one or more computers to perform a machine learning task using a trained machine learning model.

[0025] In particular, the system can receive an input for a particular machine learning task, and the system can generate one or more embedding outputs by processing the input using a trained machine learning model, where the trained machine learning model has been trained by performing operations including: obtaining a training dataset including multiple training examples; generating a batch comprising a subset of the multiple training examples, the generating including selecting the subset of training examples based on, for each training example, a respective conditional score that measures a benefit to the training of the machine learning model of including the training example in the batch given that at least a subset of the training examples of the training dataset are also included in the batch; and training the machine learning model on a contrastive loss function using the batch. The system can then perform the particular machine learning task using the one or more embedding outputs.

[0029]

[0026] In some implementations, performing the particular machine learning task using the one or more embedding outputs includes at least one of: performing audio compression, performing image compression, or performing video compression.

[0030]

[0027] In some implementations, the input is an image, and performing the particular machine learning task using the one or more embedding outputs includes performing object detection of one or more obj ects in the image.

[0031]

[0028] In some implementations, the input is an image, and performing the particular machine learning task using the one or more embedding outputs includes performing image segmentation for the image.

[0032]

[0029] In some implementations, the input is an image, and performing the particular machine learning task using the one or more embedding outputs includes performing image segmentation for the image.

[0033]

[0030] In some implementations, the input is an image, and performing the particular machine learning task using the one or more embedding outputs includes performing image captioning for the image.

[0034]

[0031] In some implementations, the input is audio data representing speech, and performing the particular machine learning task using the one or more embedding outputs includes performing speech transcription for the audio data.

[0035]

[0032] In some implementations, performing the particular machine learning task using the one or more embedding outputs includes performing command generation for a robot to perform a physical task.

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

[0036]

[0034] Training large-scale machine learning models, particularly those employing contrastive learning objectives, presents significant technical challenges. One such challenge is the computational cost (e.g., processing time on specialized hardware like GPUs or TPUs) and energy consumption required for processing vast training datasets. Another key technical challenge specific to contrastive learning is that the effectiveness of a given training example for model improvement is not an intrinsic property but is highly dependent on the other examples present in the same training batch. This is because the contrastive loss function evaluates relative similarities and dissimilarities between embeddings of examples within the batch. Conventional data selection or curation methods, which often assess examples individually or based on global statistics, frequently fail to address this critical contextual dependency efficiently or effectively

[0037]

[0035] Conventional systems can perform data curation for training a machine learning model, which involves selecting '’high-quality" training examples in order to more efficiently train the model. Some existing systems rely on manual curation by manually selecting data points, which can be relatively unscalable and time-consuming. Some other systems can select the training examples on a point-by-point basis using one or more scoring metrics. However, selecting individual training examples from a large dataset in these ways can be computationally expensive, and the scoring metrics may not accurately represent the benefit of training the machine learning model using the particular data point given the context of the other data points in the dataset. That is, these methods may perform poorly in the context of contrastive learning, where the loss for any given training example is not determined in isolation but depends critically on the other examples present in the same batch. Since contrastive loss measures relative similarity, the benefit of including a specific example can only be meaningfully assessed in relation to the surrounding examples. As such, conventional selection strategies may overlook the interactions between examples that are particularly helpful for effective contrastive learning and yield batches that are suboptimal for training.

[0038]

[0036] This sub-optimal batch composition, resulting from a failure to account for intrabatch example interactions, can lead to slower convergence during the training of the machine learning model, demanding more iterations and thus greater computational resources (e.g.. GPU hours, memory bandwidth). Ultimately, this can result in models that achieve lower performance on their intended technical tasks, such as reduced accuracy in image classification, lower precision / recall in text-to-image retrieval, or the generation of less discriminative embeddings for downstream technical applications like anomaly detection in sensor data or control signal generation.

[0039]

[0037] In contrast, the described techniques provide for generating a batch that includes a subset of training examples from a training dataset based on a respective conditional score for the subset of training examples, such that the system jointly selects high quality training examples for the batch with reference to the other training examples. Thus, by curating high quality data using the conditional scores, the system can train the machine learning model using a relatively lower amount of training examples, resulting in increased efficiency and improved training of the machine learning model.

[0040]

[0038] This method of joint example selection efficiently curates high-quality training batches for contrastive learning, leading to reduced computational resources for training, faster convergence to a high-performance model for a more effective trained machine learning model. The strategic use of an online approximation model and a reference model to determine components of the conditional score contribute to these advantages by facilitating computationally tractable yet effective data-driven batch construction

[0041]

[0039] In particular, the conditional score represents the benefit of the training example in training the machine learning model given a subset of other training examples included in the batch. The system can perform multiple iterations to select the training examples for the batch while taking into account the already selected data points in the batch. In some examples, the conditional score represents a “learnability” of each training example based on a combination of a first score and a second score.

[0042]

[0040] The first score represents data that is “hard” to leam for the machine learning model (e.g., training examples with relatively high contrastive loss). Thus, by including the first score as part of the conditional score, the system can discard “trivial” training examples that do not benefit the training of the machine learning model in comparison with the “hard” training examples. The second score represents data that is “easy” to leam for a reference model of the machine learning model (e.g., training examples with relatively low loss). By including the second score, the system can leverage the already generated outputs of the reference model to identify training examples that can benefit the training of the machine learning model while holding the machine learning model fixed.

[0043]

[0041] In some examples, the system more efficiently generates the conditional scores by determining the first score using an online approximation of the machine learning model. In particular, the system can process streamlined versions of the inputs of each training example to generate respective embeddings of the training example. For example, the training example can include an image and text, and the system can generate the first score using the online approximation by processing a lower resolution version of the image, a random patch of the image, or a combination thereof. In some examples, the system can disable one or more layers of the machine learning model disabled to reduce computational complexity and reduce latency in computing the conditional score.

[0044]

[0042] Overall, the described techniques allow for generating a batch for training a machine learning model that can enable more effective training of the machine learning model on a relatively small amount of training examples based on determining a conditional score of the training examples. Additionally, the described system can leverage the reference model, the online approximation of the machine learning model, or both to compute the conditional score of each training example and efficiently generate the batch.

[0045]

[0043] 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.

[0046] BRIEF DESCRIPTION OF THE DRAWINGS

[0047]

[0044] FIG. 1 shows an example training system.

[0048] [451 FIG. 2 shows an example batch generation system.

[0049]

[0046] FIG. 3 is a flow' diagram of an example process for training a machine learning model through contrastive learning using a batch that includes a subset of training examples from a training dataset.

[0050]

[0047] FIG. 4 is a diagram of the results of implementing joint example selection for batch learnability.

[0051]

[0048] FIG. 5 is a diagram of the results of implementing joint example selection for multimodal learnability.

[0052]

[0049] Like reference numbers and designations in the various drawings indicate like elements.

[0053] DETAILED DESCRIPTION

[0054]

[0050] FIG. 1 shows an example system 100. The system 100 is an example of a system implemented as computer programs on one or more computers in one or more locations, in which the systems, components, and techniques described below can be implemented.

[0051] The system 100 includes a batch generation system 102. a training system 104, and a training database 114.

[0055]

[0052] The system 100 trains a machine learning 112 over multiple training iterations using the batch generation system 102 and the training system 104.

[0056]

[0053] In particular, at each training iteration, the system 100 samples one or more training examples 118 to generate a batch 122, and the system 100 trains the machine learning model 1 12 on the batch 122. In some examples, the system 100 can sample from a same training data set (e.g., the entire set of training examples 118) at each training iteration to generate the batch 122 for training the machine learning model 112. In some other examples, the system 100 can sample from different subsets of the training data set at different training iterations to generate the batch 122 for training the machine learning model 112.

[0057]

[0054] The input training data itself, particularly when sourced from uncurated or diverse origins such as time-series sensor readings from an industrial process operating under variable conditions, often presents inherent technical challenges for effective machine learning model training due to noise, redundancy, or class imbalance. The described joint example selection method, by dynamically and contextually curating batches based on the conditional score, can improve the robustness of the training process against such data imperfections. For instance, by carefully balancing 'hard' examples with 'easy-for- reference' examples, the system can avoid overfitting to noisy artifacts while still facilitating the model 1 12 in extracting valuable semantic signals from challenging yet informative data. This leads to a more robustly trained machine learning model 112, better equipped to handle real-world data variations when deployed.

[0058]

[0055] The training database 114 can store training data for training a machine learning model. In particular, the training database 1 14 stores a plurality of training examples 118. Each training example 118 can include two inputs. The first input is of a first modality, and the second input is of a second modality. The modality can be one of an image, audio, video, or text. The first modality can be the same as the second modality or different than the second modality.

[0059]

[0056] The machine learning model 112 can be an embedding neural network with any appropriate architecture that includes multiple layers and processes an input to generate an embedding output for the input. An embedding, as used in this specification, is an ordered collection of numerical values, e.g., a vector of numerical values, that has a predetermined dimensionality. The machine learning model 122 can include one or more Transformer blocks. For example, the machine learning model 112 can be a text embedding neural network that can map a text input to an embedding output. In another example, the machine learning model 112 can be an image embedding neural network that can map an input image to an embedding output. In another example, the machine learning model 112 can be a multimodal embedding neural network that can map inputs of different modalities to an embedding output.

[0060]

[0057] The batch generation system 102 can process training examples to generate a batch 122 that includes a subset of training examples.

[0061]

[0058] The batch generation system 102 includes a benefit determination system 106 that can determine a conditional benefit score 122 for each training example 118 at each training iteration and a training example selection engine 108 that can select one or more training examples 1 18 to add to the batch 122 at each training iteration.

[0062]

[0059] The training system 104 can train the machine learning model 112 through contrastive learning using the batch 122 at each training iteration. .

[0063]

[0060] Advantageously, the system jointly selects training examples 118 for the batch 122 with a relatively high learnability based on respective conditional benefit scores 120 for each set of training examples 118, allowing the system 100 to train the machine learning model 112 using a relatively lower amount of training examples 118, resulting in increased efficiency.

[0064] [611 In particular, the batch generation system 102 is configured to add a respective set of training examples 1 18 to the batch at each of multiple iterations (e g., selection iterations). The selection iterations can be different than the training iterations. At each selection iteration, the batch selection engine system 106 can compute a respective conditional benefit score 122 for each set of the training examples 118 that are not included in the batch 122 as of the selection iteration.

[0065]

[0062] The conditional benefit score 122 for a given training example 118 measures a benefit to the training of the machine learning model 112 of including the training example 118 in batch 122 given that at least a subset of the training examples 118 of the training dataset are also included in the batch 122. In particular, the conditional benefit score 122 can be computed using a learnability score, an easy-reference score, or both, as described in further detail below with reference to FIG. 2.

[0066]

[0063] A technical advantage of utilizing such a conditional benefit score 122 for the joint selection of training examples 118 is a more targeted and computationally efficient training process for the machine learning model 112. By dynamically assessing and selecting examples based on their contribution within the context of an evolving batch 122, the system can assemble batches that provide a richer and more effective learning signal for the contrastive objective per unit of computation. This technical effect can manifest in several ways: firstly, as a faster convergence of the machine learning model 112 towards a desired performance level on a specific technical task; secondly, as a reduction in the total number of training examples or training epochs required to reach that performance level, thereby saving significant computational resources like processing time, energy, and data storage; and thirdly, potentially leading to a trained model 112 exhibiting improved technical performance metrics, such as enhanced accuracy in classifying medical images, greater robustness to noisy input data in an industrial monitoring system, or more precise semantic alignment in multimodal embeddings used for cross-modal information retrieval.

[0067]

[0064] The batch selection system 108 can then select the respective set of training examples 118 to be added to the batch 122 at the selection iteration based on the respective conditional benefit scores 122 for the set of the training examples 118 that are not included in the batch

[0068]

[0065] The training system 104 can then train the machine learning model 112 on a contrastive loss function using the batch 122. The goal of the contrastive loss function is to train the machine learning model 112 so that it can embed an input of a first modality and an input of a second modality into the representation space, i.e., the space of the input of the first modality and the input of second modality embeddings, in such a way that inputs with similar semantics are mapped to nearby points regardless of their modalities.

[0069]

[0066] Thus, the system 100 can train the machine learning model 112 on a contrastive loss that encourages, for all training examples 118 in the batch 122 that include a first input and a second input, an embedding of the first input and an embedding of the second input to be closer together while being farther from all other embeddings of all other inputs of the training examples 118 in the batch 122. For example, in the case where the first input is a visual input and the second input is a text input, the contrastive loss can encourage an image embedding of the visual input and a text embedding of the text input to be closer together while being farther from all other embeddings of all other visual inputs and text segments in the batch 118. For example, the contrastive loss can be a softmax loss function or a sigmoid loss function, as described in further detail below with reference to FIG. 3.

[0070]

[0067] For example, the contrastive loss can be a softmax loss, as used in ALIGN (Scaling Up Visual and Vision-Language Representation Learning With Noisy Text Supervision, Jia et al., 2021) and PaLI (PaLI: A Jointly-Scaled Multilingual Language-Image Model, Chen et al., 2022). In another example, the contrastive loss can be a sigmoid loss (Sigmoid Loss for Language Image Pre-Training. Zhai et. al.. 2023). In another example, the contrastive loss can be a cross entropy loss, as used in SimCLR (A Simple Framework for Contrastive Learning of Visual Representations, Chen at al., 2020).

[0071]

[0068] After the machine learning model 112 has been trained, embeddings generated by the trained machine learning model 122 can be used to perform one or more downstream tasks. In particular, the system can process the embeddings generated by the trained machine learning model 122 using a downstream model for the corresponding downstream task. For example, the machine learning model 112 can be used to generate embeddings for a generation task (e.g., text generation, image generation, audio signal generation, video generation, etc.), a classification task (e.g., image classification), an object detection task, an image segmentation task, a compression task, or a prediction task.

[0072]

[0069] That is, embedding representations (e.g., embeddings) generated by the trained machine learning model 112 can be provided as input to a downstream neural network for performing a downstream task.

[0073]

[0070] For example, the embeddings generated by the machine learning model 112 can be used to train a generative neural network that generates new observations (of the same ty pe as the input observations or a different ty pe) conditioned on embeddings generated using the machine learning model 112.

[0074] 1711 As yet another example, the embeddings can be used as a representation of the observation for a multi-modal task performed by a multi-modal neural network, e.g., a representation of an image or video in visual understanding tasks, e.g., image (or video)- text retrieval tasks, image (or video) classification tasks, image (or video) captioning tasks, and visual question answering tasks. The multi-modal neural network can be. e.g., a multimodal sequence generation neural network, e.g., a multi-modal large language model (LLM), or a visual language model (VLM), or a different type of multi-modal neural network.

[0075]

[0072] For example, after training the machine learning model 112. the system can receive a query input for a downstream task. The query input will generally include a query image and, optionally, other data, e.g., one or more other images, one or more inputs of a different modality, e.g., text or audio.

[0076]

[0073] The system can process the query image using the trained machine learning model 112 to generate an embedding of the query image as a set of text tokens.

[0074] The system can then provide the embedding of the query image as input to a downstream neural network configured to perform the downstream task.

[0077]

[0075] The downstream neural network can generally be any neural network that is configured to process inputs that include text tokens from the vocabulary to generate outputs for the downstream task.

[0078]

[0076] For example, the downstream neural network can be a language model neural network, e.g., a large language model neural network (LLM), or a visual language model neural network (VLM). The LLM can be, e.g., a multi-modal model that processes inputs that include tokens representing multiple different modalities of data, or can be a uni-modal model that processes inputs that include text tokens.

[0079]

[0077] For example, the query input can include the query’ image and text and the downstream neural network can be an LLM. Thus, providing the embedding of the query image as input to the downstream neural network can include providing the embedding of the query image and the text from the query input as input to the LLM instead of directly providing the query image as part of the input. For example, the LLM can have been trained on text-only data and therefore not be able to directly process image data inputs.

[0080]

[0078] As another example, the embeddings can be provided as input to a classifier, e.g., a classification neural netw ork or other type of machine learning model, that is configured to classify the input as belonging to one or more of a set of classes, e.g., object classes.

[0081] 1791 The downstream task that is performed by the downstream neural network can be any of a variety of tasks, e.g., a multi-modal dialogue task, so that the image is part of a dialogue input submitted by a user to the system and the output generated by the downstream neural network is a response to be displayed to the user.

[0082]

[0080] Other examples of downstream tasks include multi-modal zero-shot or few-shot learning tasks.

[0083]

[0081] As one 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 generative model 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. That is, the system can process the embeddings generated by the training machine learning model 112, which represent the sequence of text, to generate the translation of the 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 languages - 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.

[0084]

[0082] 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.

[0085]

[0083] As another example, the task can be a text to speech task (e.g., speech transcription), 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.

[0086]

[0084] 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 waveform. 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.

[0087]

[0085] 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 depends 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 open-vocabulary image classification, open-vocabulary object detection, image captioning, text-based image search, image-based retrieval, and so on.

[0088]

[0086] In particular, where the input is one or more images, the system can perform object detection of one or more objects in the input images by processing the one or more embedding outputs of the machine learning model 112. and the system can output one or more location indications for the one or more objects based on the detection.

[0087] 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.

[0089]

[0088] As another example, the embeddings generated by the trained machine learning model 112 can be used as part of a generative model (e.g., a language model) to solve visual understanding tasks.

[0090]

[0089] For example, the discrete representations can be used as a representation of the image in visual understanding tasks.

[0091]

[0090] One example of a visual understanding task is an image-text retrieval task, where the input includes an image or text or both and the output is an image that is received from an image datastore.

[0092]

[0091] As yet another example, the discrete representations can be used as a representation of the image in visual understanding tasks.

[0093]

[0092] One example of a visual understanding task is an image-text retrieval task, where the input includes an image or text or both and the output is an image that is received from an image datastore.

[0094]

[0093] Another example of a visual understanding task is an image classification task, where the input is an image and the output is an identification of objects depicted in the image.

[0095]

[0094] Another example of a visual understanding task is an image captioning task, where the input is an image and the output describes in natural language the objects depicted in the image.

[0096]

[0095] Another example of a visual understanding task is a visual question answering task, where the input is an image and a query about the image and the output is a response to the query.

[0097]

[0096] As another example, the downstream task can be a video processing task that requires processing respective discrete representations generated by the image processing neural network of each video frame in an input video. For example, the task can be a video question answering task, a video classification task, an action recognition task, a video generation task, and so on.

[0097] As another example, the downstream task can be a compression task, where the input can be an audio signal, an image, or a video and the output is a compressed representation of the input. For example, the system can perform audio compression by processing the audio signal using the machine learning model 112 to generate a compressed version of the audio signal (e.g., the embedding output). In another example, the system can perform image compression by processing the image using the machine learning model 112 to generate a compressed version of the image (e.g.. the embedding output). In another example, the system can perform video compression by processing the video using the machine learning model 112 to generate a compressed version of the video (e.g., the embedding output).

[0098]

[0098] As another example, the downstream task can be an image segmentation task, where the input can be an image and the output can be one or more locations (e.g., location indications) for one or more objects in the image. In particular, the system can perform object detection of one or more objects in input images by processing the one or more embedding outputs of the machine learning model 112. and the system can output one or more location indications for the one or more objects within the image based on the detection.

[0099]

[0099] As another example, the downstream task can be a command generation task for controlling a robot to perform a physical task that requires processing the one or more embedding outputs corresponding to the commands, where the input can be data associated with the robot and the output can be one or more commands for performing the physical task.

[0100]

[0100] In practice, for any of these examples, the task to be performed by the neural network can be defined by (at least a part of) the network input, e.g., that is in the form of a prompt or a request, received by the neural network. In other w ords, the neural network will be able to perform any of these tasks when an appropriate prompt or request is received based on leveraging the embeddings of the trained machine learning model 112.

[0101]

[0101] FIG. 2 shows an example batch generation system. The batch generation system 102 is an example of a system implemented as computer programs on one or more computers in one or more locations in which the systems, components, and techniques described below are implemented.

[0102]

[0102] The batch generation system 102 is configured to generate a batch 122 that includes a subset of training examples 118 from a training dataset based on a respective conditional benefit score 122 for each set of training examples 118, such that the batch generation system 102 jointly selects high quality training examples for the batch 122 with reference to the other training examples 118. In some examples, the batch generation system can provide the training example embeddings 210.

[0103]

[0103] The batch generation system 102 includes the benefit determination system 106 that can process the training examples 118 to generate the conditional benefit scores 122 over multiple training iterations and the training example selection engine 108 configured to select the training examples to add to the batch 122 based on the conditional benefit scores 122 over multiple training iterations. That is, the conditional benefit score 122 represents a “learnability” of each training example based on a combination of a first score (e.g., a learnability score 206) and a second score (e.g., an easy reference score 208).

[0104]

[0104] In some examples, the benefit determination system 106 includes an online approximation machine learning model 202 configured to process the training examples 118 to determine the learnability' scores 206 and a reference machine learning model 204 configured to process the training examples 118 to determine the easy -reference scores 208.

[0105]

[0105] The online approximation machine learning model 202 is an online approximation of the machine learning model 1 12, where the online approximation machine learning model 2202 can estimate (e.g., generate) an output of the machine learning model 1 12 with reduced computational complexity and latency. This approximation facilitates practical and efficient application of the joint example selection strategy-, particularly when dealing with very large candidate pools of training examples 118 from which batches must be formed. Employing the online approximation model 202 may lead to a significant reduction in the computational overhead — such as the number of floating-point operations (FLOPs), memory bandwidth usage, and overall latency — used to calculate the learnability score 206 for numerous potential examples. This reduction makes the iterative batch construction process computationally feasible within acceptable timeframes for production-scale machine learning model training pipelines, which would otherw ise be prohibitively slow if the full model 112 w ere used for these intermediate score calculations. In this case, the learnability score 206 measures a contrastive loss of the online approximation of the machine learning model 112. This enables more efficient operation during training or deployment, particularly in settings where rapid inference is beneficial. In particular, like the machine learning model 112, the online approximation machine learning model 202 can have any appropriate architecture to generate embedding outputs for inputs of different modalities. Since the model 202 is a lightweight “surrogate” for the model 112, the online approximation machine learning model 202 can have less layers than the full model 112. In some examples, the system can process lower resolution version of the training examples 118 using the full model 212. While this specification describes using the online approximation machine learning model 202, in some cases, the system can use the full model 212 instead.

[0106]

[0106] The reference machine learning model 204 is a pre-trained machine learning model configured to process multimodal inputs to generate an output. That is, the reference machine learning model 204 can be a fully trained and more computationally expensive model than the online approximation machine learning model 202 that the system 100 can use as a benchmark to determine the easy -reference score 208, whereas the online approximation machine learning model 202 can be a lighter, more efficient model that can generate the learnability score 206 by using less computational resources than the reference machine learning model 204. The easy reference score 208 measures a negative loss of the reference machine learning model 204 for the training example 118 given the training examples 118 included in the batch 122.

[0107]

[0107] The system can compute the conditional benefit scores 122 using the learnability scores 206, the easy reference scores 208, or both. For example, the conditional benefit score 122 can measure a difference between the learnability score 206 and the easy- reference score 208 for each set of training examples 118, as show n by Equation 1 :

[0108] (1) sbenefitB \e, et) = shardB \e) + seasy(iB \e*) = i (B\e) - i (B\e*) where sbenefltis the conditional benefit score 122, B is a sub-batch of the training examples 118, 9 represents the parameters of the online approximation machine learning model 202, 9* represents the parameters of the reference machine learning model 204, s / mrd(F|0) represents the learnability score 206, and seasy(F|0*) represents the easy -reference score 208. The contrastive loss of the batch 122 is the sum of the conditional benefit scores 120, as shown by Equation 2: where b represents the size of a set of the training examples 118 (e.g., a sub-batch).

[0109]

[0108] In particular, the learnability score 206 represents data that is "hard" to learn for the online approximation machine learning model 202. That is, the learnability score 206 measures a contrastive loss of the online approximation machine learning model 202 computed for the training examples given at least the sub-batch of the training examples 118 of the training dataset are also included in the batch 122. Thus, by including the learnability score 206 (e.g., the first score) as part of the conditional score, the system is less likely to select trivial training examples for the batch 122 that do not benefit the training of the machine learning model 112 in comparison with the hard training examples (e.g., training examples 118 with relatively high contrastive loss).

[0110]

[0109] Such 'hard' examples represent data points or relationships where the current state of the machine learning model 112 (or its online approximation 202) demonstrates a deficiency in understanding or representation. Focusing the training process on these specific examples offers the greatest potential for substantial gradient updates and. consequently, for rapid improvement in the model's ability to discern complex patterns or subtle distinctions relevant to its technical application. This targeted learning helps the model to specifically address its weaknesses, leading to a more robust and accurate embedding space capable of, for example, finer-grained differentiation between classes in a classification task.

[0111]

[0110] In some examples, each training example 118 can include a pair of inputs. In this case, the learnability score 206 can be based on a similarity of the pair of inputs in the same training example 118 and a respective similarity between the first input of the pair of inputs in a particular training example 118 and each second input of the pair of inputs from other training examples 118 that is included in batch 122. That is, for a pair of inputs, the learnability7score depends on the contrastive loss, which depends not only on the similarity7between the two inputs in the same pair of inputs but also on the similarity between the first input in the pair and the second inputs in other pairs in the batch.

[0112]

[0111] In some examples, the machine learning model 1 12 can include a first encoder and a second encoder for processing the pair of inputs to compute the conditional benefit scores 120. In particular, the system can compute the conditional benefit scores 120 by processing the respective first inputs using the first encoder on a first set of processing nodes and processing the respective second inputs using the second encoder on a second set of processing nodes. The processing nodes can be one or more layers or one or more computation units of the machine learning model 112. In some examples, the first encoder and the second encoder can correspond to another machine learning model (e.g., not the machine learning model 112 being trained).

[0113]

[0112] In some examples, at each selection iteration, the benefit determination system 106 can process streamlined (e.g., low resolution) versions of the inputs of each training example 118 using the online approximation machine learning model 202 to generate respective embeddings of the training example 118 (e.g., training example embeddings 210). For example, the training example 118 can include an image and text, and the benefit determination system 106 can generate the learnability score 206 using the online approximation machine learning model 202 by processing a lower resolution version of an embedding of the image, a random patch of the embedding of the image, or a combination thereof.

[0114]

[0113] Processing a lower resolution version of an image, for instance, directly reduces the number of input pixels (and thus the dimensionality of the initial feature maps) that may be processed by the early layers of the online approximation machine learning model 202. leading to a near-proportional decrease in computational requirements for those layers. Similarly, processing a random patch of the image, rather than the full image, reduces the input data size. These specific input data manipulation techniques facilitate faster learnability score 206 computation by reducing the input data dimensionality or spatial extent fed into the neural network, thereby decreasing the number of multiply-accumulate operations and memory accesses.

[0115]

[0114] In some examples, for each training example in the subset, the system can precompute respective embeddings of the reference machine learning model 202 for the training example 1 18 and store the respective outputs in a cache. In this case, the system can compute the second score of the conditional score using outputs retrieved from the cache.

[0116]

[0115] In some cases, the system 100 can train the machine learning model 112 with training examples of multiple different resolutions, such that the machine learning model 1 12 is exposed to both low resolution inputs and high resolution inputs. During training, the system can split batches into low resolution and high resolution (e.g., one half of the batch is low resolution and the other half of the batch is high resolution), and the system can train the machine learning model 112 on both halves of the batch. For example, the batch generation system 102 can split a set of the training examples 118 into two halves, and the batch generation system 102 can encode each half with a different resolution. In particular, the batch generation system 102 can encode a first half at a low resolution and the second half at a higher resolution, which allows the system to efficiently process training examples while enabling the machine learning model to learn from high-resolution inputs. That is, the system can train the machine learning model 112 with training examples of multiple different resolutions, such that the machine learning model 112 is exposed to both low resolution inputs and high resolution inputs. This allows for efficient scoring of training examples and allows for increased overall efficiency in training while still allowing the system to effectively process high resolution inputs after training.

[0116] In some examples, the system can disable one or more layers of the online approximation machine learning model 202 to reduce computational complexity and reduce latency in computing the conditional benefit score 122 for the training example 118. For example, the system can disable one or more Transformer blocks of the online approximation machine learning model 202.

[0117]

[0117] F urthermore, the computational load of the online approximation machine learning model 202 may be strategically reduced by, for example, decreasing the number of selfattention heads, reducing the dimensionality of Peed-forward networks within its Transformer blocks, or by employing shallower network architectures (i.e., fewer layers) compared to the main machine learning model 112. These specific architectural modifications curtail the number of learnable parameters and matrix multiplications involved in the forward pass necessary for score computation. This may substantially lower latency and reduced processing power consumption when determining the learnability score 206, thereby facilitating a larger number of candidate examples to be evaluated per unit time or enabling faster overall batch assembly for training the machine learning model 112.

[0118]

[0118] The batch generation system 102 can compute the easy -reference score 208 by evaluating a loss betw een the set of training examples 118 not in the batch 122 and training examples 118 included in the batch 122. That is, the easy-reference score 208 represents data that is “easy7’ to learn for the reference machine learning model 204. An advantage of incorporating information about examples that are “ easy” for a strong, pre-trained reference machine learning model 204 (i.e., those yielding a low7loss for the reference model) is that these examples are likely to be well-formed, high-quality, and possess clear, learnable semantic relationships. If such examples are nevertheless challenging for the current machine learning model 112 being trained, their inclusion in the batch 122 provides strong, reliable learning signals. This effectively guides the model 112 towards emulating the capabilities of the more powerful reference model 204 concerning these specific data characteristics, acting as a form of implicit knowledge distillation or curriculum learning within the batch selection process itself. This strategy can accelerate the learning of fundamental data relationships and improve the generalization capability of the model 112 on unseen data related to its technical field of application. In particular, the easy-reference score 208 can measure a contrastive loss of a reference machine learning model 204 computed for the training example 118 given that at least the subset of the training examples 118 of the training dataset are also included in the batch 122. An example contrastive loss function is shown below in Equation 3: where y is a label indicating whether the pair of inputs of the training example 118 is similar or dissimilar, is a distance between the embeddings of the pair of inputs, and m is a hyperparameter that defines a minimum distance required for dissimilar pairs.

[0119]

[0119] That is, the easy-reference score 208 represents the contrastive loss of the reference machine learning model 24 on a given training example 118 conditioned on a subset of training examples 118 included in the batch 122. In particular, for a given training example 118, the batch generation system 102 computes the contrastive loss under the reference model when that training example 118 is grouped with other examples in the current batch 122, with lower loss indicating higher alignment and thus a higher easy-reference score. By including training examples 118 with relatively low loss under the reference machine learning model 204, the batch generation system 106 can leverage the reference model 204 to identify training examples 118 that are likely to benefit the training of the machine learning model 112 , even when the batch generation system 102 holds the machine learning model 1 12 fixed.

[0120]

[0120] In some examples, for each training example 118 in the set, the benefit determination system 106 can pre-compute and store the outputs of the reference machine learning model 204 in a cache, and the benefit determination system 106 can compute the easy-reference scores 208 using the outputs retrieved from the cache, allowing the system to compute the easy-reference scores 208 efficiently during batch selection.

[0121]

[0121] The benefit determination system 106 can then provide the conditional benefit scores 120 to the batch selection engine 106. At each of multiple selection iterations, the batch selection engine 106 can then select the set of training examples 118 to be added to the batch 122 by determining a respective probability for each of the training examples 118 that are not included in the batch 122 using the respective conditional benefit scores 120 for the set of training examples 118 and sample the sub-batch of training examples 118 based on the probabilities.

[0122]

[0122] In some examples, at each selection iteration, the batch selection engine 106 can select the set of training examples to be added to the batch 122 based on a pre-selected filtering ratio, as shown by Equation 4:

[0123] (4) f = l - b / B

[0123] , where f is the filtering ratio between the size of the sub-batch and a size of the batch 122. That is, a relatively higher filtering ratio can indicate for the batch selection engine 106 to select a relatively low percentage of training examples 118 for the batch 122.

[0124]

[0124] The batch generation system 102 can then provide the batch 122 to the training system 104 for training the machine learning model 112 on a contrastive loss function over multiple training iterations. In particular, the contrastive loss function can be a softmax contrastive loss function, as shown by Equation 5: where x represents the training example 118, z"" represents the embedding representing the image of the training example 118. and z-xtrepresents the embedding representing the text of the training example 118.

[0125]

[0125] In some other examples, the contrastive loss function can be a sigmoid contrastive loss function, as shown by Equation 5 :

[0126] , where the sigmoid contrastive loss function can allow for increased scalability for training with a particularly large dataset.

[0127]

[0126] FIG. 3 is a flow diagram of an example process for training a machine learning model through contrastive learning using a batch that includes a subset of training examples from a training dataset. 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 system, e.g., the system 100 of FIG. 1, appropriately programmed, can perform the process 300.

[0128]

[0127] The system can perform the process 300 at each of multiple training iterations to train the machine learning model 112.

[0129]

[0128] The system can obtain a training dataset including multiple training examples (302).

[0130]

[0129] The system can generate a batch including a subset of the multiple training examples (304). The system can select the subset of training examples based, for each training example, a respective conditional score at multiple selection iterations. That is, the system can recompute the score for each selection iteration (e.g., the system may not compute the score of the same training examples 118 at each training iteration). The respective conditional score measures a benefit to the training of the machine learning model of including the training example in the batch given at least a subset of the training examples of the training dataset (e.g., a sub-batch) are also included in the batch.

[0131]

[0130] In particular, the system can generate the batch by adding a respective set of training examples to the batch at each of the selection iterations. At one or more of the multiple selection iterations, the system can compute a respective conditional score for each of the training examples that are not included in the batch as of the iteration. The conditional score measures a benefit to the training of the machine learning model of including the training example in the batch given the training examples that are included in the batch as of the iteration. In particular, the conditional score includes a first score that measures a contrastive loss of the machine learning model computed for the training example given that at least the subset of the training examples of the training dataset are also included in the batch. Additionally, the conditional score can include a second score that measures a loss of a reference machine learning model computed for the training example given that at least the subset of the training examples of the training dataset are also included in the batch.

[0132]

[0131] In some examples, the training examples include a respective first input and a respective second input. For example, the first input can be of a first modality, and the second input can be a second modality, where the second modality can the same as or different from the first modality. The modalities can be an image, audio, video, or text.

[0133] 11321 The system can determine the first score for each training example by processing the respective first input and the respective second input in the training example using an online approximation of the machine learning model to generate respective embeddings of the respective first input and the respective second input and computing the contrastive loss using the respective embeddings. The contrastive loss depends on a similarity between the first and second inputs in the training example and a respective similarity between the first input in the training example and each second input in each training example included in at least the subset of the training examples of the training dataset.

[0134]

[0133] In some examples, the system can process the respective first input and respective second input in the training example using the online approximation of the machine learning model by processing a lower resolution version of the respective first input in the training example using the machine learning model.

[0135]

[0134] In some other examples, the system can process the respective first input and respective second input in the training example using the online approximation of the machine learning model by processing only a randomly-selected portion of the respective first input in the training example using the machine learning model.

[0136]

[0135] In some other examples, the system can process the respective first input and respective second input in the training example using the online approximation of the machine learning model by processing the respective first input in the training example using the machine learning model with one or more layers of the machine learning model disabled.

[0137]

[0136] In some examples, for each training example in the subset, the system can precompute respective embeddings of the reference machine learning model for the training example and store the respective outputs in a cache. In this case, the system can compute the second score of the conditional score using outputs retrieved from the cache.

[0138]

[0137] The system can then select the respective set of training examples to be added to the batch at the iteration based on the respective conditional scores for the training examples that are not included in the batch.

[0139]

[0138] In particular, at each selection iteration, the system can select the respective set of training examples by determining a respective probability for each of the training examples that are not included in the batch using the respective conditional scores for the training examples not included in the batch. The system can then sample the respective set of training examples in accordance with the respective probabilities.

[0140] | 139| The system can then train the machine learning model on a contrastive loss function using the batch (306). In particular, the system can train the machine learning model on the contrastive loss function over multiple training iterations.

[0141]

[0140] FIG. 4 is a diagram of the results of implementing joint example selection for batch learnability.

[0142]

[0141] The graphs of FIG. 4 illustrate the effectiveness of joint example selection in generating learnable training batches by comparing batch learnability, sub-batch selection methods, and the impact of filtering ratios. Scaled variants (e.g., JEST++ / FlexiJEST++) use a larger "WebLI-curated++" dataset with approximately 600M additional filtered samples. Performance is evaluated using four standard benchmarks: ImageNet O-shot and 10-shot classification, and COCO image-to-text and text-to-image top-1 retrieval.

[0143]

[0142] In particular, the left-most graph on FIG. 4 shows jointly selecting examples yields batches with high learnability. Additionally, the middle graph and the right-most graph on FIG. 4 show that joint example selection is comparable to brute-force Gibbs sampling and increasing the filtering ratio by selecting from larger super-batches enhances batch quality. That is, models such as JEST and Flexi-JEST can be trained on a lOOM-scale "WebLI- curated" dataset, which is efficiently filtered for high text-image quality and alignment, such that joint example selection yields more learnable batches.

[0144]

[0143] FIG. 5 is a diagram of the results of implementing joint example selection for multimodal learnability.

[0145]

[0144] The graphs of FIG. 5 illustrate the results of joint example selection on multimodal learning by comparing training strategies across varying filtering ratios, super-batch sizes, and score prioritization. That is, the middle graph of FIG. 5 shows the results of prioritizing the learnability score, while the right-most graph of FIG. 5 shows the results of prioritizing the easy-reference score. In particular, training on the most learnable sub-batch selected from varying sizes of super-batches significantly improves training speed.

[0146]

[0145] Thus, jointly selecting batches based on learnability outperforms independent example selection, especially at higher filtering ratios, where independent selection leads to performance regression. Additionally, this technique enhances easy-reference prioritization with an increased filtering ratio.

[0147]

[0146] 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.

[0148]

[0147] 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.

[0149]

[0148] 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, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

[0150]

[0149] 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.

[0151]

[0150] 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.

[0152]

[0151] 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.

[0153]

[0152] 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.

[0154]

[0153] 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 a central processing unit for performing or executing instructions and one or more memory 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.

[0155]

[0154] 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.

[0156]

[0155] 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.

[0157]

[0156] 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.

[0158]

[0157] Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework.

[0159]

[0158] 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 amiddleware 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.

[0160]

[0159] 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-serv er 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.

[0161]

[0160] 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.

[0162]

[0161] 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 be integrated together in a single software product or packaged into multiple software products.

[0163]

[0162] 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.

Claims

WHAT IS CLAIMED IS:

1. A method performed by one or more computers and for training a machine learning model through contrastive learning, the method comprising: obtaining a training dataset comprising a plurality of training examples; generating a batch comprising a subset of the plurality of training examples, the generating comprising selecting the subset of training examples based on, for each training example, a respective conditional score that measures a benefit to the training of the machine learning model of including the training example in the batch given that at least a subset of the training examples of the training dataset are also included in the batch; and training the machine learning model on a contrastive loss function using the batch.

2. The method of claim 1, wherein generating the batch further comprises: adding a respective set of training examples to the batch at each of a plurality of iterations.

3. The method of claim 2, further comprising: at one or more of the plurality of iterations: computing a respective conditional score for each of the training examples that are not included in the batch as of the iteration, wherein the respective conditional score for each of the training examples that are not included in the batch measures a benefit to the training of the machine learning model of including the training example in the batch given the training examples that are included in the batch as of the iteration; and selecting the respective set of training examples to be added to the batch at the iteration based on the respective conditional scores for the training examples that are not included in the batch.

4. The method of claim 3, wherein selecting the respective set of training examples to be added comprises: determining a respective probability for each of the training examples that are not included in the batch using the respective conditional scores for the training examples that are not included in the batch; and sampling the respective set of training examples in accordance with the respective probabilities.

5. The method of any preceding claim, wherein the respective conditional score for each of the training examples comprises: a first score that measures a contrastive loss of the machine learning model computed for the training example given that at least the subset of the training examples of the training dataset are also included in the batch.

6. The method of claim 5, wherein each training example comprises a respective first input and a respective second input, further comprising: for each training example, determining the first score, comprising: processing the respective first input and the respective second input in the training example using an online approximation of the machine learning model to generate respective embeddings of the respective first input and the respective second input, and computing the contrastive loss using the respective embeddings.

7. The method of claim 6, wherein processing the respective first input and the respective second input in the training example using an online approximation of the machine learning model comprises: processing a lower resolution version of the respective first input in the training example using the machine learning model.

8. The method of claim 6, wherein processing the respective first input and the respective second input in the training example using an online approximation of the machine learning model comprises: processing only a randomly-selected portion of the respective first input in the training example using the machine learning model.

9. The method of claim 6, wherein processing the respective first input and the respective second input in the training example using an online approximation of the machine learning model comprises: processing the respective first input in the training example using the machine learning model with one or more layers of the machine learning model disabled.

10. The method of any one of claims 5-9, wherein, for each training example, the contrastive loss depends on (i) a similarity between the first and second inputs in the training example and (ii) a respective similarity between the first input in the training example and each second input in each training example that is included in the at least the subset of the training examples of the training dataset.

11. The method of any of claims 5-10, wherein computing the respective conditional scores for each of the training examples comprises: processing the respective first inputs using a first encoder on a first set of processing nodes, wherein the first encoder belongs to the machine learning model; and processing the respective second inputs using a second encoder on a second set of processing nodes, wherein the second encoder belongs to the machine learning model.

12. The method of any of claims 5-10, wherein training of the machine learning model comprises: processing each first input of the generated batch using a first encoder on a first set of processing nodes, wherein the first encoder belongs to the machine learning model being trained; and processing each second input of the generated batch using the second encoder on a second set of processing nodes, wherein the second encoder belongs to the machine learning model being trained.

13. The method of any preceding claim, wherein the respective conditional score for each of the training examples comprises: a second score that measures a loss of a reference machine learning model computed for the training example given that at least the subset of the training examples of the training dataset are also included in the batch.

14. The method of claim 13, wherein, for each training example in the subset, respective outputs of the reference machine learning model for the training example have been precomputed and stored in a cache, and wherein the method further comprises: computing the second score using outputs retrieved from the cache.

15. The method of any preceding claim, wherein the contrastive loss function is a softmax contrastive loss function.

16. The method of any preceding claim, wherein the contrastive loss function is a sigmoid contrastive loss function.

17. The method of any preceding claim when dependent on claim 7, wherein, for each training example, the respective first input is of a first modality and the respective second input is of a second, different modality.

18. The method of claim 17, wherein the first modality is one of an image, audio, video, or text.

19. The method of claim 17 or claim 18. wherein the second modality is one of an image, audio, video, or text.

20. A method performed by one or more computers, the method comprising: receiving an input for a particular machine learning task; generating one or more embedding outputs by processing the input using a trained machine learning model, wherein the trained machine learning model has been trained by performing operations comprising: obtaining a training dataset comprising a plurality of training examples; generating a batch comprising a subset of the plurality of training examples, the generating comprising selecting the subset of training examples based on, for each training example, a respective conditional score that measures a benefit to the training of the machine learning model of including the training example in the batch given that at least a subset of the training examples of the training dataset are also included in the batch; and training the machine learning model on a contrastive loss function using the batch; andperforming the particular machine learning task using the one or more embedding outputs.

21. The method of claim 20, wherein performing the particular machine learning task using the one or more embedding outputs comprises at least one of: performing audio compression, performing image compression, or performing video compression.

22. The method of any preceding claim, wherein the input is an image, and wherein performing the particular machine learning task using the one or more embedding outputs comprises: performing object detection of one or more objects in the image.

23. The method of any preceding claim, wherein the input is an image, and wherein performing the particular machine learning task using the one or more embedding outputs comprises: performing image segmentation for the image.

24. The method of any preceding claim, wherein the input is an image, and wherein performing the particular machine learning task using the one or more embedding outputs comprises: performing image captioning for the image.

25. The method of any preceding claim, wherein the input is audio data representing speech, and wherein performing the particular machine learning task using the one or more embedding outputs comprises: performing speech transcription for the audio data.

26. The method of any preceding claim, wherein performing the particular machine learning task using the one or more embedding outputs comprises: performing command generation for a robot to perform a physical task.

27. A system comprising: one or more computers; and one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform operations of the respective method of any one of claims 1-26.

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

Images