Input data item classification using memory data item embeddings
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.
1 . A method performed by one or more computers, the method comprising:
maintaining, for each of a plurality of memory data items, (i) a respective key embedding that has been generated by processing the memory data item using a first embedding neural network and (ii) a respective value embedding that has been generated by processing at least one of the memory data item or a corresponding data item associated with the memory data item using a second embedding neural network;
receiving a query data item; and
classifying the query data item using a semi-parametric model that uses the respective key and value embeddings for the plurality of memory data items, comprising:
processing the query data item using the first embedding neural network to generate a query embedding of the query data item;
searching through the respective key embeddings of the memory data items by using the query embedding as a query to identify a subset of memory data item;
generating, from the query embedding and the respective value embeddings for the identified subset of memory data items that were identified from searching through the respective key embeddings of the memory data items by using the query embedding as the query, an input to a classifier neural network, wherein the input comprises a set of one or more embeddings generated by fusing the query embedding and the respective value embeddings for the identified subset of memory data items, wherein the fusing comprises (i) generating the input based on the query embedding and a combined value embedding generated from the respective value embeddings, or (ii) processing the query embedding and the respective value embeddings through a sequence of one or more attention layers; and
processing the input using the classifier neural network to generate a classification output for the query data item that specifies one or more categories to which the query data item belongs.
2 . The method of claim 1 , wherein searching through the respective key embeddings of the memory data items by using the query embedding as a query to identify a subset of memory data items comprises:
performing a search of the key embeddings of the memory data items to identify, as the identified subset of memory data items, the memory data items that have the k most similar key embeddings to the query embedding according to a similarity measure, wherein k is a fixed integer that is less than a total number of memory data items.
3 . The method of claim 2 , wherein performing a search of the key embeddings comprises performing an approximate k-nearest neighbors search through the key embeddings.
4 . The method of claim 1 , wherein generating the input based on the query embedding and a combined value embedding generated from the respective value embeddings comprises:
summing the combined value embedding and the query embedding to generate the input to the classifier neural network.
5 . The method of claim 1 , wherein generating a combined value embedding from the respective value embeddings for the identified subset of memory data items comprises:
generating an initial combined value embedding from the respective value embeddings for the identified subset of memory data items.
6 . The method of claim 5 , wherein generating a combined value embedding from the respective value embeddings for the identified subset of memory data items comprises:
applying a dense neural network layer to the initial combined value embedding to generate the combined value embedding.
7 . The method of claim 6 , wherein the query embedding has a first dimensionality, each value embedding has a second, different dimensionality and wherein the dense neural network layer maps the initial combined value embedding from the second dimensionality to the first dimensionality.
8 . The method of claim 5 , wherein generating an initial combined value embedding from the respective value embeddings for the identified subset of memory data items comprises:
computing a mean of the respective value embeddings for the identified subset of memory data items.
9 . The method of claim 1 , wherein each attention layer in the sequence of one or more attention layers is configured to:
receive an input query embedding,
compute a respective attention weight for each identified subset of memory data item by, for each identified subset of memory data item, computing an attention weight between the input query embedding and the key embedding for the identified subset of memory data item,
compute an aggregated value embedding by computing a weighted sum of the value embeddings for the identified subset of memory data items in accordance with the respective attention weights, and
use the aggregated value embedding to update the input query embedding, and
the input query embedding for the first attention layer in the sequence is the query embedding.
10 . The method of claim 9 , wherein the updated query embedding generated by the last attention layer in the sequence is the input to the classifier neural network.
11 . The method of claim 9 , wherein the sequence comprises a plurality of attention layers, and wherein the input query embedding for each attention layer after the first attention layer in the sequence is the updated query embedding generated by a preceding attention layer in the sequence.
12 . The method of claim 9 , wherein using the aggregated value embedding to update the input query embedding comprises:
generating an initial updated query embedding from the aggregated value embedding; and
combining the initial updated query embedding and the query embedding to generate the updated query embedding.
13 . The method of claim 12 , wherein combining the initial updated query embedding and the query embedding to generate the updated query embedding comprises:
summing the initial updated query embedding and the query embedding to generate the updated query embedding.
14 . The method of claim 12 , wherein generating an initial updated query embedding from the aggregated value embedding comprises:
applying a dense neural network layer to the aggregated value embedding to generate the initial updated query embedding.
15 . The method of claim 14 , wherein the query embedding has a first dimensionality, each value embedding has a second, different dimensionality and wherein the dense neural network layer maps the aggregated value embedding from the second dimensionality to the first dimensionality.
16 . The method of claim 1 , wherein the query data item is an image.
17 . The method of claim 1 wherein the classifier neural network has been trained on labeled training data for a classification task.
18 . The method of claim 17 wherein the classifier neural network has been trained jointly with one or more attention layers on the labeled training data for the classification task.
19 . The method of claim 17 , wherein the first embedding neural network has been pre-trained prior to the training of the classifier neural network on the labeled training data for the classification task.
20 . The method of claim 19 , wherein the pre-trained first embedding neural network is held frozen during the training of the classifier neural network on the labeled training data for the classification task.
21 . The method of claim 19 , wherein the pre-trained first embedding neural network is fine-tuned during the training of the classifier neural network on the labeled training data for the classification task.
22 . The method of claim 17 , wherein the memory and key embeddings were generated prior to the training of the classifier neural network on the labeled training data for the classification task and are held frozen during the training of the classifier neural network on the labeled training data for the classification task.
23 . The method of claim 1 , wherein the key embeddings have been generated by processing the memory data items using the first embedding neural network after the first embedding neural network has been pre-trained and prior to training the classifier neural network.
24 . The method of claim 1 , wherein the second embedding neural network is different from the first embedding neural network.
25 . The method of claim 24 , wherein the second embedding neural network has (i) more parameters than the first embedding neural network, (ii) generates embeddings that have a higher dimensionality than the embeddings generated by the first embedding neural network, or (iii) both.
26 . The method of claim 1 , wherein the memory data item is of a first type and the corresponding data item is of a different, second type, and, for each memory data item, the value embedding has been generated by processing the corresponding data item using the second embedding neural network.
27 . The method of claim 26 , wherein the corresponding data item for each memory data item is text describing the memory data item.
28 . 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 operations comprising:
maintaining, for each of a plurality of memory data items, (i) a respective key embedding that has been generated by processing the memory data item using a first embedding neural network and (ii) a respective value embedding that has been generated by processing at least one of the memory data item or a corresponding data item associated with the memory data item using a second embedding neural network;
receiving a query data item; and
classifying the query data item using a semi-parametric model that uses the respective key and value embeddings for the plurality of memory data items, comprising:
processing the query data item using the first embedding neural network to generate a query embedding of the query data item;
searching through the respective key embeddings of the memory data items by using the query embedding as a query to identify a subset of memory data item;
generating, from the query embedding and the respective value embeddings for the identified subset of memory data items that were identified from searching through the respective key embeddings of the memory data items by using the query embedding as the query, an input to a classifier neural network, wherein the input comprises a set of one or more embeddings generated by fusing the query embedding and the respective value embeddings for the identified subset of memory data items, wherein the fusing comprises (i) generating the input based on the query embedding and a combined value embedding generated from the respective value embeddings, or (ii) processing the query embedding and the respective value embeddings through a sequence of one or more attention layers; and
processing the input using the classifier neural network to generate a classification output for the query data item that specifies one or more categories to which the query data item belongs.
29 . 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 operations comprising:
maintaining, for each of a plurality of memory data items, (i) a respective key embedding that has been generated by processing the memory data item using a first embedding neural network and (ii) a respective value embedding that has been generated by processing at least one of the memory data item or a corresponding data item associated with the memory data item using a second embedding neural network;
receiving a query data item; and
classifying the query data item using a semi-parametric model that uses the respective key and value embeddings for the plurality of memory data items, comprising:
processing the query data item using the first embedding neural network to generate a query embedding of the query data item;
searching through the respective key embeddings of the memory data items by using the query embedding as a query to identify a subset of memory data item;
generating, from the query embedding and the respective value embeddings for the identified subset of memory data items that were identified from searching through the respective key embeddings of the memory data items by using the query embedding as the query, an input to a classifier neural network, wherein the input comprises a set of one or more embeddings generated by fusing the query embedding and the respective value embeddings for the identified subset of memory data items, wherein the fusing comprises (i) generating the input based on the query embedding and a combined value embedding generated from the respective value embeddings, or (ii) processing the query embedding and the respective value embeddings through a sequence of one or more attention layers; and
processing the input using the classifier neural network to generate a classification output for the query data item that specifies one or more categories to which the query data item belongs.
30 . The method of claim 1 , wherein the respective key embeddings and the respective value embeddings of the plurality of memory data items are stored in a database, wherein the database is maintained separately from the classifier neural network, and wherein searching through the respective key embeddings of the memory data comprises searching the database.