Optimizing AI Model Performance through Single-Tag Continuous Improvement in StreamingLLM

Large language models (LLMs) such as OpenAI's ChatGPT, Meta's Llama 2, and Anthropic's Claude 2 have been a focus of attention in Silicon Valley and the corporate tech world. However, they all face some common challenges.

One of the challenges is maintaining high-quality performance consistently when engaging in a single conversation with users. This means that the replies provided by LLMs in the middle and end of a conversation should be as helpful, fast, and relevant as they were at the beginning, regardless of the length or number of exchanges in the conversation.

This is because LLMs are pretrained based on fixed-length data blocks or sequences, with Llama 2 and many other leading LLMs having a sequence length of 4,000 tokens per sequence.

Once the number of tokens from the user exceeds this length, even if input is provided across multiple different prompts, the performance of LLMs deteriorates, meaning the quality of the replies decreases. This is unacceptable for businesses that want LLMs to assist customers or employees in an open-ended manner.

A recent paper published by researchers from Meta, MIT, and CMU introduces a simple method to help LLMs maintain their performance even when the total length of user prompts exceeds the training length handled by LLMs in a single processing step.

The researchers developed a new framework called "StreamingLLM" for training and deploying LLMs, which provides important insights for other AI researchers and businesses looking to utilize LLMs for their operations.

The Problem Addressed by StreamingLLM

Anyone who has interacted with human customer support agents or IT personnel within a company knows that resolving issues often requires lengthy conversations and multiple message exchanges.

Whether you are a customer or an employee, you expect the person assisting you to be responsive, knowledgeable, and helpful throughout the entire communication process. It can be frustrating if, at some point in the conversation, after you have spent time and effort explaining the problem, your helper starts responding with single-word replies, becomes slower, or fails to provide the information you need.

While some individuals may experience this problem due to distractions, lack of motivation, or fatigue from the conversation, it is a common issue for LLMs because their performance deteriorates once the conversation exceeds the "context window" length of LLMs. The context window of an LLM is the maximum number of tokens it can process in a single response, and this value is used during pretraining. Despite most LLMs being designed to handle open-ended dialogues that can span multiple lines, this problem still exists.

Even if each line of text fits within the context window of an LLM, most LLMs have an upper limit that restricts the amount of text you can input in a single message. However, the accumulation of multiple messages in a single conversation can cause the number of tokens to exceed the initial pretrained context window of LLMs, resulting in a performance decline beyond that point.

This is similar to conversing with a human customer support agent who suddenly becomes less intelligent and attentive when the number of words you input across several sentences reaches an unknown limit.

The researchers behind the StreamingLLM framework summarized this problem in their paper: "For example, an ideal chatbot assistant should be able to stably handle the content of a conversation as long as a day. However, extending LLMs to longer sequence lengths than those they were pretrained on is highly challenging."

While it is possible to extend the token sequence length of pretrained LLMs, and many researchers have done so, it is impossible to predict how long a conversation with a given user will last.

So, when an LLM has a fixed context window length for pretraining, regardless of how long it is, how can its performance be maintained when the conversation spans multiple messages that exceed this length?

The Solution Developed by the Researchers

The researchers developed an innovative solution to maintain the performance of LLMs when the amount of information in a conversation exceeds the number of tokens used in the pretrained sequence.

The researchers found that LLMs pay more attention to the initial tokens in a conversation or training. They wrote, "Surprisingly, a large amount of attention scores are allocated to the initial tokens." Why is this?

They explained, "Due to the continuity nature of autoregressive language modeling, the initial tokens are visible to all subsequent tokens, while subsequent tokens are only visible to a limited set of subsequent tokens." They continued, "Therefore, initial tokens are easier to train as attention receivers, capturing unnecessary attention."

In other words, regardless of what content you initially provide in a conversation with an LLM, it will be used in subsequent prompts and replies, but the content you provide later may not be what the LLM chooses to focus on or reference.

However, the researchers found that if users provide some initial tokens in the subsequent conversation with the LLM, it is enough to restore the performance of the LLM to near-peak levels in the subsequent replies.

Remember the human customer support analogy mentioned earlier? Imagine if by repeating the four magical words you said at the beginning of the conversation later on, you could suddenly elicit high-quality answers from them in future conversations.

The researchers appropriately referred to these initial tokens that capture most of the LLM's attention as "attention receivers" and noted that for most LLMs, "introducing four initial tokens as attention receivers is sufficient to restore the performance of LLM...adding only one or two is not enough for full recovery."

By reintroducing attention receiver tokens in each subsequent prompt from the user, the researchers were able to maintain the performance of major models like LLama 2 and Falcon 40B, which consist of prompts with 4 million tokens (a 1,000-fold increase compared to the initial 4,000-token context window), and improve the speed of their subsequent replies by 22.2 times.

In other words, StreamingLLM "enables LLMs with a limited attention window to handle text of unlimited length without fine-tuning." It is important to note that this "unlimited" length of text still needs to be delivered to LLMs based on their context window size limitation. However, this means that LLMs can engage in endless conversations with individuals while maintaining their performance (in theory).

One Token to Rule Them All

The researchers further hypothesized and demonstrated that in reality, only adding a special token as an "attention receiver" for LLMs is sufficient to maintain their high performance. By reintroducing this token manually or automatically (behind the scenes of LLMs faced by users or employees) in subsequent conversations, the performance of LLMs can be consistently maintained at a high level.

The researchers explained, "Introducing a receiver token is highly effective for stabilizing the attention mechanism. Pairing this receiver token with the most recent token is sufficient to stabilize the model's performance...Given these findings, we suggest using receiver tokens in all samples for training future LLMs to optimize streaming deployment."

When asked about the specific data that should be used as attention receiver tokens, Guangxuan Xiao, one of the authors of the paper from MIT, wrote in an email, "'Attention receiver' can be any initial token; what matters more is their position rather than their semantics...It is not a specific word or concept; even a token without semantic meaning (e.g., a newline character "\n") can be effectively used."

As for the intended use cases for StreamingLLM, Xiao stated, "We designed StreamingLLM for continuous applications such as multi-turn dialogues. It is perfect for use cases where the model needs to persist and not rely too heavily on past data. Daily Assistant LLM is a typical example. With our method, the model can persist and gather information from recent interactions without frequent cache refreshing."

However, the researchers also explicitly acknowledged the limitations of their work and were careful to emphasize that StreamingLLM does not extend the context window of LLMs, contrary to some hype on platforms like X (formerly Twitter) about their work. It also does not guarantee that LLMs will remember everything said throughout a conversation.

Xiao stated, "In fact, we neither extend the context window of LLMs nor improve their long-term memory."