InternLM / InternLM/archspace

[ARCH-PROP] Memory Decoder

Open
#13 0 comments 0 reactions 0 assignees View on GitHub
architecture proposal under review
Dominant language
No language data
Stars
85
Forks
6
PR merge metrics
No merged PRs in 30d

Description

### 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**

1. Can Memory Decoder improve performance in a target domain without degrading the backbone's general capabilities?
2. Does token-level dynamic routing provide a better specialization-retention trade-off than direct fine-tuning or static fusion?
3. Can the same trained memory be reused across multiple compatible backbone sizes?
4. 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.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reading parent issue #12 and locating the repository's existing architecture and experiment entry points; this issue names no files, tests, or implementation path. Define the frozen backbone, parametric memory, token-level router, training data, and baselines before beginning. Done means reproducible domain gains, no material general-capability regression, and documented compute, parameter, and inference overhead, with memory reuse tested across compatible backbones.

Written by the indexing model from the issue text.

Assessment

Domain
ai, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.