PaperlayerView exact revision on arXiv

Source-linked AI summary

Brainstacks: Cross-Domain Cognitive Capabilities via Frozen MoE-LoRA Stacks for Continual LLM Learning

Mohammad R. Abu Ayyash

arXiv:2604.01152v1cs.CLcs.AI

TL;DR

Existing LLM fine-tuning is monolithic, limiting modular continual learning and selective domain-capability composition. Brainstacks addresses this with frozen, composable MoE-LoRA stacks, null-space protection, residual stacking, and outcome-based sigmoid routing. The system recovers quality from ungated stack accumulation and reveals transferable cognitive primitives, including medical prompts routing to chat+math stacks 97% of the time despite no medical training data.

  • Problem

    Monolithic fine-tuning couples domain capabilities, making continual addition, independent updating or removal, and selective activation difficult.

  • Method

    Brainstacks trains frozen domain-specific MoE-LoRA stacks with residual stacking and null-space projection, then uses independent sigmoid gates for outcome-based cross-domain composition.

  • Results

    97% of medical prompts route to chat+math stacks despite zero medical data in those stacks, while routing recovers generation quality lost under ungated stack accumulation.

  • Takeaways & Limitations

    Domain stacks encode transferable cognitive primitives rather than only domain-specific knowledge, reframing fine-tuning as capability injection.

  • Takeaways & Limitations

    Brainstacks requires a pretrained base model with coherent hidden-state geometry; stacking on randomly initialized weights produces poor results.

Abstract

from arXiv · show

We present Brainstacks, a modular architecture for continual multi-domain fine-tuning of large language models that packages domain expertise as frozen adapter stacks composing additively on a shared frozen base at inference. Five interlocking components: (1) MoE-LoRA with Shazeer-style noisy top-2 routing across all seven transformer projections under QLoRA 4-bit quantization with rsLoRA scaling; (2) an inner loop performing residual boosting by freezing trained stacks and adding new ones; (3) an outer loop training sequential domain-specific stacks with curriculum-ordered dependencies; (4) null-space projection via randomized SVD constraining new stacks to subspaces orthogonal to prior directions, achieving zero forgetting in isolation; (5) an outcome-based sigmoid meta-router trained on empirically discovered domain-combination targets that selectively weights stacks, enabling cross-domain composition. Two boundary experiments: (6) PSN pretraining on a randomly initialized model; (7) per-domain RL (DPO/GRPO) validating compatibility with post-SFT alignment. Validated on TinyLlama-1.1B (4 domains, 9 stacks) and Gemma 3 12B IT (5 domains, 10 stacks), MoE-LoRA achieves 2.5x faster convergence than parameter-matched single LoRA, residual boosting breaks through the single-stack ceiling, and the routed system recovers generation quality destroyed by ungated stack accumulation. The central finding: the outcome-based router discovers that domain stacks encode transferable cognitive primitives (instruction-following clarity, numerical reasoning, procedural logic, chain-of-thought structure) rather than domain-specific knowledge, with medical prompts routing to chat+math stacks in 97% of cases despite zero medical data in those stacks.

1 Introduction

Brainstacks replaces monolithic continual fine-tuning with frozen, composable MoE-LoRA stacks that can be selectively activated and protected from interference. Its architecture combines residual boosting, continual stacking, null-space projection, and outcome-based routing to support modular multi-domain capabilities.

  • Architecture: Brainstacks packages each domain as a frozen residual MoE-LoRA stack that can be added, removed, or updated independently.New stacks train on top of frozen stacks while null-space constraints protect prior domain directions.
  • Training: The two-loop architecture combines inner-loop residual boosting with outer-loop continual domain stacking across all seven transformer projections.It uses Shazeer-style noisy routing and rsLoRA scaling while each new stack learns residual corrections over a shifted output landscape.
  • Anti-forgetting: Null-space gradient projection uses randomized SVD to enforce orthogonality between new and prior domain subspaces, providing zero forgetting in isolation.The meta-router adds inference-time selective gating to the anti-forgetting system.
  • Routing: Outcome-based sigmoid routing discovers useful domain combinations through exhaustive loss measurement rather than relying on domain labels.Independent sigmoid gates allow multiple stacks to activate simultaneously for cross-domain composition.
  • Empirical finding: 97% of medical prompts route to chat+math stacks despite zero medical data in those stacks, indicating that stacks encode transferable cognitive primitives.The identified primitives include instruction-following clarity, numerical reasoning, procedural logic, and chain-of-thought structure.
  • Deployment: Disk-offloaded selective stack loading enables constant GPU memory as the number of domain stacks grows.Expertise is loaded on demand per prompt rather than requiring all stacks to remain resident on the GPU.

2 Related Work

Prior work separately addresses parameter efficiency, expert routing, continual learning, adapter freezing, null-space protection, or adapter composition. Brainstacks combines these directions through MoE-LoRA stacks, hard null-space constraints, residual stacking, and outcome-based routing.

  • MoE-LoRA routing: Prior MoE-LoRA methods restrict routing to selected layers or use shared attention adapters, whereas Brainstacks routes experts across all seven transformer projections.Its routing also introduces Shazeer-style learnable noise and combines it with rsLoRA rank-stabilized scaling.
  • Continual learning: Existing adapter-stacking approaches freeze task adapters or use lateral connections but do not combine MoE-LoRA internals with null-space protection and additive composition.Progressive neural networks lack a shared base, while C-LoRA and Online-LoRA select one adapter rather than composing multiple adapters additively.
  • Null-space protection: Compared with soft weight corrections, Brainstacks projects new gradients into the orthogonal complement of prior activation directions as a hard geometric constraint.Randomized SVD extracts prior principal directions before projection, preventing the active stack from writing into claimed subspaces.
  • Adapter composition: Unlike task-ID or single-adapter routing methods, Brainstacks derives routing targets from empirical loss measurements over domain combinations.This enables routing that can exclude the nominally correct domain adapter when another combination performs better.
  • Combined contribution: Brainstacks uniquely combines noisy routing, all-projection MoE-LoRA, rsLoRA scaling, residual frozen stacking, null-space projection, outcome-based gating, and disk-offloaded loading.The paper characterizes this two-level architecture as having no precedent in the literature.

3 Method

Brainstacks builds continual multi-domain learning from frozen MoE-LoRA stacks that are trained sequentially, constrained against prior domain directions, and selectively composed at inference. Its method combines sparse expert corrections, residual boosting, curriculum-ordered domain stacking, null-space protection, and outcome-based routing.

  • MoE-LoRA Building Block: MoELoRADelta replaces each targeted transformer projection with four LoRA experts, using noisy top-2 routing and a load-balance regularizer.The module combines LoRA experts, a noisy router, and an auxiliary loss; only two experts are active per token.
  • MoE-LoRA Building Block: All seven transformer projections use routed experts while the frozen base model remains in 4-bit NF4 form.The projections are q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, and down_proj.
  • Inner Loop: Residual Boosting: Residual boosting trains sequential stacks on the same domain data, allowing later stacks to learn errors left by frozen earlier stacks.The inner loop can run for up to two configurable rounds and stops on diminishing returns.
  • Outer Loop: Continual Domain Training: The outer loop trains domains sequentially in a curriculum, optionally applies per-domain DPO or GRPO, and freezes each domain’s stacks as a plugin.Previous domains are reevaluated after each domain update to check forgetting.
  • Dataset Sensitivity and Ordering: Dataset composition and domain ordering constrain training reliability, including poor medical convergence when medical precedes math and catastrophic code loss from contaminated chat data.These observations motivated a decontamination subsystem based on keyword scoring.
  • Null Space Projection: Null-space projection removes update components along prior frozen-stack directions, forcing new stacks into an orthogonal complement.The constraint is enforced by linear algebra rather than a soft regularization penalty; for Gemma 3 12B, each domain claims 64 of 3840 hidden dimensions.
  • Outcome-Based Meta-Routing: The sigmoid meta-router independently gates domain stacks from prompt representations, enabling multiple stacks to contribute simultaneously.Its input is a weighted combination of mid-layer and last-layer hidden states, with weights 0.45 and 0.55.

4 Experiments

Experiments validate MoE-LoRA as a building block, then show that residual stacking, null-space projection, and selective routing support continual multi-domain learning. Ungated accumulation causes interference, while null-space protection and meta-routing preserve domain performance.

  • 4.1 Experiment 1: 2.5× faster validation-loss convergence per step lets MoE-LoRA reach Single LoRA’s final performance by step ∼160 instead of 400.MoE-LoRA takes 20.2 versus 9.5 minutes because routing and four-expert evaluation add 2× wall-clock overhead.
  • 4.1 Experiment 1: Residual boosting improves chat validation loss from 0.8741 after 400 steps to 0.8531 after three frozen-stack rounds, a 2.4% relative improvement.Each round adds a 53M-parameter stack that learns residual behavior left uncaptured by frozen stacks.
  • 4.1 Experiment 1: MoE-LoRA matches Single LoRA quality with 53.6M versus 50.5M parameters, achieving final validation losses of 0.872 versus 0.874.The higher reported training loss includes an auxiliary load-balance term rather than task loss alone.
  • 4.2 Experiment 2: 9 stacks across 4 domains require ∼485 minutes, with residual boosting reducing code loss from 0.953 to 0.493 in two rounds.Plateau detection terminates further rounds in code, medical, and math when additional stacks provide negligible benefit.
  • 4.2.1 Ungated Interference Analysis: After all four domains, ungated validation losses reach 1.778 for chat, 1.452 for code, 1.984 for medical, and 1.041 for math through cross-stack magnitude accumulation.Isolation evaluation retains each domain’s original training-time loss because frozen weights are unchanged; chat’s sharp increase follows math training.
  • 4.2.2 Null Space Orthogonality and Zero Forgetting: Without null-space projection, final-stage losses are higher for chat, code, medical, and math than with projection, indicating faster interference accumulation.The comparison reports 1.839 versus 1.778 for chat, 1.512 versus 1.452 for code, 2.066 versus 1.984 for medical, and 1.184 versus 1.041 for math.
  • 4.2.2 Null Space Orthogonality and Zero Forgetting: Null-space projection consistently reduces interference across domains, while the full system combines orthogonal training subspaces with selective routing to achieve zero forgetting.The currently training domain can lose a small amount of performance because projection constrains its available subspace.
  • 4.2.3 Routing and Generation Analysis: Uniform inner-expert utilization of 0.232–0.274 per expert matches the expected 0.25 average for top-2-of-4 routing, while routing patterns overlap substantially across domains.Math stacks can nevertheless impose aggressive reasoning patterns on non-math prompts when all stacks fire ungated.

5 Key Findings

Brainstacks finds that empirically routed combinations of frozen domain stacks can provide transferable cognitive capabilities, while ungated accumulation damages generation quality. Boundary experiments indicate these capabilities can emerge without pretrained domain knowledge, but depend on a coherent base representation.

  • Transferable cognitive primitives: 97% of medical prompts routed to chat+math stacks, despite zero overlap between their training data and medical prompts.The oracle found chat+math combinations lower-loss than the medical stack itself.
  • Transferable cognitive primitives: Chat, code, math, and reasoning stacks encode answer structuring, numerical reasoning, procedural logic, and chain-of-thought decomposition.These capabilities transfer beyond the nominal domains that elicited them during training.
  • Outcome-based routing: The outcome-based router learns capability combinations from empirical loss measurements rather than domain labels.It translates oracle-discovered combinations into sigmoid routing targets for inference.
  • Boundary experiment: PSN v2 produces code-like structural patterns from a TinyStories-only base with zero Python training data, supporting capability transfer beyond stored domain knowledge.The generated structure uses the available TinyStories vocabulary rather than Python-specific vocabulary.
  • Outcome-based routing: 10 ungated stacks cause catastrophic generation degradation, whereas meta-routing recovers quality by gating irrelevant stacks while preserving useful compositions.For example, BMI calculation activates medical=1.0, math=1.0, and chat=0.59.
  • Routing sensitivity: Reasoning routing was initially confounded with code by code-like training formats, missing cross-domain pairs, a greedy threshold, and truncated prompts.The v2 router addressed these data and sequence-length issues with verbal-logic data and added reasoning combinations.

6 Additional Experiments

Boundary experiments test Brainstacks beyond ordinary fine-tuning and show both a pretrained-base requirement and practical deployment mechanisms. Per-domain RL is compatible in isolation but exposes instability that requires safeguards before systematic refinement.

  • Pretraining boundary: PSN trains and converges per domain from random initialization, but overall quality remains poor because LoRA corrections lack coherent base representations to refine.The experiment motivates partitioned residual streams for native continual pretraining as future work.
  • Reinforcement-learning boundary: DPO completed successfully on chat, while GRPO code training suffered loss instability that corrupted the active stack.The failure motivated stable SFT stacks before RL and a callback that restores best weights after spikes.
  • Disk-offloaded inference: The modular inference system keeps the base model and router on GPU, loads only required stacks from disk per prompt, and frees them afterward.This supports constant GPU memory with respect to the total number of stacks stored on disk.
  • Disk-offloaded inference: Cache hits eliminate load time for consecutive prompts in the same domain, while measured SSD loading takes approximately one second per stack.On Gemma 3 12B IT, interactive-session cache-hit rate exceeds 80% and router overhead is approximately 5ms per prompt.

7 Discussion

The discussion identifies data quality, training order, routing-data alignment, inference cost, null-space capacity, and pretrained-base geometry as practical boundaries. These constraints shape how Brainstacks must be trained and deployed.

  • Data and training: Short repetitive medical flashcards overfit rapidly, whereas reasoning-oriented MedQA and medical-o1-reasoning data improved generalization.The paper reports flashcard overfitting after 50 steps.
  • Data and training: Cross-domain contamination can cause destructive loss explosions, motivating decontamination and domain reassignment before training.ShareGPT examples intended for chat introduced code and medical signals into later training.
  • Data and training: Curriculum ordering matters: chat supplies formatting, code structured thinking, math numerical reasoning, and medical composes these capabilities.Training medical before math produced poor medical convergence because numerical reasoning was absent.
  • System limitations: Inference latency grows with loaded stacks, while null-space capacity may constrain many domains as each domain claims hidden-space directions.At 50+ domains, the 3840-dimensional Gemma space may become a concern.
  • System limitations: Router quality depends heavily on training data matching the learned signals of each domain stack.Code-like reasoning data contaminated routing until verbal-only sources replaced it.
  • System limitations: Brainstacks requires a pretrained base with coherent hidden-state geometry; random-initialization stacking produced poor results.The proposed partitioned residual-stream architecture remains a future direction for pretraining from scratch.

8 Future Work

Future work extends Brainstacks toward autonomous capability acquisition, native continual pretraining, smaller distilled models, improved RL refinement, and larger-scale domain capacity. Each direction targets a specific current bottleneck.

  • Autonomous expansion: Low router scores could serve as a capability-gap signal triggering autonomous acquisition of new domain stacks.The proposed loop begins when all domain scores fall below the 0.12 gate threshold.
  • Native pretraining: Partitioned residual streams could bring modularity and zero-forgetting guarantees to pretraining from scratch rather than post-hoc fine-tuning.The proposal physically assigns residual-stream subspaces to domain phases during pretraining.
  • Efficient deployment: Latent-space stack computation is proposed to reduce per-stack memory by approximately 16x, from 567MB to approximately 35MB on Gemma 3 12B.This could support consumer-GPU deployment and larger domain counts.
  • Knowledge distillation: A routing-aware distillation objective could compress the routed teacher while preserving multi-stack compositional structure.Naive distillation underperforms because a dense student cannot replicate routing-dependent expert selection.
  • Reinforcement learning: Systematic post-SFT per-domain DPO and GRPO is proposed after stability safeguards address RL-induced weight corruption.Candidate rewards include preferences for chat, execution for code, exact answers for math, and safety rubrics for medicine.
  • Scaling: Larger models and additional domains will test whether null-space capacity and Brainstacks benefits scale with model size.The immediate proposed target is Gemma 3 27B IT or a similar-capacity model.

9 Conclusion

Brainstacks packages domain expertise as frozen, composable MoE-LoRA stacks that support independent capability updates, zero forgetting, and selective inference. Its routing results further suggest that stacks encode transferable cognitive primitives, while frozen-stack deployment offers a path toward scalable modular systems.

  • 9 Conclusion: Brainstacks enables domains to be independently added, removed, or updated with zero forgetting through frozen stacks, null-space projection, and selective gating.Frozen weights preserve per-domain performance, while null-space projection separates domain subspaces and the meta-router reduces cross-stack interference.
  • 9 Conclusion: The outcome-based router shows that domain stacks learn transferable cognitive primitives rather than only domain-specific knowledge.Medical prompts preferentially use chat and math stacks for instruction-following and numerical reasoning capabilities.
  • 9 Conclusion: Frozen stacks support constant-GPU-memory deployment regardless of the number of domain capabilities by loading expertise on demand.The paper compares this deployment model to loading pages of a book and proposes composing a small set of cognitive primitives across many domains.
Loading 2604.01152v1…