US 2026/0037786 A1 · KEY VALUE NEURAL NETWORK ARCHITECTURE

Abstract

The present disclosure relates to techniques for improving inference efficiency and memory utilization in transformer-based neural networks. A modified architectural design is introduced that decouples key and value matrix generation from inter-layer dependencies, enabling statically computed or parallelizable projections across layers. The disclosed approach may eliminate the need for layer-wise prefilling, support linear-time inference, and substantially reduce the memory footprint associated with key-value (KV) caching. The disclosed architecture may use shared or layer-specific projections, with a single KV-cache serving all or subsets of layers. In some embodiments, a non-linear transformation (e.g., implemented via a feed-forward network), may preprocess input embeddings prior to query, key, and value generation. A lookup table of transformed embeddings may be precomputed to further accelerate inference. The disclosed system can enhance scalability and may allow deployment of large models on resource-constrained hardware, offering practical benefits for latency-sensitive applications and long-context processing in transformer-based models.

Claims20

Dependent claims

A computer-implemented method comprising:

accessing an input sequence comprising a plurality of input tokens to be processed by a neural network model comprising a sequence of transformer layers;
generating a sequence of embeddings corresponding to the input tokens using an embedding layer;
for each transformer layer in the sequence of transformer layers:
projecting the sequence of embeddings through a key projection matrix to generate a key sequence;
projecting the sequence of embeddings through a value projection matrix to generate a value sequence;
projecting a query token derived from a previous layer output through a query projection matrix to generate a query vector, wherein the query vector of a first layer is based on a last input of the input sequence;
computing an attention output by applying an attention mechanism using the query vector, the key sequence, and the value sequence;
and
processing the attention output through remaining components of the transformer layer to produce a layer output;
generating an output token based on a final layer output of the sequence of transformer layers;
and
wherein the key sequence and the value sequence at each layer are computed independently of outputs of any previous transformer layer, such that computation of the key sequences and the value sequences is performed with linear complexity with respect to length of the input sequence.

9 more elements

A system comprising:

one or more data processors;
and
a non-transitory computer-readable storage medium containing instructions which, when executed on the one or more data processors, cause the one or more data processors to perform a set of operations including:
accessing an input sequence comprising a plurality of input tokens to be processed by a neural network model comprising a sequence of transformer layers;
generating a sequence of embeddings corresponding to the input tokens using an embedding layer;
for each transformer layer in the sequence of transformer layers:
projecting the sequence of embeddings through a key projection matrix to generate a key sequence;
projecting the sequence of embeddings through a value projection matrix to generate a value sequence;
projecting a query token derived from a previous layer output through a query projection matrix to generate a query vector;
computing an attention output by applying an attention mechanism using the query vector, the key sequence, and the value sequence;
and
processing the attention output through remaining components of the transformer layer to produce a layer output;
generating an output token based on a final layer output of the sequence of transformer layers;
and
wherein the key sequence and the value sequence at each layer are computed independently of outputs of any previous transformer layer, such that computation of the key sequences and the value sequences is performed with linear complexity with respect to length of the input sequence.

12 more elements

A computer-program product tangibly embodied in a non-transitory machine-readable storage medium, including instructions configured to cause one or more data processors to perform a set of operations comprising:

accessing an input sequence comprising a plurality of input tokens to be processed by a neural network model comprising a sequence of transformer layers;
generating a sequence of embeddings corresponding to the input tokens using an embedding layer;
for each transformer layer in the sequence of transformer layers:
projecting the sequence of embeddings through a key projection matrix to generate a key sequence;
projecting the sequence of embeddings through a value projection matrix to generate a value sequence;
projecting a query token derived from a previous layer output through a query projection matrix to generate a query vector;
computing an attention output by applying an attention mechanism using the query vector, the key sequence, and the value sequence;
and
processing the attention output through remaining components of the transformer layer to produce a layer output;
generating an output token based on a final layer output of the sequence of transformer layers;
and
wherein the key sequence and the value sequence at each layer are computed independently of outputs of any previous transformer layer, such that computation of the key sequences and the value sequences is performed with linear complexity with respect to length of the input sequence.

9 more elements

Dependencies are inferred from the claim text.

Key value neural network architecture

Life of this patent

Key value neural network architecture

Life of this patent

Abstract

Claims20