Generating neural network outputs by enriching latent embeddings using self-attention and cross-attention operations

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 for using a neural network to generate a network output that characterizes an entity, the operations comprising:

• obtaining a representation of the entity as a set of data element embeddings;
• obtaining a set of latent embeddings;
• and
• processing:
• (i) the set of data element embeddings, and (ii) the set of latent embeddings, using the neural network to generate the network output characterizing the entity,
• wherein the neural network comprises a sequence of neural network blocks comprising:
• (i) one or more cross-attention blocks, (ii) one or more self-attention blocks, and (iii) an output block,
• wherein each cross-attention block performs operations comprising:
• updating each latent embedding in the set of latent embeddings using attention over some or all of the data element embeddings in the set of data element embeddings;
• wherein each self-attention block performs operations comprising:
• updating each latent embedding in the set of latent embeddings using attention over the set of latent embeddings;
• and
• wherein the output block performs operations comprising:
• after the set of latent embeddings are updated using the one or more cross-attention blocks and the one or more self-attention blocks, processing one or more latent embeddings from the set of latent embeddings to generate the network output characterizing the entity.

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 for using a neural network to generate a network output that characterizes an entity, the operations comprising:
• obtaining a representation of the entity as a set of data element embeddings;
• obtaining a set of latent embeddings;
• and
• processing:
• (i) the set of data element embeddings, and (ii) the set of latent embeddings, using the neural network to generate the network output characterizing the entity,
• wherein the neural network comprises a sequence of neural network blocks comprising:
• (i) one or more cross-attention blocks, (ii) one or more self-attention blocks, and (iii) an output block,
• wherein each cross-attention block performs operations comprising:
• updating each latent embedding in the set of latent embeddings using attention over some or all of the data element embeddings in the set of data element embeddings;
• wherein each self-attention block performs operations comprising:
• updating each latent embedding in the set of latent embeddings using attention over the set of latent embeddings;
• and
• wherein the output block performs operations comprising:
• after the set of latent embeddings are updated using the one or more cross-attention blocks and the one or more self-attention blocks, processing one or more latent embeddings from the set of latent embeddings to generate the network output characterizing the entity.