Distilling uncertainties into machine learning models

Abstract

A method performed by one or more computers for training a machine learning model to determine a representation of a data item. The method comprises: obtaining training examples that each comprise (i) a training input comprising a set of data items and (ii) a training output identifying a target subset of the set of data items; processing the training examples to determine, for each training example, a corresponding target probability for the target subset given the set of data items of the training input. The machine learning model determines representations of the data items and processes the representations to determine a probability for the target subset. The machine learning model is trained by optimizing an objective function that compares the predicted probabilities with the corresponding target probabilities of the training examples.

Description

Attorney Docket No.45288-0436WO1 DISTILLING UNCERTAINTIES INTO MACHINE LEARNING MODELS BACKGROUND

[0001] This specification relates to processing data using machine learning models.

[0002] 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 values of a respective set of parameters. SUMMARY

[0003] This specification describes a system and method, implemented as computer programs on one or more computers in one or more locations, that is configured to learn to generate representations of data items, such as image, audio, or text data items, through a supervised learning process that takes into account uncertainties in the training data. As a result of the supervised learning process, the representations of the data items generated by the system and method can be better aligned with the training data, such that features and characteristics of the data items that are helpful in modelling the uncertainties in the training data are included in the generated representations. Better alignment of the representations of the data items can improve performance when the representations are used by a machine learning model to perform a downstream task, such as an object recognition or image segmentation task performed by a vision model.

[0004] For example, the representations of the data items generated by the system and method can be aligned with representations of the data items that were used to generate the training data. As one particular example, in the case of human-generated training data, the representations of the data items generated by the system and method can be aligned with (e.g., mental) representations of the data items, such as important features, objects or events included in the data items, that were used by humans when generating the training data, e.g., by performing a data item similarity task.

[0005] Therefore, in one aspect, there is provided a method performed by one or more computers for training a machine learning model to determine a representation of a data item. The methodAttorney Docket No.45288-0436WO1 comprises: obtaining training examples that each comprise (i) a training input comprising a set of data items and (ii) a training output identifying a target subset of the set of data items; and processing the training examples to determine, for each training example, a corresponding target (e.g., posterior) probability for the target subset given the set of data items of the training input. The method further comprises training the machine learning model to determine a representation of a data item. The training comprises, for each of the training examples: using the machine learning model to determine, in accordance with values of parameters of the machine learning model, a representation for each of the data items in the set of data items of the training input; and processing the representations of the data items to determine a predicted probability for the target subset given the set of data items of the training input. The training further comprises adjusting the values of the parameters of the machine learning model to optimize an objective function that compares the predicted probabilities with the corresponding target probabilities.

[0006] In some implementations, the training examples can be generated by performing a data item similarity task that involves assessing similarities (or equivalently, differences) between data items. The target probability can, for example, indicate a probability that the target subset would be selected from the set of data items in the data item similarity task that was used to generate the training examples. For example, the similarity task may involve determining a data item that is the “odd-one-out” from a set of data items. In that case, each training input can comprise a respective three or more data items and the corresponding training output may identify one of the three or more data items as being the odd-one-out. That is, the training output may identify one of the data items as being less similar to the other data items in the set than the other data items are to one another. The training output may alternatively identify the other data items instead of the odd-one-out data item as this conveys the same information.

[0007] In some implementations, the training examples are human generated, e.g., the similarity task can be performed by one or more human contributors. The machine learning model can then learn to generate representations that are aligned with human preferences or judgments. For example, the machine learning model can learn to generate representations that include features or characteristics of the data items that are the same as or similar to the features or characteristics of the data items used by the human contributors to assess the similarities of the data items. For example, the target probability can then indicate a probability that the target subset would be selected from the set of data items by a human performing the data item similarity task. InAttorney Docket No.45288-0436WO1 general, each target probability can be less than one, reflecting the uncertainty inherent in the similarity task or differences between different persons who are performing the task.

[0008] In some implementations, the target probability for each training example can be an estimate of the probability that a human contributor would identify an odd-one-out data item that is the same as the data item identified in the training output of the training example. For example, some of the training examples can have training inputs that comprise the same or a similar set of data items, but different corresponding training outputs, e.g., different odd-one-out data items selected from the same or similar sets of data items. The target probabilities can be determined from a statistical analysis of the training examples. As one example, a Bayesian inference model can be used to determine the target probabilities, such as is described in Muttenthaler et al. “VICE: Variational Interpretable Concept Embeddings” in Advances in Neural Information Processing Systems, 35:33661-33675 (2022).

[0009] In general, each subset of data items may be a proper subset of the set of data items, i.e., a subset that comprises some, but not all of the data items in the set. Each subset may comprise a plurality of data items. Each subset may consist of the same number of data items, or the subsets may have different numbers of data items.

[0010] In some implementations, processing the training examples to determine the corresponding target probability for the target subset given the set of data items of the training input comprises processing the training examples to determine, for each data item, a corresponding distribution over representations of the data item. The processing may then further comprise using the distributions over representations of the data items to determine the corresponding target probability for the target subset given the set of data items of the training input.

[0011] For example, each representation of a data item can comprise a plurality of representation values and the distribution over representations of the data item can comprise a corresponding representation value distribution for each of the representation values. For example, the distribution over representations of the data item can be determined from the product of the representation value distributions. In some cases, each representation value distribution can be defined by respective mean and variance values that are learned from the training examples. Thus, a representation of a data item can be obtained by sampling respective representation values from each of the representation value distributions for the data item.Attorney Docket No.45288-0436WO1

[0012] In some implementations, the target probability for the target subset given the set of data items of the training input can be determined by Monte Carlo sampling. For example, using the distributions over representations of the data items to determine the corresponding target probability for the target subset given the set of data items of the training input can comprise: for each data item, sampling one or more representations of the data item from the corresponding distribution over representations of the data item; and processing the sampled representations of the data items to determine the corresponding target probability for the target subset given the set of data items of the training input.

[0013] In some implementations, processing the sampled representations of the data items to determine the corresponding target probability for the target subset given the set of data items of the training input comprises: using the sampled representations of the data items to determine a corresponding score for each of a plurality of subsets of the set of data items; and using the scores to determine the corresponding target probability for the target subset given the set of data items of the training input. The score for each subset can, for example, be indicative of a similarity of the data items in the subset.

[0014] Using the sampled representations of the data items to determine the corresponding score for the subset can comprise: for each pair of data items in the subset of data items, determining a pairwise score using a sampled representation of one of the data items in the pair and a sampled representation of the other one of the data items in the pair; and using the pairwise scores to determine the score for the subset. For example, the pairwise score for each pair of data items can be determined using a similarity measure between the sampled representation of one of the data items in the pair and the sampled representation of the other one of the data items in the pair.

[0015] The similarity measure can, for example, comprise a dot product.

[0016] In the case where each subset includes only a pair of data items, the score for the subset may be the pairwise score for that pair. The target probability of the target subset may then be determined, for example, using a softmax function of the pairwise scores for each of the subsets. The target probability can, for example, be determined by averaging respective target probabilities determined using successive sampled representations of the data items.

[0017] In some implementations, processing the training examples to determine, for each training example, a corresponding target probability for the target subset given the set of dataAttorney Docket No.45288-0436WO1 items of the training input can comprise, determining the target probabilities using a “teacher” machine learning model that has itself been trained using the method, e.g., using training examples generated by humans performing a data item similarity task, e.g., an odd-one-out task. Knowledge from the teacher machine learning model can then be “distilled” into the (student) machine learning model using further training examples that have not been generated by humans, but rather by the teacher machine learning model. That is, a relatively small number of human similarity judgments can be “bootstrapped” by training the teacher (or “surrogate”) machine learning model to reproduce the human similarity judgements, and then using the trained teacher machine learning model to generate further training examples for training the “student” machine learning model to generate data item representations that are aligned with the human similarity judgements.

[0018] For example, the training examples can be generated by using the teacher machine learning model to generate respective representations for each of a plurality of data items and then applying a clustering algorithm (such as k nearest neighbours) to assign each of the data items to a respective one of a plurality of clusters. Each training example can then be generated by using the cluster assignments to select a respective set of the data items and a corresponding target subset. For example, the set of data items can comprise a plurality of data items assigned to a first cluster and one or more data items assigned to one or more other clusters different from the first cluster. For example, the set of data items can comprise two data items from a first cluster and another data item from a second cluster different from the first cluster. The target subset can then, for example, comprise the two data items from the first cluster, with the other data item from the second cluster being the “odd-one-out”.

[0019] In some implementations, using the machine learning model to determine a representation for each of the data items in the set of data items of the training input comprises: using another trained machine learning model to determine an initial representation of the data item; and applying a transformation to the initial representation to determine the representation of the data item. The other trained machine learning model can, for example, be a model that has been trained to perform one or more tasks that are different from the subset prediction task used in the training, such as an object classification or image recognition task or a next token prediction task, and so on. In some instances, the other machine learning model can have (many) more parameters than the machine learning model and / or be trained on training data setsAttorney Docket No.45288-0436WO1 comprising (many) more training examples. For example, the other machine learning model can be a “foundation” machine learning model such as Pali, see, e.g., Chen et al. arXiv:2310.09199. In some cases, the transformation can be used to align the representations of the data items generated by the other machine learning model with the training examples, without the other machine learning model needing to be re-trained or fine-tuned (although re-training or fine- tuning can still be beneficial in some cases).

[0020] For example, an affine transformation can be used to transform the initial representations of the data items to allow them to be aligned to the training examples, e.g., each initialrepresentation may be transformed using an affine transformation of the form ^^ = ^^^ + ^, inwhich W is a matrix and b is a vector, the values of which are learned training of themachine learning model.

[0021] In some implementations, the objective function used to compare the predicted probabilities with the corresponding target probabilities can comprise a soft cross-entropy function. For example, the objective function can comprise a term (for each of the training examples) in which the logarithm of the predicted probability for the target subset identified in the training output given the set of data items of the training input is multiplied by the corresponding target probability of the target subset given the set of data items of the training input. The objective function can, in some cases, additionally include another term (for each of the training examples) that depends on the target probability, e.g., a term comprising the logarithm of the target probability of the target subset given the set of data items of the training input.

[0022] In some implementations where a triplet odd-one-out task is used, the predictedprobability ^^^^^^, ^^, ^^^, ^^ for the target subset identified in the training output, denoted here by^^, ^^, given the set of data items of the training input can be determined using a softmaxfunction of the pairwise similarities, e.g.: ^^^^^, ^^|^^^, ^^, ^^^, ^^ ∶= exp ^^^,^^⁄ ^exp^^^, ! + exp^^^,"! + exp ^^ ,"^!

[0023] As noted above, in some implementations, an affine transformation defined by^^ = ^^^ + ^, can be used to transform representations of the data items to allow therepresentations to be aligned to the training examples. The values of the parameters of the affine transformation can learned during training of the machine learning model, in which case theAttorney Docket No.45288-0436WO1 objective function can, for example, comprise the soft alignment objective function with the pairwise similarity matrix being determined from the transformed representations, i.e., ^^,=^^^^ + ^^$^^^ + ^!. The objective function can also comprise a regularization term, e.g.,-2 %^^, ^^ + & '^ − )* + ⁄ , / 0'1in which & is asubtracts from W the identity matrix I which was been scaled by the mean of the diagonal elements of W. The hyperparameter & may be adjusted using a grid-search, for example. Optimising this objective function may be referred to as “uncertainty distillation”.

[0024] Thus, the values of the parameters adjusted during training of the machine learning model can comprise the values of the parameters defining the affine transform, e.g., the elements of W and b.

[0025] Adjusting the values of the parameters of the machine learning model to optimize the objective function can be performed in many different ways, e.g., depending on the machine learning model that is being trained. For example, a gradient-based optimization technique can be used in which gradients of the objective function with respect to the parameters of the machine learning model are back propagated through the machine learning model and the gradients used to update the values of the parameters of the machine learning model. A conventional optimizer, e.g., stochastic gradient descent, RMSprop, or Adam optimizer can be used.

[0026] In general, the data items processed by the system, e.g., the training data items, may comprise any type of data item including, for example, an audio data item, an image data item (which, as used herein, includes a video data item), a text data item, a graph data item, or a multimodal data item.

[0027] According to a second aspect of the present disclosure, there is provided a method performed by one or more computers for training a machine learning model to determine a representation of a data item. The method comprises: obtaining training examples that each comprise a respective set of data items; processing the training examples to determine, for each training example, a corresponding target probability for each of a plurality of subsets of the set of data items given the set of data items of the training example; and training the machine learningAttorney Docket No.45288-0436WO1 model to determine a representation of a data item. The training comprises: for each of the training examples: using the machine learning model to determine, in accordance with values of parameters of the machine learning model, a respective representation for each of the data items in the set of data items of the training example; and processing the representations of the data items to determine a respective predicted probability for each of the one or more subsets of the set of data items given the set of data items of the training example; and adjusting the values of the parameters of the machine learning model to optimize an objective function that compares the predicted probabilities with the corresponding target probabilities.

[0028] In some implementations, the respective probability of each of the plurality of subsets indicates a probability that the subset would be selected (e.g., by a human) from the set of data items in a data item similarity task. Each of the probabilities can be less than one to reflect uncertainty in performing the data item similarity task. In some implementations, the plurality of subsets comprises all possible subsets that are consistent with the data item similarity task. For example, for a triplet odd-one-out task, the plurality of subsets can comprise all possible pairs of the data items.

[0029] According to a third aspect of the present disclosure there is provided a method of generating training examples. The method comprises using the method of the first aspect or the second aspect to train a machine learning model (a “teacher” machine learning model) to determine a representation of a data item, e.g., to align the representation with human similarity judgements. The method further comprises: using the trained machine learning model to generate respective representations for each of a plurality of data items; and applying a clustering algorithm to assign each of the data items to a respective one of a plurality of clusters. The method further comprises generating a plurality of training examples, the generating comprising selecting a respective set of the data items and a corresponding target subset based on the cluster to which the data items have been assigned. For example, for each training example, the training input of the training example can comprise a set of data items that comprises a plurality of data items assigned to a first cluster and one or more data items assigned to one or more other clusters different from the first cluster, and the training output can identify a target subset comprising the plurality of data items assigned to a first cluster. For example, the set of data items can comprise two data items from a first cluster and another data item from a second cluster different from theAttorney Docket No.45288-0436WO1 first cluster. The target subset can then, for example, comprise the two data items from the first cluster, with the other data item from the second cluster being the “odd-one-out”.

[0030] The method can further comprise using the training examples and the teacher machine learning model to train another machine learning model (a “student” machine learning model) using the method of the first aspect or the second aspect, wherein the target probabilities are determined using the teacher machine learning model.

[0031] According to a further aspect of the present disclosure, there is provided 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 the first aspect, the second aspect, or the third aspect.

[0032] According to a further aspect of the present disclosure, there is provided 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 the first aspect, the second aspect, or the third aspect.

[0033] According to a further aspect of the present disclosure, there is provided a computer- readable medium storing data defining a representation of a data object obtained using a machine learning model trained or further trained by performing the operations of the respective method of of the first aspect or the second aspect.

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

[0035] Implementations of the described method and system may be used to generate representations of data items, such as image, audio, or text data items, which can be processed for a variety of “downstream” tasks. That is, the trained machine learning model, or part thereof, may be used for tasks other than the subset prediction task used to train (or re-train or fine-tune) the machine learning model, i.e., tasks other than selecting a subset (e.g., odd-one-out) of a set of data items.

[0036] During the training, the machine learning model learns to generate representations of data items that are adapted for generating predictions that match the target subsets of the training examples. For example, the machine learning model can learn to generate representations thatAttorney Docket No.45288-0436WO1 include features of the data items that are helpful in selecting subsets that accurately match the target subsets. The machine learning model can also learn to generate representations that take into account variations or uncertainties in the training examples, such when there are training examples that have the same or similar training inputs, but different respective training outputs. The representations generated by the trained machine learning model can therefore be said to be “aligned” with the training examples. Such alignment can improve performance when the representations are used for other, downstream machine learning tasks.

[0037] In some implementations, the training examples can be generated by human contributors. For example, the training examples can be generated by human contributors who have been asked to select, for each of a plurality of sets of data items, a respective subset of the set of data items according to one or more selection criteria, e.g., selecting a data item that is the odd-one- out from a set of three or more data items. Different human contributors may select different odd-ones-out from the same set of three or more data items depending on which features of the data items they consider most important. For example, for a set of data items comprising images of a house cat, a golden retriever, and a lion, one contributor might select the image of the golden retriever as the odd-one-out by considering whether the animals are feline, while another might select the image of the lion by considering whether the animals are domesticated. Thus, the machine learning model can be trained to include uncertainties such as these in the representations of data items that it learns to generate. Thus, the machine learning model can learn to generate representations that are aligned with human mental representations, i.e., representations that include features and characteristics of the data items that are the same or similar as the features and characteristics that humans tend to focus on when thinking about the data items. The representations of data items generated by the trained machine learning model can therefore be used in other, downstream tasks, to make predictions that are more aligned with those of humans. The representations and / or the predictions obtained using the representations may also be more readily interpretable because they are based on features and characteristics of the data items that humans consider relevant.

[0038] More generally, the machine learning model can learn to generate representations that are aligned with the training examples regardless of whether the training examples are human- generated. For example, in some cases, the training examples can be generated by a non- deterministic or stochastic process, such that different target subsets are selected for the same orAttorney Docket No.45288-0436WO1 similar sets of data items. The machine learning model can then learn to generate representations of the data items that include features or characteristics that account for variations caused by uncertainties in the underlying generation process.

[0039] By aligning the representations with the training examples for the subset prediction task, the amount of additional training (“fine tuning”) of the machine learning model that is needed for another downstream task may be reduced and / or the accuracy with which the downstream task can be performed can be improved.

[0040] In some implementations, the machine learning model may be configured to receive a model input that comprises representations of data objects that have been generated by another machine learning model, such as a “foundation” machine learning model with many more parameters than the machine learning model, which has been trained to perform a different task than the machine learning model. For example, the other machine learning model may be a language model that generates representations of text data items, a computer vision model that generates representations of image or video data items, a multimodal model (e.g., a visual language model) that generates representations of data items of more than one modality, and so on. The machine learning model may then learn to adapt the representations generated by the other machine learning model to align them with a particular training dataset, e.g., human- generated training examples. Thus, the machine learning model may be used to customize or improve the representations generated by the other machine learning model. During the training of the machine learning model, the values of the parameters of the other machine learning model may be frozen, so that fewer computational resources are required for the training. Alternatively, the values of the parameters of the other machine learning model can also be adjusted (“fine- tuned”), which may provide better alignment of the representations with the training examples.

[0041] As one example, the system and method described herein can improve the accuracy with which pre-trained vision models can perform an odd-one-out selection task, such as selecting an odd-one-out image from a triplet of images. In particular, the system and method can increase the accuracy with which the models can perform the task compared to other methods of training the models to perform the odd-one-out selection task that do not take into account uncertainties in the training examples, such as those which use a “hard” alignment loss.

[0042] For example, the CLIP ViT-L / 14 model (https: / / huggingface.co / openai / clip-vit-large- patch14) was found to have a baseline accuracy of 0.91 on a particular odd-one-out dataset,Attorney Docket No.45288-0436WO1 whereas adapting the model and training using the gLocal / global method described in Muttenthaler et al. “Improving neural network representations using human similarity judgements” in Thirty-Seventh Conference on Neural Information Processing Systems, 2023, achieved an accuracy of 0.93. By contrast, the present “uncertainty distillation” approach using the soft alignment objective function mentioned above achieved an accuracy of 0.95. BRIEF DESCRIPTION OF THE DRAWINGS

[0043] FIG. 1 is a schematic view of a training system for training a machine learning model to determine a representation of a data item.

[0044] FIG. 2 is a schematic view of the probability model of FIG. 1.

[0045] FIG. 3 is a schematic view of an alternative probability model.

[0046] FIG. 4 is a schematic diagram illustrating a multi-stage process for training a machine learning model.

[0047] FIG. 5 is a flow diagram of a process for training a machine learning model.

[0048] FIG. 6A are 2D latent space projections providing a visualization of the representations generated by a machine learning model (I) before and (II) after the machine learning model is trained using the process of FIG. 4

[0049] FIG. 6B is a bar chart showing respective performance of two machine learning models at reproducing human answers on two data item similarity tasks before and after training using the process of FIG. 4.

[0050] FIG. 6C is a bar chart showing respective performance of representations generated using two machine learning models on two downstream tasks before and after training using the process of FIG. 4.

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

[0052] FIG. 1 shows a training system 100, which 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.

[0053] The training system 100 is configured to use training data 102 to train a machine learning model 104 to generate representations of data items, such as image, audio, or text data items, thatAttorney Docket No.45288-0436WO1 are aligned with representations of the data items that were used to generate the training data. For example, the representations of the data items generated by the machine learning model 104 can be aligned with implicit representations of the data items used by humans who generated the training data. Such alignment of the representations of the data items generated by the machine learning model can improve performance when the representations are used to perform a downstream task, such as image or audio classification task.

[0054] The training data 102 comprises a plurality of training examples 106 that each include (i) a respective training input comprising a set of data items 108 and (ii) a respective training output comprising a target subset 110 of the data items. In some implementations, the target subset 110 is generated by performing a data item similarity task, that involves assessing similarities between data items. For example, the data item similarity task can be an odd-one-out task, in which case, the plurality of data items 108 can consist of three or more data items and the target subset of the data items can identify one of the data items that is less similar to the other data items in the plurality of data items than the other data items are to one another. In general, the target subset 110 is a proper subset of the set of data items 108, i.e., the target subset 110 comprises one or more, but not all, of the data items in the set of data items 108. As one example, the training data 102 can comprise training examples from the THINGS dataset described in Hebart et al., “Things-data, a multimodal collection of large-scale datasets for investigating object representations in human brain and behavior,” eLife, 12:e82580, 2023. The THINGS dataset comprises a plurality of object triplets and corresponding human odd-one-out responses, i.e., object pairs that were chosen by a human participant to have the highest similarity.

[0055] The machine learning model 104 is configured to determine, in accordance with values of parameters 112 of the machine learning model, a respective representation 114 for each of the data items in the set of data items 108 of the training input. Each representation of a data item can, for example, be a collection of numerical values, e.g., an ordered collection of numerical values, such as a vector, matrix, or other tensor.

[0056] The machine learning model 104 may have any appropriate architecture. For example, the machine learning model may be a neural network, e.g., with a feedforward architecture, and, as one particular example, may comprise a multi-layer perceptron (MLP), such as a two-layer MLP. The neural network may alternatively or additionally comprise convolutional neuralAttorney Docket No.45288-0436WO1 network layers, or a recurrent neural network subsystem, or a transformer neural network subsystem etc. For example, the machine learning model can comprise a Vision Transformer (ViT), e.g., as described in Dosovitskiy A., et al., “An image is worth 16x16 Words; Transformers For Image Recognition at Scale”, ICLR 2021. Examples of suitable architectures for a transformer neural network subsystem are described in Ashish Vaswani et al., “Attention is all you need”, Advances in Neural Information Processing Systems, pp. 5998–6008, 2017; arXiv:1810.04805 Devlin et al. (BERT); and arXiv:1901.02860 Dai et al. (Transformer-XL).

[0057] In some examples, the machine learning model comprises an encoder neural network and one or more output neural network layers that define an affine transformation that is applied to the representations generated by the encoder neural network. For example, each representation ^^may be transformed by the one or more output layers using an affine transformation of the form^^ = ^^^ + ^, in which W is a matrix and b is a vector, the values of which are learned duringthe training of the machine learning model. The parameters of the encoder neural network can be frozen during the training, for example.

[0058] In general, the data items processed by the machine learning model 104, e.g., the training data items, can comprise any type of data item including, for example, an audio data item, an image data item (which, as used herein, includes a video data item), a text data item, a graph data item, a multimodal data item, and so on.

[0059] The training system 100 further comprises a comparison function 116 configured to process the representations 114 of the data items generated by the machine learning model 104 to determine a predicted probability 118 for the target subset 110 identified in the training output given the set of data items 108 of the training input.

[0060] As one particular example, the comparison function 116 can be configured to process pairs of the representations 114 of the data items 108 to generate respective pair-wise similarity scores for every possible pair of data items in the set of data items. For example, the pairwise score for each pair of data items can be determined using a similarity measure, such as a scalar / dot product, between the representation of the first data item in the pair and the representation of the second data item in the pair. Determining the pairwise score can comprise summing products of the representation values of the sampled representation of one of the data items in the pair with the corresponding representation values of the sampled representation of the other one of the data items in the pair. For example, the pairwise score for a pair of data itemsAttorney Docket No.45288-0436WO1can be given by ^^, = ^$^ ^ , where ^^, is an element of a pairwise similarity matrixcorresponding to respective representations ^^and ^ . However, othersimilarity can also or alternatively be used.

[0061] In some implementations, the predicted probability 118, denoted here by^^^^^^, ^^, ^^^, ^^, for the target subset 110 identified in the training output, denoted here by^^, ^^, given the set of data items 108 ^^^, ^^, ^^^ of the training input can be determined using asoftmax function of the pairwise similarities, e.g.: ^^^^^, ^^|^^^, ^^, ^^^, ^^ ∶= exp ^^^,^^⁄ ^exp^^^, ! + exp^^^,"! + exp ^^ ,"^! .

[0062] The training system 100 further comprises a probability model 120 configured to process the training example 106 to generate a corresponding target probability 124 for the target subset 110 given the set of data items 108 of the training input. Example implementations of the probability model 120 are described below in connection with FIGS. 2 and 3.

[0063] The training system 100 further comprises an optimizer 126 configured to generate updated values 130 of the parameters of the machine learning model 130 by optimizing an objective function 122 that compares the predicted probabilities 118 for the training examples 106 with the corresponding target probabilities 124 generated using the probability model 122. In some implementations, the training system 100 can process batches of the training examples 106 when optimizing the objective function 122. The optimizer 126 can be any appropriate optimizer, e.g., stochastic gradient descent, RMSprop, or Adam optimizer can be used.

[0064] As one example, the objective function can comprise a term (for each of the training examples) in which the logarithm of the predicted probability for the target subset identified in the training output given the set of data items of the training input is multiplied by the corresponding target probability of the target subset given the set of data items of the training input. The objective function can, in some cases, additionally include another term (for each of the training examples) that depends on the target probability, e.g., a term comprising the logarithm of the target probability of the target subset given the set of data items of the training input.

[0065] In some implementations, the training system 100 is configured to determine respective probabilities of each possible subset of the set of data items. That is, the training system 100 can be configured to use the probability model 122 to determine respective target probabilities 124 ofAttorney Docket No.45288-0436WO1 each subset and to use the machine learning model 104 to determine respective predicted probabilities 118 of each subset. The objective function 128 can then comprise terms comparing each predicted probability 118 with the corresponding target probability 124. For example, in some implementations, the objective function can comprise a soft-alignment objective (loss) function of the form: D 1%456789:;<=^^^ ∶= * log ,∗^^^^^, ^^, ^^^^ − ,∗^^^^^, ^^, ^^^^ log ^^^^^^, ^^, ^^^, ^^in the target^^, ^^^, ^^are the corresponding predicted probabilities 118 for each of the subsets of the set of data items. The target probabilities 124 and the predicted probabilities 118 can each define a probability distribution over the subsets of the set of data items. The objective function 128 can then be regarded as measuring an extent to which the probability distribution predicted by the machine learning model 104 matches the target probability distribution generated by the probability model 122. The first term in the soft-alignment objective function can, for example, indicate the human uncertainty in identifying the correct subset, whilst the second term indicates the cross entropy between the predicted and target probability distributions.

[0066] The objective function can also comprise a regularization term. For the example of a learned affine transformation, the optimization performed by the training system can be expressed as, -2arg min^,^ ^^ + '^ − * +M 1 in which & is a hyperparameter thatfrom W the identity matrix I which was been scaled by the mean of the diagonal elements of W. The hyperparameter & may be adjusted using a grid-search, for example.

[0067] FIG. 2 shows an exemplary implementation of the probability model 122 of the training system 100. The probability model 122 comprises an uncertainty model 202 configured to process the data items 108 of the training example 106 to generate a respective distribution 204 over representations for each data item 108. For example, each representation of a data item can comprise a plurality of representation values (numeric values) and the distribution overAttorney Docket No.45288-0436WO1 representations of the data item can comprise a corresponding representation value distribution for each of the representation values. For example, the distribution 204 over representations of the data item can be determined from the product of the representation value distributions. In some cases, each representation value distribution can be defined by respective mean and variance values that are learned from the training examples. Thus, a representation of a data item can be obtained by sampling respective representation values from each of the representation value distributions for the data item 108.

[0068] As one example, the uncertainty model can comprise a Bayesian inference model (e.g., “VICE model”), such as is described in Muttenthaler et al. “VICE: Variational Interpretable Concept Embeddings” in Advances in Neural Information Processing Systems, 35:33661-33675 (2022). To capture the uncertainties in human triplet responses, the Bayesian inference model can learn a respective mean and variance for each dimension of the representation of each data object.For example, the Bayesian inference model can learn a matrix of mean values, N ∈ ℝQ×S and amatrix of variance values, T ∈ ℝQ×S, where d is the size of each representation and m is thenumber of data objects. As one particular example, the Bayesian inference model can comprise a VICE model trained on the THINGS triplet dataset, e.g., using the hyperparameters mentioned in Muttenthaler et al.

[0069] The probability model 122 comprises a sampling function 206 for obtaining a respective plurality of representations 208 for each of the data items 108 by sampling from the corresponding distributions 204 over representations. The plurality of representations 208 can be

[0070] The probability model 122 further comprises a comparison function 210, which can be the same as, or different from, the comparison function 116 of the training system 100. The probability model 122 is configured to process the pluralities of representations 208 of the data objects to determine the target probability 124. For example, the target probability 124,,^^^^, ^^|^^^, ^^, ^^^^^, can be determined as a Monte Carlo average over the target probabilitiesover the representations of the data items, e.g., 1 Z !,where X^Y^,the number oftriplets of representations, and , ^V^^ ^^^, ^^|^^^, ^^, ^^^^^ is a target probability for the r-th triplet,Attorney Docket No.45288-0436WO1 which can be determined from the sampled representations, in the same way as for the predicted probability, i.e., by computing a pair-wise similarity matrix and then applying a softmax function, as described above.

[0071] FIG. 3 shows another probability model 322 that can be used to generate the target probabilities 124 for training the machine learning model 104 using the training system 100 of FIG. 1. The probability model 322 comprises a teacher machine learning model 324 configured to generate respective representations 308 of the data items 108 in each training example 106. The teacher machine learning model 324 can have been trained to generate representations that are aligned with representations that were used to generate the training data 102. For example, the teacher machine learning model 324 can have been trained to generate such representations using the training system 100 and the probability model 122 of FIGS. 1 and 2. The probability model further comprises a comparison function 310 for converting representations 308 of the data items 108 generated by the teacher machine learning model into target probabilities 124. The comparison function can, for example, be the same as, or different from the, the comparison functions 116, 210 described with respect to FIGS. 1 and 2.

[0072] Use of the probability model 322 comprising the teacher machine learning model 324 to generate the target probabilities 124 may be referred to as “distillation” or “knowledge transfer”, as knowledge from the teacher machine learning model 324 can be transferred to the machine learning model 104 via the target probabilities 124. The machine learning model 104 of the training system 100 may then be referred to as a “student” machine learning model.

[0073] In some implementations, the teacher machine learning model 124 can be used to generate training data 102 for training the student machine learning model 104. For example, the teacher machine learning model 124 can be used to generate respective representations of each data item 108 in a collection of data items. The collection of data items may comprise data items that were not used in the training of the teacher machine learning model 124. For example, the collection of data items can comprise (many) more data items than were used in the training of the teacher machine learning model 124. As one particular example, the teacher machine learning model 124 can be trained using training data 102 from the THINGS dataset, whilst the collection of data items for generating other training data 102 for training the student machine learning model can comprises data items from the ImageNet dataset described by RussakovskyAttorney Docket No.45288-0436WO1 et al., “Imagenet large scale visual recognition challenge”, International Journal of Computer Vision, 115:211–252, 2015.

[0074] To transfer the structure of the pairwise similarity matrix of a teacher machine learning model, into a different student machine learning model [\] (e.g., a Transformer-based visual foundation model), the training system 100 can, for be configured to use Kullback-Leibler divergence based objective function %9:;<=^7^^_, ^`^!^

[0075] In matrices of therepresentations generated by the teacher and student machine learning models, T is a softmax function that transforms the similarities into probabilities with temperature values, d_and d`for the teacher and student machine learning models respectively (optimum values for which can be determined, e.g., using a hyperparameter grid search) , and B is the total number of samples used. In some implementations, an overall objective function of the following form can be used: %^[ ]! + &gh∗ ` 29:;<=^7 \ − h g2Here, h∗are the parameters oflearning model before fine-tuning by alignment and h`are the parameters of the fine-tuned student machine learning model. The right-hand side of the equation helps to preserve the fine-grained structure of the pretrained representation space as a function of & which determines the strength of the regularization.

[0076] FIG. 4 is a schematic view of a method 400 for training a student machine learning model 104C using the training system 100. The method 400 comprises a first stage 400A in which the training system 400 is used to train a pre-trained model 104A to generate aligned representations of data objects from a first set of training data 102A, in this example, images from the THINGS database. In this first stage, the training system 100 uses the probability model 122 described above in connection with FIG. 2 to train the model 104A, e.g., using a VICE model that has been trained using the first set of training data 102A.

[0077] In a second stage 400B of the method 400, the trained model 104A is used as a “linearly aligned” teacher machine learning model 324 to generate representations of the data items from aAttorney Docket No.45288-0436WO1 collection of data items 402, in this example, images from the ImageNet database. A clustering algorithm, such as k-means clustering, is then applied to the representations to assign each representation to a respective “superordinate” cluster. An optimal number of clusters k can be determined using the Elbow criterion, for example. Data items are then sampled from the clusters are then used to generate a plurality of training examples (“cluster-border triplets”) 106 that each comprise two data items 108 from a first of the clusters and a third data item 108 from a second of the clusters. The plurality of training examples is then used as a second set of training data 102C (“AligNet”).

[0078] In a third stage 400C of the method 400, the training system 100 is used to train a student machine learning model 104C on the second set of training data 102C using the teacher machine learning model 324 to determine the target probabilities 124.

[0079] FIG. 5 shows a process 500 which can be performed by one or more computers. For example, the process 500 can be implemented by the training system 100 shown in FIG. 1.

[0080] The process 500 comprises obtaining (step 502) training examples, e.g., training examples that each comprise (i) a training input comprising a set of data items and (ii) a training output identifying a target subset of the set of data items.

[0081] The process 500 further comprises processing (step 504) the training examples to determine, for each training example, a corresponding target probability for the target subset of the set of data items given the set of data items of the training input, and training (step 506) the machine learning model to determine a representation of a data item.

[0082] The training comprises, for each of the training examples, using (step 508) the machine learning model to determine, in accordance with values of parameters of the machine learning model, a representation for each of the data items in the set of data items of the training input, and processing (step 510) the representations of the data items to determine a predicted probability for the target subset of the set of data items given the set of data items of the training input.

[0083] The training further comprises adjusting (step 512) the values of the parameters of the machine learning model to optimize an objective function that compares the predicted probabilities with the corresponding target probabilities.

[0084] FIG. 6A shows two 2D latent space projections of representations of images generated by a vision transformer model, here a ViT-B classifier, (I) before and (II) after the model is trainedAttorney Docket No.45288-0436WO1 (fine-tuned) using the process of FIG. 4. Before the training, the representations of the images belonging to different categories (“animal”, “clothes”, “food”, etc,) overlap one another in the latent space, whereas after the training, the representations of the image belonging to different categories are clustered in different regions of the latent space. Respective projections of each image in a triplet of images comprising a spider, a beaver, and a landscape scene, are also shown in FIG. 6A. Before the training, the 2D distances between the projections are relatively small and indicate that the image of the spider is the odd-one-out. After training, the 2D distances are much larger and indicate that the landscape scene should be considered the odd-one-out, consistent with typical human classifications of the triplet of images.

[0085] FIG. 6B shows alignment scores for two different vision transformer machine learning models before and after training using the process 400 of FIG. 4. After fine-tuning on the AligNet dataset 102C, the models achieved odd-one-out accuracies of 60.45% (ViT-B) − 62.54% (SigLIP ViT-B) of which 62.54% is believed to be the best alignment score of a neural network model on the THINGS dataset 102A to date. In contrast, the base model representations achieved odd-one-out accuracies of 34.87% (ViT-B) − 51.21% (DINOv2 ViT-B). The odd-one-out accuracy of ViT-B increased substantially from 34.87% to 60.45%. This is a relative increase in performance of 73.35%. All performance increases are statistically significant (α = 0.05).

[0086] FIG. 6C shows accuracy scores for the trained vision transformer machine learning models when performing (i) a few-shot learning task in which the objective is to classify images given only a single labelled example per class or (ii) a distribution shift task in which the training images are sampled from a different subpopulation from the test images.

[0087] The alignment technique described in this specification can be used in many different domains, including the visual domain, but also natural language processing. For instance, in natural language processing, machine learning models can be trained with objectives that focus on distinguishing between close matches (e.g., prediction objectives that primarily distinguish words that occur in the current context from those that would occur in slightly different contexts, rather than considering their relations to less probable concepts). Applying the present alignment techniques to natural language processing can therefore help to better capture the global structure of semantic and syntactic relationships among language inputs that these objectives miss.Attorney Docket No.45288-0436WO1

[0088] When used to perform a downstream task, the values of the parameters of the trained machine learning model may be frozen, or subject to further training on the downstream task.

[0089] The task (i.e., the subsequent task) that the trained neural network, or part thereof, is used to perform may generally correspond to a type of the training data item. For example where the training data item comprises an audio data item, an image data item, a multimodal data item, a text data item, or a graph data item, the trained machine learning model, or part thereof, may be used, correspondingly, to process input data comprising audio data, image data, multimodal data, text data, or graph data respectively to perform an audio signal processing task, an image processing task, a multimodal processing task, a text processing task, or a graph processing task.

[0090] As one example, the training data item, and input data, may comprise audio data representing values of a digitized audio waveform, e.g., a time sequence of waveform- representing elements. Such a representation may comprise, e.g., samples representing digitized amplitude values of the waveform or a time-frequency domain representation of the waveform such as a STFT (Short-Term Fourier Transform) or MFCC (Mel-Frequency Cepstral Coefficient) representation. The audio waveform may comprise, e.g., a speech waveform or a waveform of a sound, e.g., a captured sound. As some examples of transformations that may be used, transformed views of the training data item may be obtained by transformations including: time or pitch warps; random crops in the time or frequency domain, e.g., selections of portions of the audio data item with random start and end times or with randomly selected upper and lower frequencies; modifications to the amplitude of a data item, e.g., by randomly increasing or diminishing the amplitude of the audio; or modifications to the frequency characteristics of the audio, e.g., by randomly filtering the audio. Objects in the audio waveform may comprise, e.g., speech elements such as words, syllables, or phonemes; or events or other distinguishable audio objects in the sound.

[0091] The audio signal processing task may comprise, e.g.: processing audio data representing speech to provide output data that detects words or phonemes in the speech or categorizes words or phonemes in the speech into one or more of a plurality of categories; or processing audio data representing a sound to provide output data, e.g., likelihood data, that detects presence of a particular sound or audio object or event in the sound, e.g., in a hotword detection or identification task; or processing audio data representing a sound to provide output data that categorizes a content of the sound into one or more of a plurality of categories (i.e., classifying aAttorney Docket No.45288-0436WO1 sound). In some further examples the audio signal processing task may comprise, e.g.: an identification or classification task such as a speech or sound recognition task, e.g., a hotword detection or identification task, a speaker or natural language classification task, or an audio tagging task, in which case the output data may comprise a category score or tag for the audio or for a segment of the audio; or a similarity determination task, e.g., an audio copy detection or search task, in which case the output data may comprise a similarity score.

[0092] In some implementations the training data item, and input data, may comprise sensor data representing values of a digitized sensor waveform, i.e., a sensor other than an audio sensor may be used to obtain the digitized waveform. The digitized sensor waveform may be treated similarly to a digitized audio waveform. The sensor data may be generated by sensors configured to monitor the real-world state, condition or environment of a physical system, e.g., of a mechanical or electronic physical system or machine, e.g., sensing force, pressure, movement, temperature, or vibration. The objects may comprise events or other distinguishable objects in the sensor data, or conditions of the physical system. The signal processing task may be to process the input data to provide output data that identifies the presence of one or more of the events, objects, conditions or environments.

[0093] As another example, the training data item, and input data, may comprise image data representing a still or moving image, i.e., an image or video, e.g., an image or video that has been captured using a camera. Elements of the image data may comprise monochrome or color pixels of the image or video. As defined herein an “image” includes a point cloud, e.g., from a LIDAR system, and a “pixel” includes a point of the point cloud. Similarly “video” includes a time sequence of point clouds. Objects in the image or video may comprise objects, e.g., physical objects, represented by the image or video.

[0094] The image processing task may comprise, e.g.: processing the image data to provide output data that identifies the location of one or more specified or unspecified objects in the image or video, e.g., output data that defines one or more object bounding shapes or boxes; or processing the image data to provide output data that segments pixels of the image or video into regions that represent one or more objects in the image or video signal; or processing the image data to provide output data that categorizes a content of the image or video into one or more of a plurality of categories; or processing the image data to provide output data that predicts depth values for pixels of the image or video. A task that segments the pixels may be, e.g., a semanticAttorney Docket No.45288-0436WO1 segmentation task that associates each pixel with a category representing a class of objects, or an instance segmentation task that associates each pixel with a category representing an instance of an object, i.e., to distinguish between different instances of the same category of object.

[0095] Where the image data comprises pixels of a video the image processing task may comprise, e.g.: processing the image data to provide output data that identifies the location of one or more actions represented in the video; or processing the image data to provide output data that categorizes one or more actions, e.g., gestures, represented in the video into one or more of a plurality of categories.

[0096] In general the image processing task may include any sort of image processing or computer vision task such as an image classification or scene recognition task, an image segmentation task, e.g., a semantic or instance segmentation task, an object localization or detection task, or a depth estimation task. When performing such a task the input data may be derived from pixels of the image.

[0097] For an image classification or scene recognition task the output may comprise a classification output providing a score for each of a plurality of image or scene categories, e.g., representing an estimated likelihood that the image data or an object represented in the image data, or that an action within image data representing a video, belongs to a category of a set of categories.

[0098] For an image segmentation task the output may comprise, for each pixel, an assigned segmentation category or a probability that the pixel belongs to a segmentation category, e.g., to an object or action represented in the image or video. For an object localization or detection task the output may comprise data defining coordinates of a bounding box or region for one or more objects represented in the image. Such a bounding box or region may be defined in two, three or more dimensions (time counting as a dimension). For a depth estimation task the output may comprise, for each pixel, an estimated depth value. The output may define a continuous value or it may define a probability distribution over discrete depth value buckets, such that the output pixels define a (spatial 3D) depth map for the image. Such tasks may also contribute to higher level tasks, e.g., to object tracking across video frames; or to gesture recognition, i.e., recognition of gestures that are performed by entities depicted in a video. As another example, the image processing task may include an image keypoint detection task in which the output comprises the coordinates of one or more image keypoints, such as landmarks of an object represented in theAttorney Docket No.45288-0436WO1 image, e.g., a human pose estimation task in which the keypoints may define the positions of body joints. A further example is an image similarity determination task, in which the output may comprise a value representing a similarity between two images, e.g., as part of an image search task.

[0099] In some applications, the image data item may be a medical image, such as an X-ray image, CAT scan or MRI image. For example, the machine learning model may be trained to output segmentation data indicating the location(s) of cancerous matter in the medical image

[0100] As another example the training data item, and input data, may comprise text data; elements of the text data may comprise, e.g., sentences, words, or parts of words, e.g., wordpieces. The text processing task may comprise, e.g.: a part-of-speech tagging task, in which case the output data may comprise, e.g., a category score or tag for the text or for a segment of the text; or a dependency parsing task, in which case the output data may comprise data representing a dependency parse of the text; or a text segmentation task, in which case the output data may comprise data that associates elements of the text with one or more of a plurality of categories for the text. Other example tasks include an identification or classification task, or a similarity determination task, e.g., to generate a category score, a similarity score, or a tag as described above; or a machine translation task.

[0101] As another example the training data item, and input data, may comprise multimodal data. In general such multimodal data is a combination of two or more different types of data, where the different types of data represent the same or overlapping objects using the different modalities (types). As one example the multimodal 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 multimodal 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. Elements of the multimodal data may correspond to elements of the data types making up the combination. Optionally, but not necessarily, when processing multimodal data the data may be mapped into a common embedding space.

[0102] In general the multimodal processing task may correspond to any of the tasks previously described for any of the types of data making up the multimodal combination. For example, an accuracy of the previously described tasks may be increased when the task is applied toAttorney Docket No.45288-0436WO1 multimodal 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.

[0103] As one particular example, where the multimodal data comprises audio-visual data the multimodal processing task may comprise: processing the combination, i.e., the image / video and audio, to provide output data that detects presence of a particular multimodal object or event in the combination (e.g., to identify a phoneme or viseme when lip reading); or processing the combination to provide output data that categorizes the combination into one or more of a plurality of categories, e.g., by defining a score for each category of a plurality of possible categories for the combination. As another particular example, where the multimodal data comprises a combination of text data and image or video or audio data the multimodal processing task may comprise processing the combination to provide output data that defines whether the image or video or audio waveform is described by the text, e.g., by a particular caption, e.g., by defining a score for the text or caption.

[0104] Some example multimodal machine learning models with which the techniques described herein may be used include: Flamingo (Alayrac et al. arXiv:2204.14198); ALIGN (Jia et al., arXiv:2102.05918); PaLI (Chen et al. arXiv:2209.06794); and PaLI-X (Chen et al. arXiv:2305.18565).

[0105] As another example the training data item, and input data, may comprise graph data; in such implementations the neural networks described herein may comprise graph neural networks. In general the graph data may define a graph structure having a set of nodes with associated node feature vectors connected by edges which may have associated edge feature vectors. A graph may, but need not be, defined by an adjacency matrix, e.g., where N is thenumber of nodes, an i × i matrix defining which nodes are connected by edges. Elements ofthe graph data may comprise, e.g., nodes or edges of a graph represented by the graph data.

[0106] A graph may represent a real-world physical system; merely as some examples, a mechanical structure in which bodies are connected by joints, or a structure of a molecule such as a drug molecule. The objects may comprise, e.g., physical bodies or parts of a molecule, e.g., chemical moieties. The graph processing task may comprise e.g.: characterizing a physical entity represented by the graph to provide output data that defines a predicted stability of the physical structure or molecule, or the binding affinity of a molecule represented by the graph with anotherAttorney Docket No.45288-0436WO1 molecule, e.g., to identify a drug candidate (which may then be evaluated by synthesizing the molecule and, e.g., testing the molecule in vitro or in vivo). The predicted stability of the physical structure may be used, e.g., to design or evaluate a structure; the result may then be used to construct a structure to the design. As another example the graph may be a scene graph that represents a scene; the scene graph may have been generated from a captured real-world image. The graph processing task may then comprise generating output data that identifies or classifies the scene or one or more objects within the scene, e.g., to facilitate object / scene editing or information extraction for scene interpretation.

[0107] In some implementations, the downstream task may comprise any one or more of:

[0108] (i) Data item classification, i.e., assigning an input data item (e.g., image) to one (or more) of set of classes based on content of the data item. Note that “content” here means content defined by pixelated intensity values in the case that the data item is an image, or defined by amplitude values in the case that the media item is a sound.

[0109] (ii) Captioning (generating a caption, such as a sequence of tokens selected from a vocabulary) describing the content of the data item.

[0110] (iii) Question answering, e.g., visual question answering (VQA), of generating an answer which is an appropriate response to a question about the data item;

[0111] (iv) A segmentation task of identifying a (proper) subset of the data item having specific properties, e.g., for a data item which is an image, “a penguin in the foreground”, “cancerous material”; or for a data item which is a sound signal, “a child crying after a dog has barked”).

[0112] (v) An agent control task of generating control data for an agent in an environment (e.g., an electro-mechanical agent such as a robot) based on a data item describing the environment (e.g., received sensor data, such as an image of an environment captured by a camera).

[0113] Some implementations may be used to perform neural machine translation. Thus in some implementations the data item comprises tokens that represent words, wordpieces, or characters in a first natural language and the downstream task is to generate output tokens that represent words, wordpieces or characters in a second, different natural language. That is, the data item may represent input text in the first language and the output sequence may represent a translation of the input text into the second language.

[0114] Some implementations may be used for automatic code generation. For example, the data item can comprise input tokens that represent words, wordpieces or characters in a first naturalAttorney Docket No.45288-0436WO1 language and the downstream task can be to generate output tokens that represent instructions in a computer programming or markup language, or instructions for controlling an application program to perform a task, e.g., build a data item such as an image or web page.

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

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

[0117] 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. Thus aAttorney Docket No.45288-0436WO1 system, artificial neural network, or trained artificial neural network as described herein, can be implemented in hardware using electronic circuitry, e.g., in a physical box. Similarly computer code as described herein can be code to emulate such hardware or code for a hardware description language.

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

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

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

[0121] 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 elements of a computer are a central processing unit for performingAttorney Docket No.45288-0436WO1 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.

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

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

[0124] 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.Attorney Docket No.45288-0436WO1

[0125] Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework, a Microsoft Cognitive Toolkit framework, an Apache Singa framework, or an Apache MXNet framework.

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

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

[0128] 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.Attorney Docket No.45288-0436WO1

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

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

[0131] What is claimed is:

Claims

Attorney Docket No.45288-0436WO1 CLAIMS 1. A method performed by one or more computers for training a machine learning model to determine a representation of a data item, the method comprising: obtaining training examples that each comprise (i) a training input comprising a set of data items and (ii) a training output identifying a target subset of the set of data items; processing the training examples to determine, for each training example, a corresponding target probability for the target subset given the set of data items of the training input; and training the machine learning model to determine a representation of a data item, the training comprising: for each of the training examples: using the machine learning model to determine, in accordance with values of parameters of the machine learning model, a representation for each of the data items in the set of data items of the training input; and processing the representations of the data items to determine a predicted probability for the target subset identified in the training output given the set of data items of the training input; and adjusting the values of the parameters of the machine learning model to optimize an objective function that compares the predicted probabilities with the corresponding target probabilities.

2. The method of claim 1, wherein processing the training examples to determine the corresponding target probability for the target subset given the set of data items of the training input comprises: processing the training examples to determine, for each data item, a corresponding distribution over representations of the data item; and using the distributions over representations of the data items to determine the corresponding target probability for the target subset given the set of data items of the training input.Attorney Docket No.45288-0436WO1 3. The method of claim 2, wherein using the distributions over representations of the data items to determine the corresponding target probability for the target subset given the set of data items of the training input comprises: for each data item, sampling one or more representations of the data item from the corresponding distribution over representations of the data item; and processing the sampled representations of the data items to determine the corresponding target probability for the target subset given the set of data items of the training input.

4. The method of claim 3, wherein processing the sampled representations of the data items to determine the corresponding target probability for the target subset given the set of data items of the training input comprises: using the sampled representations of the data items to determine a corresponding score for each of a plurality of subsets of the set of data items; and using the scores to determine the corresponding target probability for the target subset given the set of data items of the training input.

5. The method of claim 4, wherein the score for each subset is indicative of a similarity of the data items in the subset.

6. The method of claim 4 or 5, wherein using the sampled representations of the data items to determine the corresponding score for the subset comprises: for each pair of data items in the subset of data items, determining a pairwise score using a sampled representation of one of the data items in the pair and a sampled representation of the other one of the data items in the pair; and using the pairwise scores to determine the score for the subset.

7. The method of claim 6, wherein the pairwise score for each pair of data items is determined using a similarity measure between the sampled representation of one of the data items in the pair and the sampled representation of the other one of the data items in the pair.

8. The method of claim 7, wherein the similarity measure comprises a dot product.Attorney Docket No.45288-0436WO1 9. The method of any preceding claim, wherein using the machine learning model to determine a representation for each of the data items in the set of data items of the training input comprises: using another trained machine learning model to determine an initial representation of the data item; and applying a transformation to the initial representation to determine the representation of the data item.

10. The method of claim 9, wherein the parameters of the machine learning model comprise parameters defining the transformation and wherein adjusting the values of the parameters of the machine learning model to optimize the objective function comprises adjusting the values of the parameters defining the transformation.

11. The method of claim 9 or 10, wherein the transformation is an affine transformation.

12. The method of any preceding claim, wherein the number of data items in each target subset is one less than the number of data items in the corresponding set of data items.

13. The method of any preceding claim, wherein each set of data items is a triplet of data items.

14. The method of any preceding claim, wherein the objective function that compares the predicted probabilities with the corresponding target probabilities is a soft cross-entropy function.

15. The method of claim 14, wherein the training examples comprise one or more groups of training examples, each group comprising training examples having identical or similar respective sets of data items, but different respective target subsets.

16. The method of any preceding claim, wherein each target subset is selected by one or more humans.Attorney Docket No.45288-0436WO1 17. The method of any preceding claim, wherein each data item comprises one or more of: an audio data item; an image data item, a multimodal data item; a text data item; and a graph data item, 18. The method of claim 17, wherein the method further comprises, after the training of the machine learning model: using the machine learning model to process input data comprising one or more data items to perform an audio signal processing task, an image processing task, a multimodal processing task, a text processing task, or a graph processing task.

19. A method performed by one or more computers for performing a machine learning task, the method comprising: providing input data comprising one or more data items to a machine learning model, the machine learning model having been trained or further trained by performing the respective operations of the method of any one of claims 1-18; determining a representation of each of the one or more data items using the machine learning model; and processing the representation of each of the one or more data items to perform the machine learning task.

20. The method of claim 18 or 19, wherein: (i) the audio data comprises data representing values of an audio waveform, and wherein the audio signal processing task comprises: processing audio data representing speech to provide output data that detects words or phonemes in the speech or categorizes words, syllables or phonemes in the speech into one or more of a plurality of categories, or processing audio data representing a sound to provide output data that detects presence of a particular sound or event in the sound, or processing audio data representing a sound to provide output data that categorizes a content of the sound into one or more of a plurality of categories; or (ii) the image data comprises pixels of an image or of video, and wherein the image processing task comprises: processing the image data to provide output data that identifies the location of one or more objects in the image or video, or processing the image data to provide output data that segments pixels of the image or video into regions that represent one or moreAttorney Docket No.45288-0436WO1 objects in the image or video signal, or processing the image data to provide output data that categorizes a content of the image or video into one or more of a plurality of categories, or processing the image data to provide output data that predicts depth values for pixels of the image or video; or (iii) the image data comprises pixels of a video and the image processing task comprises: processing the image data to provide output data that identifies the location of one or more actions represented in the video, or processing the image data to provide output data that categorizes one or more actions represented in the video into one or more of a plurality of categories; or (iv) the multimodal data comprises audio-visual data comprising a combination of pixels of an image or of video and audio data representing values of an audio waveform, and wherein the multimodal processing task comprises: processing the combination to provide output data that detects presence of a particular event in the combination, or processing the combination to provide output data that categorizes the combination into one or more of a plurality of categories; or (v) the multimodal data comprises a combination of text data representing text in a natural language and pixels of an image or of video or audio data representing values of an audio waveform, and wherein the multimodal processing task comprises: processing the combination to provide output data that defines whether the image or video or audio waveform is described by the text; or (vi) the text data comprises a input sequence of text in a source language and wherein the text processing task comprises: processing the input sequence to provide an output sequence of text in a target language that is a translation of the input sequence.

21. A method performed by one or more computers for processing a data item, the method comprising: providing an input data item to a machine learning model, the machine learning model having been trained or further trained by performing the respective operations of the method of any one of claims 1-20; determining a representation of the input data item using the machine learning model;Attorney Docket No.45288-0436WO1 and processing the representation of the input data item.

22. 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 the operations of the respective method of any one of claims 1-21.

23. 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 the operations of the respective method of any one of claims 1-21.

24. A computer-readable medium storing data defining a representation of a data object obtained using a machine learning model trained or further trained by performing the operations of the respective method of any one of claims 1-21.

Images