The transformer architecture is arguably one of the most impactful innovations in modern deep learning. Proposed in the famous 2017 paper “Attention Is All You Need,” it has become the go-to approach for most language-related modeling, including all Large Language Models (LLMs), such as the GPT family, as well as many computer vision tasks.

As the complexity and size of these models grow, so does the need to optimize their inference speed, especially in chat applications where the users expect immediate replies. Key-value (KV) caching is a clever trick to do just that – let’s see how it works and when to use it.

Transformer architecture overview

Before we dive into KV caching, we will need to take a short detour to the attention mechanism used in transformers. Understanding how it works is required to spot and appreciate how KV caching optimizes transformer inference.

We will focus on autoregressive models used to generate text. These so-called decoder models include the GPT family, Gemini, Claude, or GitHub Copilot. They are trained on a simple task: predicting the next token in sequence. During inference, the model is provided with some text, and its task is to predict how this text should continue.

From a high-level perspective, most transformers consist of a few basic building blocks:

• A tokenizer that splits the input text into subparts, such as words or sub-words.
• An embedding layer that transforms the resulting tokens (and their relative positions within the texts) into vectors.
• A couple of basic neural network layers, including dropout, layer normalization, and regular feed-forward linear layers.

The last building block missing from the list above is the slightly more involved self-attention modules.

The self-attention module is, arguably, the only advanced piece of logic in the transformer architecture. It is the cornerstone of every transformer, enabling it to focus on different parts of the input sequence when generating the outputs. It is this mechanism that gives transformers the ability to model long-range dependencies effectively.

Let’s inspect the self-attention module in more detail.

Basic self-attention module

Self-attention is a mechanism that allows the model to “pay attention” to specific parts of the input sequence as it generates the next token. For example, in generating the sentence “She poured the coffee into the cup,” the model might pay more attention to the words “poured” and “coffee” to predict “into” as the next word since these words provide context for what is likely to come next (as opposed to “she” and “the”).

Mathematically speaking, the goal of self-attention is to transform each input (embedded token) into a so-called context vector, which combines the information from all the inputs in a given text. Consider the text “She poured coffee”. Attention will compute three context vectors, one for each input token (let’s assume tokens are words).

To calculate the context vectors, self-attention computes three kinds of intermediate vectors: queries, keys, and values. The diagram below shows step by step how the context vector for the second word, “poured,” is calculated:

Let’s denote the three tokenized inputs as x1, x2, and x3, respectively. The diagram pictures them as vectors with three elements, but in practice, they will be hundreds or thousands of elements long.

As the first step, self-attention multiplies each input separately with two weight matrices, Wk and Wv. The input for which the context vector is now being computed (x2 in our case) is additionally multiplied with a third weight matrix, Wq. All three W matrices are your usual neural network weights, randomly initialized and optimized in the learning process. The outputs of this step are the keys (k) and values (v) vectors for each input, plus an additional query (q) vector for the input being processed.

In step two, the key vector of each input is multiplied by the query vector of the input being processed (our q2). The output is then normalized (not shown in the diagram) to produce the attention weights. In our example, a21 is the attention weight between the inputs “She” and “poured.”

Finally, each attention weight is multiplied by its corresponding value vector. The outputs are then summed to produce the context vector z. In our example, the context vector z2 corresponds to the input x2, “poured.” The context vectors are the outputs of the self-attention module.

If it’s easier for you to read code than diagrams, take a look at this implementation of the basic self-attention module by Sebastian Raschka. The code is part of his book, “Build A Large Language Model (From Scratch)”:

import torch

class SelfAttention_v2(torch.nn.Module):

    def __init__(self, d_in, d_out, qkv_bias=False):
        super().__init__()
        self.W_query = torch.nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key   = torch.nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = torch.nn.Linear(d_in, d_out, bias=qkv_bias)

    def forward(self, x):
        keys = self.W_key(x)
        queries = self.W_query(x)
        values = self.W_value(x)
        
        attn_scores = queries @ keys.T
        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)

        context_vec = attn_weights @ values
        return context_vec

Sebastian’s code operates on matrices: the x in his forward() method corresponds to our x1, x2, and x3 vectors stacked together as a matrix with three rows. This allows him to simply multiply x with W_key to obtain keys, a matrix consisting of three rows (k1, k2, and k3 in our example).

The important takeaway from this brief explanation of self-attention is that in each forward pass, we multiply keys with the queries and then later with the values. Keep this in mind as you read on.

Advanced self-attention modules

The variant of self-attention described above is its simplest vanilla form. Today’s largest LLMs typically use slightly modified variations that typically differ from our basic flavor in three ways:

Causal attention means that the model should only consider previous tokens in the sequence when predicting the next one, preventing it from “looking ahead” at future words. Going back to our example, “She poured coffee.”, when the model was given the word “She” and is now attempting to predict the next one (“poured” would be correct), it should not compute or have access to attention weights between “coffee” and any other word since the word “coffee” has not appeared in the text yet. Causal attention is typically implemented by masking the “look-ahead” part of the attention weights matrix with zeros.

Next, to reduce overfitting during training, dropout is often applied to the attention weights. This means that some of them are randomly set to zero in each forward pass.

Finally, basic attention can be referred to as single-head, meaning that there is just one set of Wk, Wq, and Wv matrices. An easy way to increase the model’s capacity is to switch to multi-head attention. This boils down to having multiple sets of the W-matrices and, consequently, multiple query, key, and value matrices, as well as multiple context vectors for each input.

Additionally, some transformers implement additional modifications of the attention module with the goal of improving speed or accuracy. Three popular ones are:

• Grouped-query attention: Instead of looking at every input token individually, tokens are grouped, allowing the model to focus on related groups of words at once, which speeds up processing. This is used by Llama 3, Mixtral, and Gemini.
• Paged attention: Attention is broken down into “pages” or chunks of tokens, so the model processes one page at a time, making it faster for very long sequences.
• Sliding-window attention: The model only attends to nearby tokens within a fixed “window” around each token, so it focuses on the local context without needing to look at the entire sequence.

All of these state-of-the-art approaches to implementing self-attention don’t change its basic premise and the fundamental mechanism it relies on: one always needs to multiply the keys by the queries and then later by the values. And as it turns out, at inference time, these multiplications show major inefficiencies. Let’s see why that’s the case.

What is key-value caching?

During inference, transformers generate one token at a time. When we prompt the model to start generation by passing “She,” it will produce one word, such as “poured” (for the sake of avoiding distractions, let’s keep assuming one token is one word). Then, we can pass “She poured” to the model, and it produces “coffee.” Next, we pass “She poured coffee” and obtain the end-of-sequence token from the model, indicating that it considers generation to be complete.

This means we have run the forward pass three times, each time multiplying the queries by the keys to obtain the attention scores (the same applies to the later multiplication by the values).

In the first forward pass, there was just one input token (“She”), resulting in just one key vector and one query vector. We multiplied them to obtain the q1k1 attention score.

Next, we passed “She poured” to the model. It now sees two input tokens, so the computation inside our attention module looks as follows:

We did the multiplication to compute three terms, but q1k1 was computed needlessly—we had already calculated it before! This q1k1 element is the same as in the previous forward pass because:

• q1 is calculated as the embedding of the input (“She”) times the Wq matrix,
• k1 is calculated as the embedding of the input (“She”) times the Wk matrix,
• Both the embeddings and the weight matrices are constant at inference time.

Note the grayed-out entries in the attention scores matrix: these are masked with zero to achieve causal attention. For example, the top-right element where q1k3 would have been is not shown to the model as we don’t know the third word (and k3) at the moment of generating the second word.

Finally, here is the illustration of the query-times-keys calculation in our third forward pass.

We make the computational effort to calculate six values, half of which we already know and don’t need to recompute!

You may already have a hunch about what key-value caching is all about. At inference, as we compute the keys (K) and values (V) matrices, we store their elements in the cache. The cache is an auxiliary memory from which high-speed retrieval is possible. As subsequent tokens are generated, we only compute the keys and values for the new tokens.

For example, this is how the third forward pass would look with caching:

When processing the third token, we don’t need to recompute the previous token’s attention scores. We can retrieve the keys and values for the first two tokens from the cache, thus saving computation time.

Assessing the impact of key-value caching

Key-value caching may have a significant impact on inference time. The magnitude of this impact depends on the model architecture. The more cachable computations there are, the larger the potential to reduce inference time.

Let’s analyze the impact of K-V caching on generation time using the GPT-Neo-1.3B model from EleutherAI, which is available on the Hugging Face Hub.

We will start by defining a timer context manager to calculate generation time:

import time

class Timer:

   def __enter__(self):
       self._start = time.time()
       return self

   def __exit__(self, exc_type, exc_value, traceback):
       self._end = time.time()
       self.duration = self._end - self._start

   def get_duration(self) -> float:
       return self.duration

Next, we load the model from the Hugging Face Hub, set up the tokenizer, and define the prompt:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "EleutherAI/gpt-neo-1.3B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

input_text = "Why is a pour-over the only acceptable way to drink coffee?"

Finally, we can define the function to run model inference:

def generate(use_cache):
    input_ids = tokenizer.encode(
        input_text,
        return_tensors="pt").to(device),
    )
 output_ids = model.generate(
     input_ids,
     max_new_tokens=100,
     use_cache=use_cache,
 )

Note the use_cache argument we pass to model.generate: It controls whether K-V caching is employed.

With this setup, we can measure the average generation time with and without K-V caching:

for use_cache in (False, True):
   gen_times = []
   for _ in range(10):
     with Timer() as t:
       generate(use_cache=use_cache)
     gen_times += [t.duration]
   print(f"Average inference time with use_cache={use_cache}: {np.round(np.mean(gen_times), 2)} seconds")

I have executed this code on Google Colab using their free-tier T4 GPU using torch==2.5.1+cu121 and transformers==4.46.2 on Python 3.10.12 and obtained the following output:

Average inference time with use_cache=False: 9.28 seconds
Average inference time with use_cache=True: 3.19 seconds

As you can see, in this case, the speedup from caching is almost threefold.

Challenges and trade-offs

As is usually the case, there is no such thing as a free lunch. The generation speedup we have just seen can only be achieved at the cost of increased memory usage, and it requires considerate management in production systems.

Latency-memory trade-off

Storing data in the cache uses up memory space. Systems with limited memory resources may struggle to accommodate this additional memory overhead, potentially resulting in out-of-memory errors. This is especially the case when long inputs need to be processed, as the memory required for the cache grows linearly with the input length.

Another aspect to keep in mind is that the additional memory consumed by the cache is not available for storing the batches of data. As a result, one might need to reduce the batch size to keep it within the memory limits, thus decreasing the throughput of the system.

If the memory consumed by the cache becomes a problem, one can trade additional memory for some of the model accuracy. Specifically, one can truncate the sequences, prune the attention heads, or quantize the model:

• Sequence truncation refers to limiting the maximum input sequence length, thus capping the cache size at the expense of losing long-term context. In tasks where this long context is relevant, the model’s accuracy might suffer.

• Reducing the number of layers or attention heads, thereby decreasing both the model size and cache memory requirements, is another strategy to reclaim some memory. However, reducing model complexity may impact its accuracy.

• Finally, there is quantization, which means using lower-precision data types (e.g., float16 instead of float32) for caching to reduce memory usage. Yet again, model accuracy can suffer.

To sum up, faster latency provided by K-V caching comes at the cost of increased memory usage. If there is sufficient memory, it’s a non-issue. If the memory becomes the bottleneck, however, one can reclaim it by simplifying the model in various ways, thus transitioning from a latency-memory trade-off to a latency-accuracy trade-off.

KV cache management in production systems

In large-scale production systems with many users, the K-V cache needs to be properly managed to ensure consistent and reliable response time while preventing excessive memory consumption. The two most critical aspects of this are cache invalidation (when to clear it) and cache reuse (how to use the same cache multiple times).

Cache invalidation

Three of the most popular cache invalidation strategies are session-based clearing, time-to-live invalidation, and contextual relevance-based approaches. Let’s explore them in this order.

The most basic cache invalidation strategy is session-based clearing. We simply clear the cache at the end of a user session or conversation with the model. This simple strategy is a perfect fit for applications where conversations are short and independent of each other.

Think about a customer support chatbot application in which each user session typically represents an individual conversation where the user seeks assistance with specific issues. In this context, the contents of this cache are unlikely to be needed again. Clearing the K-V cache once the user ends the chat or the session times out due to inactivity is a good choice, freeing up memory for the application to handle new users.

In situations where individual sessions are long, however, there are better solutions than session-based clearing. In time-to-live (TTL) invalidation, cache contents are automatically cleared after a certain period. This strategy is a good choice when the relevance of cached data diminishes predictably over time.

Consider a news aggregator app that provides real-time updates. Cached keys and values might only be relevant for as long as the news is hot. Implementing a TTL policy where cached entries expire after, say, one day ensures that responses to similar queries about fresh developments are generated fast while old news doesn’t fill up memory.

Finally, the most sophisticated of the three popular cache invalidation strategies is based on contextual relevance. Here, we clear the cache contents as soon as they become irrelevant to the current context or user interaction. This strategy is ideal when the application handles diverse tasks or topics within the same session, and the previous context doesn’t contribute value to the new one.

Think about a coding assistant that works as an IDE plug-in. While the user is working on a particular set of files, the cache should be retained. As soon as they switch to a different codebase, however, the previous keys and values become irrelevant and can be deleted to free memory. Contextual relevance-based approaches might be challenging to implement, though, as they require pinpointing the event or point in time at which the context switch occurs.

Cache reuse

Another important aspect of cache management is its reuse. On some occasions, a once-generated cache can be used again to speed up generation and save memory by avoiding storing the same data multiple times in different users’ cache instances.

Cache reuse opportunities typically show up when there is shared context and/or a warm start is desirable.

In scenarios where multiple requests share a common context, one can reuse the cache for that shared portion. In e-commerce platforms, certain products may have standard descriptions or specifications that are frequently asked about by multiple customers. These may include product details (“55-inch 4K Ultra HD Smart LED TV”), warranty information (“Comes with a 2-year manufacturer’s warranty covering parts and labor.”), or customer instructions (“For best results, mount the TV using a compatible wall bracket, sold separately.”). By caching the key-value pairs for these shared product descriptions, a customer support chatbot will generate responses to common questions faster.

Similarly, one can precompute and cache the initial K-V pairs for frequently used prompts or queries. Consider a voice-activated virtual assistant application. Users frequently start interactions with phrases like “What’s the weather today?” or “Set a timer for 10 minutes.” The assistant can respond more quickly by precomputing and caching the key-value pairs for these frequently used queries.

Conclusion

Key-value (K-V) caching is a technique in transformer models where the key and value matrices from previous steps are stored and reused during the generation of subsequent tokens. It allows for the reduction of redundant computations and speeding up inference time. This speedup comes at the cost of increased memory consumption. When memory is a bottleneck, one can reclaim some of it by simplifying the model, thus sacrificing its accuracy. Implementing K-V caching in large-scale production systems requires careful cache management, including choosing the strategy for cache invalidation and exploring the opportunities for cache reuse.

Scaling large language model (LLM) operations is a challenge that many of us are facing right now. For those navigating similar waters, I recently shared some thoughts about it on the Data Exchange Podcast based on our journey at neptune.ai over the last few years. 

Six years ago, we were mainly focused on MLOps when machine learning in production was still evolving. Experiment tracking back then was straightforward—dealing mostly with single models or small-scale distributed systems. Reinforcement learning was one of the few areas pushing the boundaries of scale. In that reinforcement learning, we wanted to run multiple agents and send data from multiple distributed machines to our experiment tracker. This was a huge challenge.

Scaling LLMs: from ML to LLMOps

The landscape changed two years ago when people started training LLMs at scale. LLMOps has taken center stage, and the importance of scaling large language models has grown with research becoming more industrialized. While researchers continue to lead the training process, they are also adjusting to the transition toward commercial applications.

LLMOps isn’t just MLOps with bigger servers, it is a paradigm shift for tracking experiments. We’re not tracking a few hundred metrics for a couple of hours anymore; we’re tracking thousands, even tens of thousands, over several months. These models are trained on GPU clusters spanning multiple data centers, with training jobs that can take months to complete.

Due to time constraints, training frontier models has become a production workflow rather than experimentation. When a training from scratch run takes 50,000 GPUs over several months in different data centers, you don’t get a second chance if something goes wrong—you need to get it right the first time. 

Another interesting aspect of LLM training that only a few companies have truly nailed is the branch-and-fork style of training—something that Google has implemented effectively. This method involves branching off multiple experiments from a continuously running model, requiring a significant amount of data from previous runs. It’s a powerful approach, but it demands infrastructure capable of handling large data inheritance, which makes it feasible only for a handful of companies.

From experiment tracking to experiment monitoring

Now we want to track everything—every layer, every detail—because even a small anomaly can mean the difference between success and failure and many hours of work wasted. During this time, we should not only consider pre-training and training time; post-training takes a huge amount of time and collaborative human work. Grasping this issue, we have re-engineered Neptune’s platform to efficiently ingest and visualize massive volumes of data, enabling fast monitoring and analysis at a larger scale.

One of the biggest lessons we’ve learned is that experiment tracking has evolved into experiment monitoring. Unlike MLOps, tracking is no longer just about logging metrics and reviewing them later or restarting your training from a checkpoint a few steps back. It’s about having real-time insights to keep everything on track. With such long training times, a single overlooked metric can lead to significant setbacks. That’s why we’re focusing on building intelligent alerts and anomaly detection right into our experiment tracking system.

Think of it like this—we’re moving from being reactive trackers to proactive observers. Our goal is for our platform to recognize when something is off before the researcher even knows to look for it.

Fault tolerance in LLMs

When you’re dealing with LLM training at this scale, fault tolerance becomes a critical component. With thousands of GPUs running for months, hardware failures are almost inevitable. It’s crucial to have mechanisms in place to handle these faults gracefully. 

At Neptune, our system is designed to ensure that the training can resume from checkpoints without losing any data. Fault tolerance does not only mean preventing failures; it also includes minimizing the impact when they occur, so that time and resources are not wasted.

What does this mean for enterprise teams?

If you’re creating your own LLMs from scratch, or even if you’re an enterprise fine-tuning a model, you might wonder how all this is relevant to you. Here’s the deal: strategies originally designed for handling the massive scale of training LLMs are now being adopted in other areas or by smaller-scale projects. 

Today, cutting-edge models are pushing the boundaries of scale, complexity, and performance, but these same lessons are starting to matter in fine-tuning tasks, especially when dealing with compliance, reproducibility, or complex domain-specific models.

For enterprise teams, there are three key focuses to consider:

1. Prompt Engineering: Fine-tune your prompts to get the most effective outputs. This is crucial for adapting large models to your specific needs without having to train from scratch.
2. Implement guardrails in your application: Ensuring your models behave safely and predictably is key. Guardrails help manage the risks associated with deploying AI in production environments, especially when dealing with sensitive data or critical tasks.
3. Observability in your system: Observability is vital to understanding what’s happening inside your models. Implementing robust monitoring and logging allows you to track performance, detect issues early, and ensure your models are working as expected. Neptune’s experiment tracker provides the observability you need to stay on top of your model’s behavior.

The future: what we’re building next

At Neptune, we’ve nailed the data ingestion part—it’s fast, reliable, and efficient. The challenge for the next year is making this data useful at scale. We need more than just filtering; we need smart tools that can surface the most critical insights and the most granular information automatically. The goal is to build an experiment tracker that helps researchers discover insights, not just record data.

We’re also working on developing a platform that combines monitoring and anomaly detection with the deep expertise researchers acquire over years of experience. By embedding that expertise directly into the tool (either automatically or by defining rules manually), less experienced researchers can benefit from the patterns and signals that would otherwise take years to learn.

In our paper, Understanding LLMs Requires More Than Statistical Generalization, we argue that current machine learning theory cannot explain the interesting emergent properties of Large Language Models, such as reasoning or in-context learning. From prior work (e.g., Liu et al., 2023) and our experiments, we’ve seen that these phenomena cannot be explained by reaching globally minimal test loss – the target of statistical generalization. In other words, model comparison based on the test loss is nearly meaningless.

We identified three areas where more research is required:

• Understanding the role of inductive biases in LLM training, including the role of architecture, data, and optimization.
• Developing more adequate measures of generalization.
• Using formal languages to study language models in well-defined scenarios to understand transfer performance.

In this commentary, we focus on diving deeper into the role of inductive biases. Inductive biases affect which solution the neural network converges to, such as the model architecture or the optimization algorithm. For example, Stochastic Gradient Descent (SGD) favors neural networks with minimum-norm weights.

How do language complexity and model architecture affect generalization ability?

In their Neural Networks and the Chomsky Hierarchy paper published in 2023, Delétang et al. showed how different neural network architectures generalize better for different language types.

Following the well-known Chomsky hierarchy, they distinguished four grammar types (regular, context-free, context-sensitive, and recursively enumerable) and defined corresponding sequence prediction tasks. Then, they trained different model architectures to solve these tasks and evaluated if and how well the model generalized, i.e., if a particular model architecture could handle the required language complexity.

In our position paper, we follow this general approach to expose the interaction of architecture and data in formal languages to gain insights into complexity limitations in natural language processing. We study popular architectures used for language modeling, e.g., Transformers, State-Space Models (SSMs) such as Mamba, the LSTM, and its novel extended version, the xLSTM.

To investigate how these models deal with formal languages of different complexity, we use a simple setup where each language consists only of two rules. During training, we monitor how well the models perform next-token prediction on the (in-distribution) test set, measured by accuracy.

However, our main question is whether these models generalize out-of-distribution. For this, we introduce the notion of rule extrapolation.

Can models adapt to changing grammar rules?

To understand rule extrapolation, let’s start with an example. A simple formal language is the aⁿbⁿ language, where the strings obey two rules:

Examples of valid strings include “ab” and “aabb,” whereas strings like “baab” (violates rule 1) and “aab” (violates rule 2) are invalid. Having trained on such strings, we feed the models an out-of-distribution (OOD) string, violating rule 1 (e.g., a string where the first token is b). 

We find that most models still obey rule 2 when predicting tokens, which we call rule extrapolation – they do not discard the learned rules entirely but adapt to the new situation in which rule 1 is seemingly no longer relevant. 

This finding is surprising because none of the studied model architectures includes conscious choices to promote rule extrapolation. It emphasizes our point from the position paper that we need to understand the inductive biases of language models to explain emergent (OOD) behavior, such as reasoning or good zero-/few-shot prompting performance.

Efficient LLM training requires understanding what is a complex language for an LLM

According to the Chomsky hierarchy, the context-free aⁿbⁿ language is less complex than the context-sensitive aⁿbⁿcⁿ language, where the n a’s and n b’s are followed by an equal number of c’s.

Despite their different complexity, the two languages seem very similar to humans. Our experiments show that, e.g., Transformers can learn context-free and context-sensitive languages equally well. However, they seem to struggle with regular languages, which are deemed to be much simpler by the Chomsky hierarchy.

Based on this and similar observations, we conclude that language complexity, as the Chomsky hierarchy defines it, is not a suitable predictor for how well a neural network can learn a language. To guide architecture choices in language models, we need better tools to measure the complexity of the language task we want to learn.

It’s an open question what these could look like. Presumably, we’ll need to find different complexity measures for different model architectures that consider their specific inductive biases.

What’s next?

Understanding how and why LLMs are so successful paves the way to more data-, cost- and energy efficiency. If you want to dive deeper into this topic, our position paper’s “Background” section is full of references, and we discuss numerous concrete research questions.

If you’re new to the field, I particularly recommend Same Pre-training Loss, Better Downstream: Implicit Bias Matters for Language Models (2023) by Liu et al., which nicely demonstrates the shortcomings of current evaluation practices based on the test loss. I also encourage you to check out SGD on Neural Networks Learns Functions of Increasing Complexity (2023) by Nakkiran et al. to understand more deeply how using stochastic gradient descent affects what functions neural networks learn.