MOSAIC: Multi-Objective Search for Allocating Inference Cache Across LLM Decoder Layers
MOSAIC is a multi-objective neural architecture search framework that optimizes per-layer KV-cache budgets for long-context LLM inference, consistently outperforming uniform allocation across various models and eviction methods by achieving higher task scores with significantly reduced memory usage.
Original paper dedicated to the public domain under CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). This is an AI-generated explanation of the paper below. It is not written by the authors. For technical accuracy, refer to the original paper. Read full disclaimer
Large language models are the engines behind many of today's most advanced artificial intelligence tools, capable of reading vast documents or holding long conversations. However, these models face a physical limitation when they try to process long sequences of text: they must store a growing amount of temporary information in their computer memory to keep track of what they have already read. This storage area, known as a key-value cache, expands linearly with the length of the text. For a model trying to understand a book-length document, this memory requirement can become so massive that it exceeds the capacity of standard computer hardware, forcing the system to either stop or discard information.
To solve this, researchers have developed methods to be selective about what information to keep. These techniques act like a filter, deciding which parts of the text are important enough to retain in memory and which can be safely forgotten. Until now, these filters have applied the same strict rules to every part of the model's internal processing. The assumption was that every layer of the model's brain needed the same amount of memory to do its job. This new research challenges that assumption, proposing that different layers actually need different amounts of space, and that giving them equal shares is a waste of resources.
A team of researchers has introduced a new method called MOSAIC, which automatically figures out the perfect amount of memory to assign to each layer of a large language model. Instead of treating the model as a uniform block where every section gets an equal slice of the memory pie, MOSAIC treats the layers as individuals with unique needs. The system uses a sophisticated search process to test thousands of different ways to distribute memory. It looks for a specific pattern where some layers get more space while others get less, all while staying within a total memory limit. The goal is to find the arrangement that allows the model to answer questions or summarize text most accurately without running out of memory.
The researchers tested this approach on two popular large language models and compared it against the standard method of giving every layer the same amount of memory. They evaluated the models on a variety of difficult tasks, including complex reading comprehension, code generation, and summarizing long documents. The results showed that the new method consistently outperformed the old one in the vast majority of cases, beating the uniform allocation in 146 of 154 test settings. In some instances, the improvement was dramatic, with the model gaining nearly fifty points on a standardized test for long-context understanding. Perhaps more importantly, the new method allowed the models to achieve the same high performance using up to four times less memory than before. This means the models could handle longer texts or run on smaller, less expensive computers.
One of the most surprising findings was that the best memory distribution was not random. The search revealed a consistent pattern: a small group of layers in the middle of the model consistently needed more memory than the layers at the beginning or the end. The layers at the start and finish could function well with less space, while the middle layers acted as the critical bottleneck where information needed to be preserved. By shifting memory from the layers that didn't need it to the ones that did, the system became much more efficient. This discovery suggests that the internal structure of these models is not uniform; different parts of the network play different roles and require different resources to function at their best.
The researchers also found that this optimal pattern could be discovered once and then applied to many different memory limits. They did not need to run a new, expensive search for every single memory size a user might want. Instead, they found a "winning shape" for the memory distribution and simply scaled it up or down. This means that a single, one-time search could provide the solution for any future memory constraint, making the method practical for real-world use. The system works by learning the ideal balance offline, before the model is ever used to answer a question, so it adds no extra delay during actual use.
When compared to other recent attempts to solve this problem, this new approach explored a vastly larger space of possibilities. While other methods tested only a few variations at a time, this system evaluated a tiny fraction of the total possible configurations—scanning a search space billions of times larger than previous methods while evaluating only a few hundred candidates to find the best solution. Despite looking at so many more options, it required fewer computer hours to find the best solution. The study confirms that the old idea of treating every layer the same is inefficient. By recognizing that some parts of the model are more critical than others and allocating resources accordingly, we can make these powerful tools faster, cheaper, and capable of understanding much longer contexts than previously thought possible.
Drowning in papers in your field?
Get daily digests of the most novel papers matching your research keywords — with technical summaries, in your language.