Efficient embedding table storage and lookup
The present disclosure provides systems, methods, and computer program products for providing efficient embedding table storage and lookup in machine-learning models. A computer-implemented method may include obtaining an embedding table comprising a plurality of embeddings respectively associated with a corresponding index of the embedding table, compressing each particular embedding of the embedding table individually allowing each respective embedding of the embedding table to be decompressed independent of any other embedding in the embedding table, packing the embedding table comprising individually compressed embeddings with a machine-learning model, receiving an input to use for locating an embedding in the embedding table, determining a lookup value based on the input to search indexes of the embedding table, locating the embedding based on searching the indexes of the embedding table for the determined lookup value, and decompressing the located embedding independent of any other embedding in the embedding table.
1 . A computer-implemented method, comprising:
receiving, by a computing system comprising one or more processors and a memory, an input for processing using a machine-learned model;
accessing, by the computing system, an embedding table associated with the machine-learned model, the embedding table comprising a plurality of embeddings;
locating, by the computing system based on one or more lookup values associated with the input, one or more embeddings in the embedding table that respectively correspond to the one or more lookup values;
loading, by the computing system and into the memory, and based on the locating of the one or more embeddings, the one or more embeddings without loading the entirety of the embedding table into the memory; and
processing, by the computing system, the one or more embeddings in the memory using the machine-learned model.
2 . The computer-implemented method of claim 1 , wherein the embedding table is associated with a plurality of items included in a vocabulary for the machine-learned model.
3 . The computer-implemented method of claim 2 , wherein the plurality of items comprise a plurality of tokens.
4 . The computer-implemented method of claim 3 , wherein the machine-learned model is configured to perform natural language processing.
5 . The computer-implemented method of claim 1 , wherein the one or more embeddings are embeddings compressed using pruning.
6 . The computer-implemented method of claim 1 , wherein the one or more lookup values correspond to one or more index values that index the one or more embeddings in the embedding table.
7 . The computer-implemented method of claim 1 , comprising:
loading, by the computing system and into the memory, the one or more embeddings without loading, into the memory, any other embedding that is not related to the input.
8 . The computer-implemented method of claim 1 , comprising:
accessing the one or more embeddings using an index value associated with a logical storage unit storing at least one of the one or more embeddings.
9 . The computer-implemented method of claim 1 , wherein the one or more embeddings correspond to features of word content that are mapped to vectors of real numbers used by the machine-learned model to generate a prediction output.
10 . A computing system, comprising:
one or more processors; and
a memory comprising one or more non-transitory computer-readable media that store instructions that, when executed by the one or more processors, cause the computing system to perform operations comprising:
receiving an input for processing using a machine-learned model;
accessing an embedding table associated with the machine-learned model, the embedding table comprising a plurality of embeddings;
locating, based on one or more lookup values associated with the input, one or more embeddings in the embedding table that respectively correspond to the one or more lookup values;
loading, into the memory, and based on the locating of the one or more embeddings, the one or more embeddings without loading the entirety of the embedding table into the memory; and
processing the one or more embeddings in the memory using the machine-learned model.
11 . The computing system of claim 10 , wherein the embedding table is associated with a plurality of items included in a vocabulary for the machine-learned model.
12 . The computing system of claim 11 , wherein the plurality of items comprise a plurality of tokens.
13 . The computing system of claim 12 , wherein the machine-learned model is configured to perform natural language processing.
14 . The computing system of claim 10 , wherein the one or more lookup values correspond to one or more index values that index the one or more embeddings in the embedding table.
15 . The computing system of claim 10 , the operations comprising:
loading, by the computing system and into the memory, the one or more embeddings without loading, into the memory, any other embedding that is not related to the input.
16 . The computing system of claim 10 , the operations comprising:
accessing the one or more embeddings using an index value associated with a logical storage unit storing at least one of the one or more embeddings.
17 . The computing system of claim 10 , wherein the one or more embeddings correspond to features of word content that are mapped to vectors of real numbers used by the machine-learned model to generate a prediction output.
18 . One or more non-transitory computer-readable media that store instructions that, when executed by one or more processors, cause a computing system to perform operations comprising:
receiving an input for processing using a machine-learned model;
accessing an embedding table associated with the machine-learned model, the embedding table comprising a plurality of embeddings;
locating, based on one or more lookup values associated with the input, one or more embeddings in the embedding table that respectively correspond to the one or more lookup values;
loading, into a memory of the computing system, and based on the locating of the one or more embeddings, the one or more embeddings without loading the entirety of the embedding table into the memory; and
processing the one or more embeddings in the memory using the machine-learned model.
19 . The one or more non-transitory computer-readable media of claim 18 , the operations comprising:
accessing the one or more embeddings using an index value associated with a logical storage unit storing at least one of the one or more embeddings.
20 . The one or more non-transitory computer-readable media of claim 18 , wherein the one or more embeddings correspond to features of word content that are mapped to vectors of real numbers used by the machine-learned model to generate a prediction output.