When a human-AI conversation involves many rounds of continuous dialogue, the powerful large language machine-learning models that drive chatbots like ChatGPT sometimes begin to collapse, causing the bots’ performance to rapidly deteriorate.
A team of researchers from MIT and elsewhere has pinpointed a surprising reason behind this problem and developed a straightforward solution that permits a chatbot to keep up a nonstop conversation without crashing or slowing down.
Their method involves a tweak to the key-value cache (which is sort of a conversation memory) on the core of many large language models. In some methods, when this cache needs to carry more information than it has capability for, the primary pieces of knowledge are bumped out. This will cause the model to fail.
By ensuring that these first few data points remain in memory, the researchers’ method allows a chatbot to maintain chatting irrespective of how long the conversation goes.
The tactic, called StreamingLLM, enables a model to stay efficient even when a conversation stretches on for greater than 4 million words. When put next to a different method that avoids crashing by continually recomputing a part of the past conversations, StreamingLLM performed greater than 22 times faster.
This might allow a chatbot to conduct long conversations throughout the workday without having to be continually rebooted, enabling efficient AI assistants for tasks like copywriting, editing, or generating code.
“Now, with this method, we are able to persistently deploy these large language models. By making a chatbot that we are able to at all times chat with, and that may at all times reply to us based on our recent conversations, we could use these chatbots in some latest applications,” says Guangxuan Xiao, an electrical engineering and computer science (EECS) graduate student and lead writer of a paper on StreamingLLM.
Xiao’s co-authors include his advisor, Song Han, an associate professor in EECS, a member of the MIT-IBM Watson AI Lab, and a distinguished scientist of NVIDIA; in addition to Yuandong Tian, a research scientist at Meta AI; Beidi Chen, an assistant professor at Carnegie Mellon University; and senior writer Mike Lewis, a research scientist at Meta AI. The work will probably be presented on the International Conference on Learning Representations.
A puzzling phenomenon
Large language models encode data, like words in a user query, into representations called tokens. Many models employ what’s often called an attention mechanism that uses these tokens to generate latest text.
Typically, an AI chatbot writes latest text based on text it has just seen, so it stores recent tokens in memory, called a KV Cache, to make use of later. The eye mechanism builds a grid that features all tokens within the cache, an “attention map” that maps out how strongly each token, or word, relates to one another token.
Understanding these relationships is one feature that permits large language models to generate human-like text.
But when the cache gets very large, the eye map can change into much more massive, which slows down computation.
Also, if encoding content requires more tokens than the cache can hold, the model’s performance drops. As an example, one popular model can store 4,096 tokens, yet there are about 10,000 tokens in an educational paper.
To get around these problems, researchers employ a “sliding cache” that bumps out the oldest tokens so as to add latest tokens. Nonetheless, the model’s performance often plummets as soon as that first token is evicted, rapidly reducing the standard of newly generated words.
On this latest paper, researchers realized that in the event that they keep the primary token within the sliding cache, the model will maintain its performance even when the cache size is exceeded.
But this didn’t make any sense. The primary word in a novel likely has nothing to do with the last word, so why would the primary word be so vital for the model to generate the most recent word?
Of their latest paper, the researchers also uncovered the reason behind this phenomenon.
Attention sinks
Some models use a Softmax operation of their attention mechanism, which assigns a rating to every token that represents how much it relates to one another token. The Softmax operation requires all attention scores to sum as much as 1. Since most tokens aren’t strongly related, their attention scores are very low. The model dumps any remaining attention rating in the primary token.
The researchers call this primary token an “attention sink.”
“We want an attention sink, and the model decides to make use of the primary token as the eye sink since it is globally visible — every other token can see it. We found that we should always keep the eye sink within the cache to keep up the model dynamics,” Han says.
In constructing StreamingLLM, the researchers discovered that having 4 attention sink tokens at first of the sliding cache results in optimal performance.
In addition they found that the positional encoding of every token must stay the identical, at the same time as latest tokens are added and others are bumped out. If token 5 is bumped out, token 6 must stay encoded as 6, regardless that it’s now the fifth token within the cache.
By combining these two ideas, they enabled StreamingLLM to keep up a continuous conversation while outperforming a well-liked method that uses recomputation.
As an example, when the cache has 256 tokens, the recomputation method takes 63 milliseconds to decode a latest token, while StreamingLLM takes 31 milliseconds. Nonetheless, if the cache size grows to 4,096 tokens, recomputation requires 1,411 milliseconds for a latest token, while StreamingLLM needs just 65 milliseconds.
“The revolutionary approach of StreamingLLM, centered around the eye sink mechanism, ensures stable memory usage and performance, even when processing texts as much as 4 million tokens in length,” says Yang You, a presidential young professor of computer science on the National University of Singapore, who was not involved with this work. “This capability is just not just impressive; it’s transformative, enabling StreamingLLM to be applied across a wide selection of AI applications. The performance and flexibility of StreamingLLM mark it as a highly promising technology, poised to revolutionize how we approach AI-driven generation applications.”
Tianqi Chen, an assistant professor within the machine learning and computer science departments at Carnegie Mellon University who also was not involved with this research, agreed, saying “Streaming LLM enables the graceful extension of the conversation length of huge language models. We now have been using it to enable the deployment of Mistral models on iPhones with great success.”
The researchers also explored using attention sinks during model training by prepending several placeholder tokens in all training samples.
They found that training with attention sinks allowed a model to keep up performance with just one attention sink in its cache, reasonably than the 4 which can be often required to stabilize a pretrained model’s performance.
But while StreamingLLM enables a model to conduct a continuous conversation, the model cannot remember words that aren’t stored within the cache. In the long run, the researchers plan to focus on this limitation by investigating methods to retrieve tokens which have been evicted or enable the model to memorize previous conversations.
StreamingLLM has been incorporated into NVIDIA’s large language model optimization library, TensorRT-LLM.
This work is funded, partly, by the MIT-IBM Watson AI Lab, the MIT Science Hub, and the U.S. National Science Foundation.
