Architecture Name
Memory Decoder
Parent issue
#12
Motivations
Domain-specific knowledge is often long-tailed and continuously evolving. Although a modern language model can provide strong general reasoning and instruction-following capabilities, no fixed checkpoint can fully cover every specialized field, task protocol, or newly emerging domain.
Directly fine-tuning the backbone for each new domain is undesirable because the same parameter updates that improve specialized performance may perturb the model's general reasoning, instruction following, and other broadly useful capabilities. We propose to turn domain extension from backbone rewriting into modular memory attachment: the backbone remains frozen, while specialized knowledge and task patterns are supplied by an independently trained external parametric memory.
Proposed Architecture
Memory Decoder extends a frozen backbone language model with a plug-and-play parametric memory. The backbone continues to provide general reasoning and instruction-following capabilities, while the memory complements it with knowledge and behaviors learned from a target professional domain.
Memory training compresses retrieval-based domain evidence into a reusable parametric module. A token-level datastore is constructed from a domain instruction-tuning corpus. For each answer-side target token, nearest-neighbor retrieval produces a soft next-token teacher distribution. The memory model is trained with a combination of retrieval-distribution distillation and standard supervision on the gold answer token, allowing it to internalize recurring domain knowledge and task patterns.
At inference time, the frozen backbone and the memory process the same decoding context in parallel and produce separate next-token distributions. A lightweight token-level router uses their hidden representations together with confidence and entropy features to predict a dynamic fusion weight:
$$p_{\mathrm{final}}(\cdot\mid c_t)=(1-\lambda_t)p_{\mathrm{base}}(\cdot\mid c_t)+\lambda_t p_{\mathrm{memory}}(\cdot\mid c_t).$$
The router is trained on a mixture of domain and general instruction data while both language-model branches remain frozen. It learns to increase memory contribution for domain-dependent predictions and preserve backbone-dominant generation for general reasoning and instruction-following tokens.
New capabilities can be introduced by attaching independently trained memories without modifying the backbone. A domain memory may also be reused across compatible backbone models that share the same tokenizer and output vocabulary, with only a lightweight router trained for each backbone.
Preliminary Results (if any)
No response
Experiments Plan
We will evaluate Memory Decoder as a general mechanism for adding specialized capabilities to a frozen backbone.
Research questions
- Can Memory Decoder improve performance in a target domain without degrading the backbone's general capabilities?
- Does token-level dynamic routing provide a better specialization-retention trade-off than direct fine-tuning or static fusion?
- Can the same trained memory be reused across multiple compatible backbone sizes?
- How do memory size, domain-data scale, and router-training data affect specialization quality and general-capability retention?
Comparisons
- Frozen backbone without domain adaptation.
- Full-parameter supervised fine-tuning.
- LoRA or another parameter-efficient fine-tuning baseline.
- Memory Decoder with dynamic token-level routing.
- Static or fixed-weight backbone-memory fusion as a routing ablation.
Evaluation
- Select one or more specialized benchmarks appropriate to each target domain.
- Evaluate general capability retention on a broad suite including math, instruction-following, knowledge, and reasoning benchmarks.
- Report domain-task metrics, per-benchmark general scores, the average change relative to the frozen backbone, training compute, memory and router parameter counts, and inference latency or throughput overhead.
- Use the same evaluation prompts, decoding configuration, and scoring protocol for the frozen backbone and all adaptation methods.
Success criteria
- A clear and reproducible improvement on the selected domain benchmarks.
- No material regression in the aggregate general benchmark score.
- Better specialization-retention trade-off than full fine-tuning and parameter-efficient fine-tuning baselines.
- Successful reuse of the same domain memory on at least two compatible backbone sizes, subject to available compute.
Architecture Name
Memory Decoder
Parent issue
#12
Motivations
Domain-specific knowledge is often long-tailed and continuously evolving. Although a modern language model can provide strong general reasoning and instruction-following capabilities, no fixed checkpoint can fully cover every specialized field, task protocol, or newly emerging domain.
Directly fine-tuning the backbone for each new domain is undesirable because the same parameter updates that improve specialized performance may perturb the model's general reasoning, instruction following, and other broadly useful capabilities. We propose to turn domain extension from backbone rewriting into modular memory attachment: the backbone remains frozen, while specialized knowledge and task patterns are supplied by an independently trained external parametric memory.
Proposed Architecture
Memory Decoder extends a frozen backbone language model with a plug-and-play parametric memory. The backbone continues to provide general reasoning and instruction-following capabilities, while the memory complements it with knowledge and behaviors learned from a target professional domain.
Memory training compresses retrieval-based domain evidence into a reusable parametric module. A token-level datastore is constructed from a domain instruction-tuning corpus. For each answer-side target token, nearest-neighbor retrieval produces a soft next-token teacher distribution. The memory model is trained with a combination of retrieval-distribution distillation and standard supervision on the gold answer token, allowing it to internalize recurring domain knowledge and task patterns.
At inference time, the frozen backbone and the memory process the same decoding context in parallel and produce separate next-token distributions. A lightweight token-level router uses their hidden representations together with confidence and entropy features to predict a dynamic fusion weight:
The router is trained on a mixture of domain and general instruction data while both language-model branches remain frozen. It learns to increase memory contribution for domain-dependent predictions and preserve backbone-dominant generation for general reasoning and instruction-following tokens.
New capabilities can be introduced by attaching independently trained memories without modifying the backbone. A domain memory may also be reused across compatible backbone models that share the same tokenizer and output vocabulary, with only a lightweight router trained for each backbone.
Preliminary Results (if any)
No response
Experiments Plan
We will evaluate Memory Decoder as a general mechanism for adding specialized capabilities to a frozen backbone.
Research questions
Comparisons
Evaluation
Success criteria