PaceLLM: Brain-Inspired Large Language Models for Long-Context Understanding

While Large Language Models (LLMs) demonstrate strong performance across domains, their long-context capabilities are limited by transient neural activations causing information decay and unstructured feed-forward network (FFN) weights leading to semantic fragmentation. Inspired by the brain’s working memory and cortical modularity, we propose PaceLLM, featuring two innovations: (1) a Persistent Activity (PA) Mechanism that mimics prefrontal cortex (PFC) neurons’ persistent firing by introducing an activation-level memory bank to dynamically retrieve, reuse, and update critical FFN states, addressing contextual decay; and (2) Cortical Expert (CE) Clustering that emulates task-adaptive neural specialization to reorganize FFN weights into semantic modules, establishing cross-token dependencies and mitigating fragmentation. Extensive evaluations show that PaceLLM achieves 6% improvement on LongBench’s Multi-document QA and 12.5-17.5% performance gains on Infinite-Bench tasks, while extending measurable context length to 200K tokens in Needle-In-A-Haystack (NIAH) tests. This work pioneers brain-inspired LLM optimization and is complementary to other works. Besides, it can be generalized to any model and enhance their long-context performance and interpretability without structural overhauls.

💡 Research Summary

PaceLLM introduces a brain‑inspired architecture to tackle the long‑context bottleneck that plagues modern large language models (LLMs). While prior work has focused on augmenting attention mechanisms, compressing inputs, or attaching external memory modules, this paper identifies two overlooked internal limitations: (1) transient activations in feed‑forward networks (FFNs) that cause information decay over long sequences, and (2) unstructured FFN weights that fragment semantics across tokens. Drawing inspiration from the prefrontal cortex’s persistent firing (working memory) and the cerebral cortex’s functional modularity, the authors propose two complementary mechanisms: Persistent Activity (PA) via an Activation Memory Bank (AMB) and Cortical Expert (CE) clustering of FFN parameters.

Persistent Activity (PA) – Activation Memory Bank
Each selected FFN layer is equipped with a memory bank M = {K, V, u}, where K and V store keys and values derived from past intermediate activations, and u tracks usage frequency. When processing a new token, the current activation Xc is compared to stored keys using cosine similarity. The top‑k most similar entries (positive memory) and the bottom‑k′ least similar entries (negative memory) are retrieved. The final output is a weighted combination of the current activation, the positive memory mean µpos, and the negative memory mean µneg, modulated by a similarity‑dependent gating function with thresholds θlow and θhigh. Memory updates follow a similarity‑aware policy: high similarity → only increment usage; medium similarity → merge current activation into the stored slot; low similarity → replace the least‑used slot (LRU). This design mimics the brain’s working memory, allowing relevant activation traces to persist and be re‑used across distant parts of a document, thereby mitigating the decay of information that typically occurs in standard transformer pipelines.

Cortical Expert (CE) Clustering
The second innovation restructures the FFN weight matrices into semantically coherent “expert” modules. The rows of the first FFN projection matrix W1 (dff × dmodel) are normalized and clustered using K‑meansConstrained, enforcing equal‑size clusters. Each cluster corresponds to a cortical expert, a group of neurons that respond to similar patterns. After clustering, a permutation π is derived that groups rows (and the corresponding columns of the second projection matrix W2) by expert. The weight matrices are then reordered so that each expert occupies a contiguous block in both W1 and W2. This reorganization does not alter the learned parameters; it simply changes their layout, enabling the model to process inputs with a modular, specialist‑like structure. The authors argue that this reduces semantic fragmentation because tokens that belong to the same semantic domain are more likely to be processed by the same expert block, preserving coherence over long spans.

Experimental Setup
The authors evaluate PaceLLM on two popular long‑context benchmarks: LongBench (covering single‑question answering, multi‑question answering, summarization, few‑shot learning, and coding) and ∞‑Bench (encompassing English dialogue and multi‑choice tasks). They also conduct a Needle‑In‑A‑Haystack (NIAH) test to probe the maximum usable context length. Base models are Qwen‑2‑7B‑Instruct and Llama‑2‑7B‑chat. Experiments are performed in a training‑free setting (no weight updates) and in a fine‑tuning scenario where the base model is further trained on the target tasks. Hyper‑parameters for the memory bank (M = 512 slots, top‑k = 8, k′ = 4, θlow = 0.4, θhigh = 0.8) and the number of experts (K = 64–128) are selected via a small validation sweep.

Results
In the training‑free regime, adding CE alone yields modest gains (≈0.2–0.4 % absolute) across most LongBench tasks, while PA alone provides slightly larger improvements (≈0.3–0.5 %). When both mechanisms are combined, the gains are additive: Qwen‑2‑7B‑Instruct achieves a 6 % absolute increase on the Multi‑document QA sub‑task (from 38.49 to ~44.5) and modest improvements on the other four sub‑tasks. On ∞‑Bench, the combined model improves English Dialogue by 12.5 % and English Multi‑Choice by 17.5 % relative to the vanilla baseline, indicating that the modular expert layout helps maintain coherence in dialogue and reasoning over many options. The NIAH test demonstrates that PaceLLM can reliably retrieve and answer queries from contexts up to 200 K tokens, surpassing the 128 K token limit of the strongest prior method (Activation Beacon). Memory overhead remains modest because the AMB is dynamically pruned and only a subset of FFN layers are equipped with it.

Ablation and Analysis
A detailed ablation shows that the size of the memory bank exhibits diminishing returns beyond 512 slots; larger banks increase latency without proportional accuracy gains. The top‑k / bottom‑k′ trade‑off influences the balance between reuse and diversity: too small a top‑k reduces the benefit of persistent activation, while too large a bottom‑k′ introduces noise. For CE, the number of experts K must be tuned to model capacity: overly fine‑grained clustering (K > 256) fragments the network and harms performance, whereas too coarse a clustering (K < 32) fails to capture semantic specialization. The authors also provide visualizations of expert activation patterns, showing that certain experts consistently fire for domain‑specific terminology (e.g., medical, legal) across long documents, supporting the claim of functional modularity.

Discussion
PaceLLM’s contributions are twofold. First, by storing intermediate activations rather than raw token embeddings, the AMB operates at a finer granularity, enabling the model to “remember” abstract feature patterns that recur across a document. This is more efficient than external retrieval‑augmented generation, which requires separate indexing and incurs latency. Second, the CE clustering offers a lightweight way to impose a modular structure on otherwise monolithic FFNs, improving interpretability (experts can be inspected) and potentially facilitating future conditional computation (e.g., routing only relevant experts). The approach is model‑agnostic, requiring no retraining, and can be combined with other long‑context strategies such as sliding‑window attention or KV‑cache compression.

Limitations and Future Work
The current implementation equips only a subset of FFN layers with the AMB, leaving potential gains untapped in deeper layers. Memory management, while efficient, still adds overhead that may be prohibitive for extremely resource‑constrained environments. The method is evaluated solely on text; extending PA and CE to multimodal transformers (vision‑language, speech) remains an open question. Future research directions include (1) integrating PA/CE with advanced sparse‑attention mechanisms, (2) exploring dynamic expert selection based on input content, and (3) applying the persistent activation concept to episodic memory modules for continual learning scenarios.

Conclusion
PaceLLM demonstrates that brain‑inspired mechanisms—persistent activity and cortical modularity—can be translated into practical, training‑free enhancements for LLMs. By addressing internal activation decay and semantic fragmentation, the model achieves consistent performance gains on long‑context benchmarks and pushes the usable context window to 200 K tokens. The work opens a promising avenue for neuro‑inspired architecture design, offering both accuracy improvements and greater interpretability without extensive architectural overhauls.