Apple - Pretraining with hierarchical memories - separating long-tail and common knowledge

Pretraining with hierarchical memories: separating long-tail and common knowledge

Apple's Secret Weapon for On-Device AI: An 'External Brain'

Pretraining with hierarchical memories: separating long-tail and common knowledge

Introduction: Why is running powerful AI on an iPhone so difficult?

We all know that today's Large Language Models (LLMs) like GPT-4 are incredibly powerful, largely due to their massive number of parameters—hundreds of billions, or even trillions. These parameters are like the neurons in a model's brain, storing the vast world knowledge and complex reasoning abilities learned from the internet.

But this creates immense challenges:

1. Sheer Size: Such enormous models require tens or even hundreds of gigabytes of storage and consume staggering amounts of memory during inference. This is simply unthinkable for resource-constrained edge devices like an iPhone or a smartwatch.
2. Inefficiency: When responding to a simple question, the model has to engage all its parameters, even though the vast majority of that knowledge is irrelevant. It's like reading the entire Oxford English Dictionary from cover to cover just to look up a single word—a colossal waste of effort.

Therefore, figuring out how to make LLMs lightweight and efficient enough to run locally on personal devices, all while retaining their power, has become a central issue in the AI field. This is especially true for companies like Apple, which prioritize on-device experiences and user privacy. This paper proposes a highly creative solution to this very problem.

The Core Idea: Building an 'External Brain' for AI

Imagine that instead of trying to build a single, omniscient "super-brain," we adopt a smarter, more collaborative division of labor:

1. A lean "generalist brain" (the Anchor Model): This is a small language model with a low parameter count that runs incredibly fast. It isn't responsible for memorizing a sea of trivial facts. Instead, it focuses on core common knowledge (like grammar and basic facts) and reasoning abilities. Think of it as a smart generalist, skilled at thinking, analyzing, and organizing language.
2. A vast "external knowledge base" (the Hierarchical Memory Bank): This is a huge, meticulously organized parameterized memory bank. It's specifically designed to store concrete, precise, but non-universal long-tail knowledge (e.g., the atomic number of a specific chemical element, the director of an obscure film). It acts like a well-organized, giant digital library.
3. A lightning-fast "librarian" (the Memory Retriever): When a user asks a question, this "librarian" instantly analyzes the type of knowledge needed and, in milliseconds, dynamically retrieves a small, highly relevant knowledge module from the "external knowledge base."

Finally, this compact "generalist brain" combines with the freshly retrieved "knowledge module" to generate a precise answer.

Figure: Diagram of the Memory-Augmented Transformer Architecture

As shown in the figure above, when a query (context ) is input, the Memory Retriever selects relevant parameters from the vast memory bank. These are then combined with the Anchor Model to generate the final answer.

The genius of this architecture is that at any given moment, the only components that need to be loaded into active memory are the small Anchor Model and a tiny snippet of temporarily summoned "memory parameters." This makes the system's operational cost extremely low.

A Technical Deep Dive: How Does This 'External Brain' Work So Efficiently?

1. How is knowledge organized? A file system-like hierarchical structure.

To enable the "librarian" to find knowledge quickly, this massive memory bank must be impeccably organized. The researchers used the k-means clustering algorithm to hierarchically cluster a massive corpus of pre-training documents.

You can think of it like organizing files on your computer:

• Layer 1 (L1): 16 top-level categories, such as "Science," "History," and "Art."
• Layer 2 (L2): Each major category is further divided into 16 sub-categories. For example, "Science" might be split into "Physics," "Chemistry," and "Biology."
• Layers 3 (L3) & 4 (L4): This process continues, forming a massive knowledge tree with  (over 65,000) leaf nodes.

When a query comes in, the retriever only needs to traverse this tree four steps down from the root to pinpoint the most relevant knowledge cluster, resulting in extremely low retrieval complexity.

Deeper memories signify more specialized knowledge. Experiments confirmed that this hierarchical structure is highly effective.

Figure: Relationship Between Memory Hierarchy and Performance

The figure above shows that for a fixed amount of retrieved memory, the deeper the memory layer (e.g., L4), the higher the model's accuracy, as it provides more specific and relevant knowledge.

2. How is knowledge 'injected'? The highly efficient FFN-Memories.

Once a knowledge module is found, how is it fused with the Anchor Model? The research team tested several methods and found that FFN-Memories performed best.

To understand this, you just need to know that in each layer of a Transformer model (the foundational architecture of modern LLMs), there is a component called a Feed-Forward Network (FFN). You can think of it as an "information processing workshop." This paper's method essentially involves dynamically adding a dedicated "knowledge production line" to this workshop when processing specific knowledge.

Figure: Performance Comparison of Different Memory Types

The chart above clearly demonstrates that when handling Specific-Knowledge tasks, FFN-Memories (the blue line) significantly outperform other methods (like LoRA or KV memories) across all parameter scales.

3. How is knowledge stored? A deployment strategy optimized for mobile devices.

This is one of the most exciting highlights of the research, perfectly embodying its design philosophy for edge devices. The researchers leveraged the properties of hierarchical memory to propose an exceptionally clever hardware deployment strategy:

• Shallow memories (L1, L2): These are coarser-grained, potentially more frequently used knowledge clusters. Their total size is relatively small, allowing them to be stored in fast but expensive RAM.
• Deep memories (L3, L4): This is where the vast, highly specific long-tail knowledge resides. They can be "offloaded" to the slower but much larger flash storage (i.e., the phone's internal drive).

Figure: Deployment Efficiency Advantage of Hierarchical Memory

The left chart shows that loading a hierarchical memory chunk (partly from RAM, partly from flash) takes only 38ms. In contrast, loading a flat memory of the same size entirely from flash would take 198ms—more than 5 times slower. This is critically important for on-device AI that requires near-instantaneous responses.

Showcasing the Results: Small Models with Big Power

The theory sounds great, but how does it perform in practice? The experimental data provides some astonishing answers.

1. Performance Leap: Punching Above Its Weight

A 160M-parameter Anchor Model, when dynamically calling about 18M in memory parameters (for a total runtime size of 178M), can match or even surpass the performance of a conventionally trained 410M-parameter model on specific-knowledge tasks.

Figure: Performance Across Different Memory Configurations

This chart illustrates the relationship between performance and two dimensions: runtime fetched memory size (X-axis) and total memory bank size (color/shape). The conclusion is clear: bigger is better for both. This demonstrates the method's excellent scalability.

2. Knowledge Controllability: Opening the Door to Privacy and Editing

Because knowledge is stored in discrete, addressable memory blocks, we can perform fine-grained operations on it. The researchers conducted an interesting experiment: when answering questions about the atomic numbers of elements, they maliciously "blocked" or "masked" the specific memory block containing that knowledge.

Figure: Results of the Memory Blocking Experiment

As the figure shows, the model's accuracy plummeted from 70% to 20%. This provides strong evidence that knowledge is indeed localized within specific memory blocks. This also implies that in the future, we could:Protect privacy: Store a user's private data in memory blocks that can be removed at any time.Edit knowledge: When a fact becomes outdated or is found to be incorrect, we can simply update or replace the corresponding memory block without having to retrain the entire massive model.

3. Universality: 'Plug-and-Play' for Existing Models

This technology isn't just for training models from scratch. The research shows that it can be applied post-hoc as an "enhancement patch" to already-trained open-source models.

Table: Results of Post-Hoc Memory Addition to Open-Source Models

Table 2 Summary

• After post-hoc addition of about 10% in memory parameters, popular models like Gemma 3 270M, Qwen 2.5 0.5B, and Llama 3.2 1B all showed significant performance gains on specific-knowledge (Avg-SK) tasks.
• This demonstrates the method's powerful universality, allowing it to serve as a plug-and-play enhancement module that can revitalize the existing ecosystem of models.

Head-to-Head: How Is This Different from the Popular RAG?

You might be thinking this sounds a bit like Retrieval-Augmented Generation (RAG). While there are similarities in the "retrieve + generate" philosophy, there's a fundamental difference:

• RAG: When a user asks a question, RAG retrieves relevant raw text passages from a text database (like Wikipedia) and feeds them, along with the question, to an LLM. The LLM then uses this context to formulate an answer. It's more like giving the model an "open-book exam."
• This paper's method (Learned Memory): It doesn't retrieve raw text. Instead, it retrieves parameterized knowledge representations that have been learned and compressed by the model. This is more akin to the model activating a specialized, already internalized "knowledge region" of its own brain when needed.

Table: Comparison with Vanilla RAG

Table 3 Summary

• Using RAG with high-quality data from Wikipedia can improve specific-knowledge (Avg-SK) performance, but it offers no help for common-knowledge (Avg-CK) tasks.
• In contrast, this paper's memory method, with far fewer runtime parameters (178M vs. 410M), achieves superior performance on both common and specific knowledge tasks, and does so with lower computational overhead (FLOPs).

Conclusion: What Does This Mean for Apple and the AI Industry?

The strategic value of this research extends far beyond its technical merits:

1. The 'Holy Grail' of On-Device AI: This technology is a critical step toward achieving truly powerful AI that can run smoothly on devices like the iPhone. It perfectly aligns with Apple's long-standing strategy of on-device computation and user privacy protection. In the future, Siri might be able to answer all sorts of complex questions without ever needing to contact the cloud.
2. Challenging the 'Brute-Force' Approach: It offers an alternative path for the development of Edge AI beyond endlessly piling on more parameters—the path of efficient, elegant architectural design. This could change the game for the Edge AI chip and model markets, making smaller, highly efficient models far more valuable.

In summary, by cleverly separating an LLM's general reasoning abilities from its externalized long-tail knowledge memory, Apple's research not only solves the daunting challenge of deploying large models on edge devices but also brings enormous advantages in efficiency, privacy, and controllability. It paints a more sustainable and accessible blueprint for the future of artificial intelligence.