EP4699016A1 - Retrieval-enhanced vision-text models - Google Patents

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for performing a classification task on a data item. In particular, a system classifies an input data item using key and value embeddings of memory data items.

Description

RETRIEVAL-ENHANCED VISION-TEXT MODELS

CROSS-REFERENCE TO RELATED APPLICATIONS

This application claims the benefit of priority to U.S. Provisional Patent Application No. 63/469,358, filed on May 26, 2023, the contents of which are hereby incorporated by reference.

BACKGROUND

This specification relates to processing images and text using machine learning models.

As one example, 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 another layer in the network, e.g., 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 weights.

SUMMARY

This specification describes a system implemented as computer programs on one or more computers that performs a multi-modal machine learning task using text embeddings and image embeddings.

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

Contrastive image-text models form the building blocks of many state-of-the-art systems for performing various image processing tasks. While these models excel at recognizing common generic concepts, they still struggle on fine-grained entities which are rare, or even absent from the pre-training dataset. Hence, a key ingredient to their success has been the use of large-scale curated pre-training data aiming at expanding the set of concepts that they can memorize during the pre-training stage.

This specification, on the other hand, describes an alternative to encoding finegrained knowledge directly into the model’s parameters: the techniques instead retrieve this knowledge from an external memory. Specifically, the described techniques equip a vision-text model with the ability to refine initial embeddings of images and text with cross-modal retrieved information from a memory at inference time, which greatly improves the performance of the model, e.g., when zero-shot predictions are required.

Additionally, this increase in accuracy can be accomplished with minimal increased computational overhead relative to the existing model. For example, the performance gains can be achieved by including a kNN search and a light-weight fusion model on top of a frozen text-vision model.

As a particular example, by performing uni-modal search followed by cross- modal fusion, the system can effectively incorporate context from memory text in a memory when performing an image classification task, e.g., a zero-shot image classification task.

As another particular example, by performing uni-modal search followed by cross-modal fusion, the system can effectively incorporate context from memory images in a memory when performing an image retrieval task, e.g., a task that requires retrieving or identifying images that are described by a given input text.

Moreover, because the system only retrieves a very small number, e.g., between 1-20 data items, of context data items from the memory even when the memory is very large, e.g., includes more than 500 million pairs, the system can effectively leverage the very⁷ large memory' with only minimal additional computational overhead when processing any given query.

Additionally, after training, the system can use the described techniques to effectively perform tasks using memory data items that are entirely or partially different from the data items on which the neural networks in the system were trained. Thus, the system can (i) incorporate additional context by adding new pairs to the memory after training is completed or (ii) use an entirely different data set as context after training without requiring any computationally expensive retraining of any of the neural networks in the system.

The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims.

BRIEF DESCRIPTION OF THE DRAWINGS

FIG. 1 shows an example neural network system. FIG. 2 is a flow diagram of an example process for generating a prediction for a query image.

FIG. 3 shows an example of the operation of the system when receiving a query image

FIG. 4 is a flow diagram of an example process for generating a prediction for a query text.

FIG. 5 shows an example of the operation of the system when receiving a query image.

FIG. 6 shows an example of training the query' fusion neural networks.

FIG. 7 shows an example of the performance of the described techniques on an image classification task.

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

DETAILED DESCRIPTION

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

The system 100 performs a multi-modal machine learning task in response to received inputs 102 using text embeddings and image embeddings. Generally, the task is referred to as a multi-modal task because it requires processing both images and text.

An embedding, as used in this specification, is an ordered collection of numerical values having a fixed dimensionality. The space of embeddings having the fixed dimensionality is referred to as an embedding space. For example, each embedding can be a vector having a fixed number of entries, with the space of vectors that have the same, fixed number of entries being the embedding space.

In general, the multi-modal task can be any task that requires processing both images and text.

Examples of such tasks include image classification, e.g., zero-shot image classification, where the input is an image and the output (the “prediction”) identifies one or more text labels for the input image that identify objects that are depicted in the image; image-to-text retrieval, where the input is an input image and the output identifies one or more text segments in a repository that are most similar to, i.e., best describe, the input image; text-to-image retrieval, where the input is text and the output identifies one or more images in a repository that are most similar to. i.e., best described by. the input text; and visual question-answering, where the input is an image and a text question that relates to the contents of the image and the output is a text answer to the question.

More specifically, the system 100 performs a multi-modal task or multiple different multi-modal tasks by augmenting an image embedding neural network 130. a text embedding neural network 140. or both with a memory 150.

The memory 1 0 stores, for each of a plurality of memory image - memon text pairs, (i) a respective memon' image embedding 160 of the memon' image in the memon' image - memory text pair in an embedding space and (ii) a respective memory text embedding 170 of the memory text in the memory image - memon text pair in the embedding space.

The image embedding neural network 130 can be any appropriate neural network that can map images to an embedding. For example, the image embedding neural network 130 can be a convolutional neural network or a vision Transformer neural network.

The text embedding neural network 140 can be any appropriate neural network that can map text to an embedding. For example, the text embedding neural network 140 can be a Transformer neural network or a recurrent neural network (RNN).

As a particular example, the system or another training system can have jointly trained the image embedding neural network 130 and the text embedding neural network 140 through unsupervised learning, e g., through contrastive learning or another representation learning task that aligns embeddings between images and text.

To generate the embeddings of the memory' images in the memory' 150, the system 100 can process each memory image using the (trained) image embedding neural network 130. That is, the memory⁷ image embeddings 160 of the memory images have been generated using the image embedding neural network 130.

To generate the embeddings of the memory' texts in the memory 150, the system can process each memory text using the (trained) text embedding neural network 140. That is. the memory text embeddings 170 of the memory texts have been generated using the text embedding neural network 140.

More specifically, when performing any given task, the system 100 uses the memory' 150 to augment a query embedding 104 of an input 102 with context embeddings 106 retrieved from the memory 150 in order to generate a fused embedding 108. For example, for some tasks, the system 100 can receive a query text as part of the input 102 and process the query text 102 using the text embedding neural network 140 to generate a query text embedding of the query text 102 in the embedding space.

As another example, for some tasks, the system 100 receives a query image as part of the input 102 and processes the query image 102 using the image embedding neural network 130 to generate a query image embedding of the query image 102 in the embedding space.

The system 100 can then perform a uni-modal search through the memory 150 to identify context data items that are same modality as the query input, but then perform cross-modal retrieval to retrieve, as the context embeddings 106, the memory embeddings of the other modality that are in the same memory pairs as the context data items. That is, when the query is a query image, the system 100 retrieves memory text embeddings of memory texts. When the query is text, the system 100 retrieves memory image embeddings of memory' images.

The system 100 can then perform cross-modal fusion by using the queryembedding 104 and the context embeddings 106, i.e., the retrieved memory- embeddings of the other modality ,to generate the fused embedding 108.

Thus, performing cross-modal fusion allows the system 100 to complement the query embedding with cross-modal knoyvledge retrieved from the memory 150. This cross-modal fusion can provide a significantly greater performance boost than performing uni-modal fusion because the signal brought by cross-modal fusion is richer due to the complementarity- of the different modalities. Hoyvever, because the retrieval is still uni- modal, the system 100 can more readily find the elements that are most relevant to a given query- input.

Finally, the system 100 can process the fused embedding 108 to generate a prediction 112 for the multi-modal task. In some cases, generating the prediction 112 can require generating respective embeddings of one or more additional query- inputs, e.g., inputs that are the opposite modality of the query input used to generate the queryembedding 104.

Generating predictions for tasks and performing cross-modal fusion will be described in more detail belo v.

FIG. 2 is a flow diagram of an example process 200 for generating a prediction for a query image. For convenience, the process 200 will be described as being performed by a system of one or more computers located in one or more locations. For example, a neural network system, e.g., the neural network system 100 of FIG.1, appropriately programmed, can perform the process 200.

That is, in the example 200, the system is performing a task that requires generating a prediction for an input that includes a query image.

The system maintains, in a memory. memory' data (step 202). As described above, the memory data includes, for each of a plurality of pairs, (i) a respective memory image embedding of the memory' image in the pair in the embedding space and (ii) a respective memory text embedding of the memory- text in the pair in the embedding space.

The system receives a query image (step 204).

The system processes the query image using an image embedding neural network to generate a query image embedding of the uery image in the embedding space (step 206).

The system identifies, as context images, a subset of the plurality of memory images using the query image embedding and the respective memory image embeddings for the plurality' of memory' images (step 208).

Generally, to identify the context images, the system performs a search of the memory' image embeddings to identify, as the context images, the memory images that have the k most similar memory image embeddings to the query- image embedding according to a similarity measure. Here, k is a fixed integer that is less than the total number of memory images. In general, k can be significantly smaller than the total number of memory' images.

The system generates, from the query' image embedding and respective memory' text embeddings for each memory text that is in the same pair as one of the context images, i.e., as the memory' images in the subset, a fused query image embedding (step 210). That is, rather than directly using the context images to generate the fused query' image embedding, the system instead uses memory' text embeddings for the memory texts that are in the same pair as one of the context images.

The system can generate the fused embedding in any of a variety of ways.

As one example, the system can directly combine, e g., concatenate, sum, or average, the query image embedding and the respective memory' text embeddings.

As one example, the system can process the query image embedding and the respective memory text embeddings using an image query’ fusion neural network to generate the fused query image embedding. For example, the image query' fusion neural network can be a Transformer neural network. The Transformer neural network can process the query image embedding and the memory text embeddings through one or more self-attention layer blocks, i.e., layer blocks that each apply self-attention over the embeddings, to generate the fused query image embedding.

In some cases, making use of the query’ image fusion neural network can improve the performance of the system, i.e., relative to directly combining the embeddings, because it allows the model to learn how to incorporate the context from the memory' embeddings, by, for example, learning how to omit or enhance some of the retrieved elements.

The system generates a prediction for a machine learning task for the query image using the fused query' image embedding (step 212).

As one example, to perform the task, the system can identify a set of query texts.

For example, the set of query texts can be some or all of the memory’ texts in the pairs in the memory data. That is, in this example, the machine learning task is image-to- text retrieval from the memory texts in at least a subset of the pairs.

As another example, some or all of the query texts can be texts that are not in the memory' data. For example, the machine learning task can be image classification, and the set of query texts can be received along with the query image.

The system can then obtain either fused query text embeddings for each of the query’ texts or query text embeddings for each of the query texts. That is, the system can obtain either a fused query’ text embedding generated through cross-modal fusion as described below with reference to FIG. 4 or a query text embedding generated by processing the query text using the text embedding neural network.

When the query' text is one of the memory texts, the system can have precomputed the fused query text embedding or query text embedding of the query text and can access the corresponding embedding from memory. When the query text is not one of the memory texts, the system can generate the fused query text embedding or query text embedding of the query text as described below with reference to FIG. 3.

In either case, the system can then generate the prediction for the machine learning task by identifying, as most similar query' texts, one or more query' texts that have query text embeddings or fused query' text embeddings that are most similar to the fused query image embedding according to a similarity measure, e.g., inner product or Euclidean distance, and then providing, as the prediction, data identifying the one or more most similar query texts.

FIG. 3 shows an example 300 of the operation of the system when receiving a query image. As show n in the example 300, the system performs a uni-modal search to identify context images for the query image in the memory, but then performs cross- modal fusion by fusing the embedding of the query⁷ image with memory⁷ text embeddings of the texts that are in the same pairs as the context images. This is in contrast to uni- modal fusion, which would fuse the embedding of the query' image with the memory image embeddings of the context images.

FIG. 4 is a flow diagram of an example process 400 for generating a prediction or a query text. For convenience, the process 400 will be described as being performed by a system of one or more computers located in one or more locations. For example, a neural network system, e.g., the neural network system 100 of FIG. 1, appropriately programmed, can perform the process 400.

That is, in the example 400, the system is performing a task that requires generating a prediction for an input that includes a query text.

The system maintains, in a memory, memory data (step 402) as described above.

The system receives a query text (step 404).

The system processes the query⁷ text using a text embedding neural netw ork to generate a query text embedding of the query text in the embedding space (step 406).

The system identifies, as context texts, a subset of the plurality of memory texts using the query text embedding and the respective memory text embeddings for the plurality' of memory' texts (step 408).

Generally, to identify the context texts, the system performs a search of the memory text embeddings to identify, as the context texts, the memory⁷ texts that have the k most similar memory text embeddings to the query text embedding according to a similarity⁷ measure. Here, k is a fixed integer that is less than the total number of memory⁷ texts. In general, k can be significantly smaller than the total number of memory texts.

The system generates, from the query text embedding and respective memory image embeddings for each memory image that is in the same pair as one of the context texts, i.e., as the memory⁷ texts in the subset, a fused query⁷ text embedding (step 410). That is, rather than directly using the context texts to generate the fused query' text embedding, the system instead uses memory image embeddings for the memory images that are in the same pair as one of the context texts.

The system can generate the fused embedding in any of a variety of w ays. As one example, the system can directly combine, e g., concatenate, sum, or average, the query text embedding and the respective memory’ image embeddings.

As one example, the system can process the query text embedding and the respective memory⁷ image embeddings using a text query fusion neural network to generate the fused query text embedding.

For example, the text query⁷ fusion neural network can be a Transformer neural network. The Transformer neural network can process the query image embedding and the memory text embeddings through one or more self-attention layer blocks, i.e., layer blocks that each apply self-attention over the embeddings, to generate the fused query⁷ text embedding.

The system generates a prediction for a machine learning task for the query text using the fused query⁷ text embedding (step 412).

As one example, to perform the task, the system can identify a set of query images.

For example, the set of query images can be some or all of the memory images in the pairs in the memory data. That is, in this example, the machine learning task is text- to-image retrieval from the memory images in at least a subset of the pairs.

As another example, some or all of the query⁷ images can be images that are not in the memory data. For example, the set of query images can be received along with the query text.

The system can then obtain either fused query image embeddings for each of the query⁷ images or query⁷ image embeddings for each of the query images. That is, the system can obtain either a fused query image embedding generated through cross-modal fusion as described above with reference to FIG. 2 or a query image embedding generated by processing the query image using the image embedding neural network.

When the query⁷ image is one of the memory⁷ images, the system can have precomputed the fused query⁷ image embedding or query image embedding of the query image and can access the corresponding embedding from memory⁷. When the query image is not one of the memory images, the system can generate the fused query image embedding or query image embedding of the query image as described above with reference to FIG. 2.

In either case, the system can then generate the prediction for the machine learning task by identifying, as most similar query images, one or more query⁷ images that have query image embeddings or fused query image embeddings that are most similar to the fused query text embedding according to a similarity measure, e.g., inner product or Euclidean distance, and then providing, as the prediction, data identifying the one or more most similar query images.

FIG. 5 shows an example 500 of the operation of the system when receiving a query text. As shown in the example 500, the system performs a uni-modal search to identify context texts for the query text in the memory, but then performs cross-modal fusion by fusing the embedding of the query text with memory image embeddings of the images that are in the same pairs as the context texts. This is in contrast to uni-modal fusion, which would fuse the embedding of the query text with the memory' text embeddings of the context texts.

As described above, in some implementations, the system uses the text query fusion neural network or the image query fusion neural network when generating the fused query' embedding of a given query' input. In these implementations, prior to using the text query' fusion neural network or the image query fusion neural network, the system or another training system trains these neural networks to improve the quality of the fused query embeddings.

In particular, the training system can train the image query' fusion neural network and the text query' fusion neural netw ork jointly to minimize a contrastive loss on a training data set.

Generally, the training data set for this training includes, for each of a plurality of training image - training text pairs, (i) a respective training image embedding of the training image in the training image - training text pair in the embedding space and (ii) a respective training text embedding of the training text in the training image - training text pair.

The training pairs can generally be the same as the memory pairs or different from the memory pairs, i.e., at least some of the training image - memory text pairs are different from the memory image - memory text pairs, to ensure that the learned fusion generalizes after training. When some of the training pairs are the same as some of the memory pairs, the system can use. as the training embeddings of a given training pair, the memory embeddings of the given pair, i.e., reuse the embeddings rather than recomputing new' embeddings.

In other words, after training, the system can effectively use the trained neural networks to retrieve information from a set of memory pairs that is entirely or partially different from the set of pairs on which the neural networks were trained without any further retraining of the neural networks.

The contrastive loss generally includes a first contrastive loss term that operates on fused image embeddings generated by the image query fusion neural network for the training images and fused text embeddings generated by the text query⁷ fusion neural network for the training texts.

Optionally, the contrastive loss can also include one or more additional terms.

For example, the loss can include a second contrastive loss term that operates on fused image embeddings generated by the image query fusion neural network for the training images and the training text embeddings for the training texts.

As another example, the loss can include a third contrastive loss term that operates on training image embeddings for the training images and fused text embeddings generated by the text query fusion neural network for the training texts.

Including these additional contrastive losses can align the fused embeddings with the (unfused) image and text embeddings, allowing similarity between fused and unfused embeddings to be effectively compared when performing a given task after training.

The contrastive loss terms can generally be any appropriate contrastive loss terms. As one example, the loss terms can be InfoNCE losses. As another example, the loss terms can be focal contrastive losses.

During this training, the system generally keeps the text embedding neural network and the image embedding neural network fixed, i.e., does not train the text and image embedding neural networks.

FIG. 6 shows an example 600 of the training of the query fusion neural networks on a training a pair.

As shown in FIG. 6, the system receives a training pair that includes an image I and a corresponding text T. For the training image I, the system uses the image embedding neural network (“image encoder”) to process the image I to generate an image embedding v. The system then uses uni-modal search and cross-modal retrieval with the memory to return a set of text embeddings, as described above.

The system then processes the image embedding v and the text embeddings using the image query fusion neural network to generate a fused image embedding v.

Similarly, for the training image T, the system uses the text embedding neural network (“text encoder”) to process the text T to generate a text embedding t. The system then uses uni-modal search and cross-modal retrieval with the memory to return a set of image embeddings, as described above.

The system then processes the text embedding t and the image embeddings using the text query' fusion neural network to generate a fused text embedding t.

The system can then apply a contrastive loss, e.g., the InfoNCE loss in the example of FIG .6. that operate on the fused image embedding v and the fused text embedding t.

Optionally, the system can also apply a contrastive loss that operates on the fused text embedding t and the image embedding v and a contrastive loss that operates on the fused image embedding v and the text embedding t.

FIG. 7 shows an example 700 of the performance of the described techniques on an image classification task and, more specifically, a zero-shot image classification task. In particular, the example 700 shows the performance on five different data sets (Cars, CUB, Flowers, IMlk, P1365) in terms of top-1 accuracy of various techniques that incorporate a memory relative to a baseline technique that directly makes predictions from the embedding of the query image without performing any fusion. The performance of the baseline technique for each data set is shown in the second row of the table directly underneath the name of the data set and the values in parentheses show the delta relative to the baseline technique.

Results 1-4 are for variants that generate the fused embedding using transformer fusion neural networks with different search and fusion strategies. As can be seen from the example 700, the described techniques (uni -modal search and cross-modal fusion) significantly outperform not only the baseline technique but also outperform the other possible search and fusion strategies across all five data sets.

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. 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.

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.

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.

In this specification, the term “database” is used broadly to refer to any collection of data: the data does not need to be structured in any particular way, or structured at all, and it can be stored on storage devices in one or more locations. Thus, for example, the index database can include multiple collections of data, each of which may be organized and accessed differently.

Similarly, in this specification the term “engine” is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine will be implemented as one or more software modules or components, installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in other cases, multiple engines can be installed and running on the same computer or computers.

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.

Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read only memory or a random access memory or both. The essential elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special purpose logic circuitry. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto optical disks, or optical disks. However, a computer need not have such devices. Moreover, a computer can be embedded in another device, e.g., a mobile telephone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few. 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.

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.

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.

Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework or a Jax framework.

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 sen- er, or that includes a front end component, e.g., a client computer having a graphical user interface, a web brow ser, or an app through which a user can interact with an implementation of the subj ect 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.

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.

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.

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 softw are products.

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. What is claimed is:

Claims

1. A method performed by one or more computers, the method comprising: maintaining, for each of a plurality of memory image - memory text pairs, (i) a respective memory image embedding of the memory image in the memory image - memory text pair in an embedding space and (ii) a respective memory text embedding of the memory text in the memory image - memory text pair in the embedding space; receiving a query image; processing the query image using an image embedding neural network to generate a query image embedding of the query image in the embedding space; identifying, as context images, a subset of the plurality of memory images using the query image embedding and the respective memory image embeddings for the plurality of memory images; generating, from the query image embedding and respective memory text embeddings for each memory text that is in a same memory image - memory' text pair as one of the memory⁷ images in the subset, a fused query image embedding; and generating a prediction for a machine learning task for the query⁷ image using the fused query image embedding.

2. The method of claim 1, wherein generating, from the query image embedding and respective memory⁷ text embeddings for each memory⁷ text that is in a same memory⁷ image - memory text pair as one of the memory images in the subset, a fused queryimage embedding comprises: processing the query image embedding and the respective memory text embeddings for the texts that are in a same memory image - memory⁷ text pair as one of the memory⁷ images in the subset using an image query⁷ fusion neural network to generate the fused query image embedding.

3. The method of claim 2, wherein the image query⁷ fusion neural network is a Transformer neural network.

4. The method of any preceding claim, wherein the memory image embeddings have been generated by processing the memory images using the image embedding neural network.

5. The method of any preceding claim, wherein identifying, as context images, a subset of the plurality of memory images using the query image embedding and the respective memory image embeddings for the plurality of memory images comprises: performing a search of the memory’ image embeddings to identify, as the context images, the memory images that have the k most similar memory image embeddings to the query image embedding according to a similarity measure, wherein k is a fixed integer that is less than a total number of memory’ images.

6. The method of any preceding claim, further comprising: receiving data identifying a plurality of query texts; obtaining a respective query text embedding of each query' text in the embedding space, wherein generating a prediction for a machine learning task for the query' image using the fused query’ image embedding comprises: identifying, as most similar query texts, one or more query texts that have query’ text embeddings that are most similar to the fused query image embedding according to a similarity measure; and providing, as the prediction, data identifying the one or more most similar query texts.

7. The method of any one of claims 1-5, further comprising: receiving data identifying a plurality of query' texts; obtaining a respective fused query text embedding of each query text in the embedding space, wherein the respective fused query’ text embedding of each query text has been generated by: processing the query text using a text embedding neural network to generate a query text embedding of the query text in the embedding space; identifying, as context texts, a subset of the plurality of memory texts using the query’ text embedding and the respective memory text embeddings for the plurality of memory texts; generating, from the query text embedding and respective memory image embeddings for each image that is in a same memory image - memory text pair as one of the memory texts in the subset, the fused query' text embedding for the query text, wherein generating a prediction for a machine learning task for the query image using the fused query' image embedding comprises: identifying, as most similar query texts, one or more query texts that have fused query' text embeddings that are most similar to the fused query' image embedding according to a similarity measure; and providing, as the prediction, data identifying the one or more most similar query texts.

8. The method of claim 6 or claim 7, wherein the machine learning task is image-to- text retrieval from the memory texts in at least a subset of the memory image - memory text pairs, wherein the query texts identified by the received data are the memory texts, and wherein the respective query text embeddings for the query' texts are the memory text embeddings for the memory texts.

9. The method of claim 8 when dependent on claim 7, wherein the respective fused query' text embeddings have been pre-computed prior to receiving the query image.

10. The method of claim 6 or claim 7, wherein the machine learning task is image classification, wherein the respective query text embeddings for the query texts are generated in response to receiving the query image.

11. The method of any preceding claim, when dependent on claim 7, wherein generating, from the query text embedding and respective memory image embeddings for each image that is in a same memory' image - memory text pair as one of the memory texts in the subset, a fused query' text embedding comprises: processing the query text embedding and the respective memory' image embeddings for the images that are in a same memory image - memory text pair as one of the memory texts in the subset using a text query fusion neural network to generate the fused query text embedding.

12. The method of claim 11, wherein the text query fusion neural network is a Transformer neural network.

13. A method performed by one or more computers, the method comprising: maintaining, for each of a plurality of memory image - memory text pairs, (i) a respective memory image embedding of the memory image in the memory image - memory text pair in an embedding space and (ii) a respective memory text embedding of the text in the memory image - memory text pair; receiving a query text; processing the query text using a text embedding neural network to generate a query text embedding of the query text in the embedding space; identifying, as context texts, a subset of the plurality of memory' texts using the query text embedding and the respective memory text embeddings for the plurality of memory texts; generating, from the query text embedding and respective memory image embeddings for each image that is in a same memory image - memory text pair as one of the memory' texts in the subset, a fused query' text embedding; and generating a prediction for a machine learning task for the query’ text using the fused query text embedding.

14. The method of claim 13, wherein generating, from the query' text embedding and respective memory image embeddings for each image that is in a same memory image - memory text pair as one of the memory texts in the subset, a fused query text embedding comprises: processing the query’ text embedding and the respective memory' image embeddings for the images that are in a same memory image - memory text pair as one of the memory’ texts in the subset using a text query fusion neural network to generate the fused query text embedding.

15. The method of claim 14, wherein the text query’ fusion neural network is a Transformer neural network.

16. The method of any one of claims 13-15, wherein the memory text embeddings have been generated by processing the memory texts using the text embedding neural network.

17. The method of any one of claims 13-16, wherein identifying, as context texts, a subset of the plurality of memory texts comprises: performing a search of the memory text embeddings to identify, as the context texts, the memory texts that have the k most similar memory text embeddings to the query text embedding according to a similarity measure, wherein Ar is a fixed integer that is less than a total number of memory texts.

18. The method of any one of claims 13-17, further comprising: receiving data identifying a plurality⁷ of uery images; obtaining a respective query image embedding of each query⁷ image in the embedding space, wherein generating a prediction for a machine learning task for the query' text using the fused query text embedding comprises: identifying, as most similar query images, one or more query images that have query image embeddings that are most similar to the fused query⁷ text embedding according to a similarity measure; and providing, as the prediction, data identifying the one or more most similar query⁷ images.

19. The method of any one of claims 13-17, further comprising: receiving data identifying a plurality of query image; obtaining a respective fused query image embedding of each query image in the embedding space, wherein the respective fused query image embedding of each query image has been generated by: processing the query image using an image embedding neural network to generate a query image embedding of the query image in the embedding space; identifying, as context images, a subset of the plurality of memory images using the query image embedding and the respective memory image embeddings for the plurality of memory texts; generating, from the query image embedding and respective memory text image embeddings for each memory⁷ text that is in a same memory' image - memory' text pair as one of the memory images in the subset, the fused query image embedding for the query text, wherein generating a prediction for a machine learning task for the query text using the fused query' text embedding comprises: identifying, as most similar query’ images, one or more query images that have fused query image embeddings that are most similar to the fused query text embedding according to a similarity measure; and providing, as the prediction, data identifying the one or more most similar query images.

20. The method of claim 18 or claim 19, wherein the machine learning task is text-to- image retrieval from the images in at least a subset of the memory' image - memory' text pairs, wherein the query images identified by the received data are the memory images, and wherein the respective query’ image embeddings for the query’ images are the memory image embeddings for the memory' images.

21. The method of claim 20 when dependent on claim 19, wherein the respective fused query’ image embeddings have been pre-computed prior to receiving the query text.

22. The method of claim 18 or claim 19, wherein the machine learning task is text-to- image retrieval from images not in the memory’ image - memory’ text pairs, wherein the respective query image embeddings for the query image are generated in response to receiving the query text.

23. The method of any preceding claim, when dependent on claim 19, wherein generating, from the query image embedding and respective memory text embeddings for each image that is in a same memory’ image - memory’ text pair as one of the memory images in the subset, a fused query image embedding comprises: processing the query image embedding and the respective memory text embeddings for the texts that are in a same memory image - memory text pair as one of the memory’ images in the subset using an image query fusion neural network to generate the fused query image embedding.

24. The method of any one of claims 13-23, wherein generating a prediction for a machine learning task for the query' text using the fused query text embedding comprises: processing an input comprising the fused query text embedding using a model for the machine learning task to generate the prediction.

25. The method of any one of claims 1-12, wherein generating a prediction for a machine learning task for the query image using the fused query image embedding comprises: processing an input comprising the fused query image embedding using a model for the machine learning task to generate the prediction.

26. The method of any preceding claim, wherein the image query fusion neural network and the text query fusion neural network have been jointly trained to minimize a contrastive loss on a training data set that comprises, for each of a plurality of training image - training text pairs, (i) a respective training image embedding of the training image in the training image - training text pair in the embedding space and (ii) a respective training text embedding of the training text in the training image - training text pair, and wherein the contrastive loss comprises a first contrastive loss term that operates on fused image embeddings generated by the image query fusion neural network for the training images and fused text embeddings generated by the text query fusion neural network for the training texts.

27. The method of claim 26, wherein the contrastive loss comprises a second contrastive loss term that operates on fused image embeddings generated by the image query' fusion neural network for the training images and the training text embeddings for the training texts.

28. The method of claim 26 or claim 27, wherein the contrastive loss comprises a third contrastive loss term that operates on training image embeddings for the training images and fused text embeddings generated by the text query' fusion neural network for the training texts.

29. The method of any one of claims 26-28, wherein at least some of the training image - memory' text pairs are different from the memory' image - memory' text pairs.

30. The method of any one of claims 26-29, wherein at least some of the training image - memory text pairs are the same as at least some of the memory image - memory’ text pairs, and wherein, for each training image - training text pair that is the same as a corresponding memory image - memory’ text pair, (i) the respective training image embedding of the training image in the training image - training text pair is the same as the respective memory’ image embedding of the memory image in the corresponding memory' image - memory⁷ text pair and (ii) the respective training text embedding of the training text in the training image - training text pair is the same as the respective memory text embedding of the memory text in the corresponding memory image - memory text pair.

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

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