Memory in LLMs: Weights and Activations - Jack Morris, Cornell

The Problem: Knowledge Gaps in LLMs

Current Large Language Models (LLMs) like ChatGPT, despite their impressive general knowledge, have significant limitations. They are constrained by knowledge cut-offs, rendering them unaware of recent events unless connected to web search. They struggle with niche, long-tail tasks, such as optimizing code for specific hardware (e.g., AMD GPUs), due to insufficient representation in their vast training data. Furthermore, they cannot access private, proprietary, or real-time information, making them unsuitable for tasks involving company-specific data, personal preferences, or confidential negotiations.

Three Paradigms for Injecting Knowledge

The central question is how to effectively inject new, specific knowledge into an LLM. The speaker outlines three primary methods:

1. Full Context: Cramming all available information into the model’s context window.
2. Retrieval-Augmented Generation (RAG): Retrieving the most relevant documents from a large corpus and placing them in the context window.
3. Training into Weights: Directly modifying the model’s parameters to permanently store new information.

The talk argues for the superiority and untapped potential of the third approach.

The Limitations of Context and RAG

Full Context

While simple, relying solely on the context window is impractical for substantial knowledge bases.

• Cost and Latency: Performance degrades dramatically with larger contexts. An example given is a drop from 10,000 tokens/second with a 1k context to just 130 tokens/second with a 128k context.
• Architectural Bottleneck: The quadratic complexity of the self-attention mechanism in Transformers makes it computationally infeasible to scale to millions of tokens.
• Reasoning Failure (Context Drowning): Model performance degrades as the context window fills with information, even if it’s relevant. A Chroma report showed that beyond 10,000 tokens, models effectively fail to solve the task, a phenomenon known as the “needle in a haystack” problem.
• Efficient Architectures are Not a Panacea: Newer architectures like Mamba or linear attention models offer efficiency gains but often at the cost of performance, with some state-of-the-art models reverting to standard quadratic attention for its superior quality.

Retrieval-Augmented Generation (RAG)

RAG, powered by vector databases, is the current industry standard for enterprise LLM applications.

• Prevalence: It is described as “the file system of LLMs today” and is lauded for its ease of implementation.
• Fundamental Flaws: Despite its popularity, the audience and speaker agree it is far from a perfect solution.
  • Security Risk: Embeddings are not secure; the speaker’s research demonstrates the ability to reconstruct original text from its vector representation with over 90% accuracy.
  • Lack of Adaptability: Standard, universal embeddings cluster related but distinct concepts (e.g., Visa and MasterCard documents) too closely, harming search relevance in specialized domains.
  • Limited Reasoning: RAG struggles with queries that require synthesizing information across multiple documents or understanding implied relationships that aren’t explicitly stated.

The Future: Training Knowledge into Weights

The most promising frontier is to directly encode knowledge into the model’s parameters, effectively swapping out useless trivia for valuable, domain-specific information.

The Naive Approach Fails

Simply performing next-token prediction on a specific document (e.g., a 3M financial report) leads to catastrophic forgetting. The model memorizes the document perfectly but loses its general reasoning abilities, becoming unable to perform even simple creative tasks related to the new content.

The Solution: Synthetic Data and Parameter-Efficient Tuning

The key is a two-pronged approach:

1. Synthetic Data Generation: Instead of training on a small, raw dataset, the model should be trained on a large, diverse, synthetically generated dataset that represents the core facts and relationships from the original data. This technique, demonstrated in works like “Synthetic Continued Pre-training,” allows a model to learn a domain deeply without overfitting and can even outperform larger models like GPT-4 on specialized tasks.
2. Parameter-Efficient Fine-Tuning (PEFT): To avoid catastrophic forgetting, updates should be made to a small subset of the model’s parameters, not the entire model. Several architectures enable this:
   • LoRA (Low-Rank Adaptation): Trains small, low-rank matrices to adapt the model.
   • Prefix Tuning: Modifies the KV cache.
   • Memory Layers: Adds a differentiable lookup table as a new expert layer.

The evidence for which PEFT method is best is currently conflicting, but the underlying principle is consistent: make small, targeted updates to preserve the model’s core capabilities while injecting new knowledge. The speaker’s research shows that for RL-based tasks, an incredibly small number of parameters (as few as one) can yield significant performance gains.

Conclusion: A Hybrid Future

The path forward involves a hybrid strategy. RAG will remain useful for handling highly dynamic or infrequently accessed data. However, it will be augmented by a new paradigm of periodically and efficiently training core, stable knowledge directly into the model’s weights. This approach promises to create models that are not only more knowledgeable and capable but also cheaper to run at inference time, paving the way for true personalization at scale.

• Tags:
• LLM
• Synthetic Data
• RAG
• Transformers
• Fine-Tuning
• Feature
• LoRA
• Memory
• Model Architecture
• Context Window
• Parameter-Efficient Fine-Tuning
• PEFT
• Catastrophic Forgetting