Architecture Name
MLP Memory
Parent issue
#1
Motivations
Decoder-only language models achieve strong performance across many generation and reasoning tasks, but they can still produce fluent responses that fail to use factual knowledge accurately. Existing approaches to improving knowledge utilization involve a fundamental trade-off between non-parametric retrieval and parametric adaptation.
Retrieval-augmented generation provides flexible access to external corpora, but introduces nearest-neighbor search, longer retrieved contexts, and substantial inference latency. Its retrieval component is also only shallowly integrated with the language model. Continued pretraining and LoRA avoid retrieval at inference time, but modify model parameters and may cause catastrophic forgetting or degradation on general tasks.
MLP Memory addresses this trade-off with an external parametric module pretrained to imitate a token-level $k$NN retriever over the training corpus. It compresses retrieval behavior into an all-MLP memory that can be used without explicit document access, datastore storage, or neighbor search at inference time.
Proposed Architecture
MLP Memory augments a pretrained language model with a standalone stack of feed-forward blocks. For each decoding position, the backbone produces a contextual hidden representation from a selected intermediate layer. The memory processes this single-token representation independently, without attention or other token-mixing operations, and maps it directly to a distribution over the model vocabulary.
The memory is trained to imitate a token-level $k$NN retriever. A datastore is first constructed by running the pretrained language model over the training corpus and recording pairs of contextual hidden states and their next tokens. For each training position, nearest-neighbor search produces a soft next-token distribution. The query itself is excluded from its retrieved neighbor set to prevent trivial self-retrieval, and the resulting hidden-state and retrieval-distribution pairs are cached as offline supervision.
Given context $c_t$, the memory predicts $p_{\mathrm{MLP}}(\cdot\mid c_t)$ from the backbone representation. Its training objective combines retrieval-distribution distillation with standard next-token supervision:
$$\mathcal{L}(c_t)=\alpha\mathrm{KL}\left(p_{k\mathrm{NN}}(\cdot\mid c_t)|p_{\mathrm{MLP}}(\cdot\mid c_t)\right)+(1-\alpha)\left[-\log p_{\mathrm{MLP}}(w_t\mid c_t)\right].$$
The KL term teaches the memory to reproduce the richer, multi-token distribution produced by retrieval, while the cross-entropy term anchors it to the observed next token and stabilizes language modeling. The default experiments use $\alpha=0.4$. The backbone and the MLP Memory are pretrained separately with different pretraining tasks.
At inference time, the datastore and nearest-neighbor search are removed. The language model and MLP Memory produce next-token distributions from the same decoding state, which are combined through probability interpolation:
$$p_{\mathrm{final}}(\cdot\mid c_t)=(1-\lambda)p_{\mathrm{LM}}(\cdot\mid c_t)+\lambda p_{\mathrm{MLP}}(\cdot\mid c_t).$$
where $\lambda\in[0,1]$ controls the memory contribution and is selected on validation data. This preserves the fluent, general-purpose distribution of the backbone while injecting the sharper long-tail knowledge distribution learned from retrieval.
Preliminary Results (if any)
No response
Experiments Plan
We will evaluate MLP Memory as a retrieval-free parametric memory that learns to reproduce retrieval behavior during pretraining.
Research questions
- Can MLP Memory improve language modeling and knowledge-intensive task performance?
- How does it compare with parametric adaptation and explicit retrieval methods?
- How do memory architecture and training objectives affect performance?
- Can it retain its performance and efficiency advantages across different backbones and data scales?
Comparisons
- Pretrained backbone without memory augmentation.
- Continued pretraining and parameter-efficient fine-tuning.
- Explicit retrieval methods, including RAG and $k$NN-LM.
- MLP Memory with its main architectural and training ablations.
Evaluation
- Evaluate language modeling, knowledge-intensive tasks, and general capabilities.
- Report task performance, training cost, memory size, and inference overhead.
- Use consistent evaluation and decoding settings across all compared methods.
Success criteria
- Clear and reproducible improvements over the pretrained backbone.
- A favorable performance-efficiency trade-off compared with adaptation and retrieval baselines.
- No material regression in general capabilities.
- Lower inference overhead than explicit retrieval without requiring datastore access.
Architecture Name
MLP Memory
Parent issue
#1
Motivations
Decoder-only language models achieve strong performance across many generation and reasoning tasks, but they can still produce fluent responses that fail to use factual knowledge accurately. Existing approaches to improving knowledge utilization involve a fundamental trade-off between non-parametric retrieval and parametric adaptation.
Retrieval-augmented generation provides flexible access to external corpora, but introduces nearest-neighbor search, longer retrieved contexts, and substantial inference latency. Its retrieval component is also only shallowly integrated with the language model. Continued pretraining and LoRA avoid retrieval at inference time, but modify model parameters and may cause catastrophic forgetting or degradation on general tasks.
MLP Memory addresses this trade-off with an external parametric module pretrained to imitate a token-level $k$NN retriever over the training corpus. It compresses retrieval behavior into an all-MLP memory that can be used without explicit document access, datastore storage, or neighbor search at inference time.
Proposed Architecture
MLP Memory augments a pretrained language model with a standalone stack of feed-forward blocks. For each decoding position, the backbone produces a contextual hidden representation from a selected intermediate layer. The memory processes this single-token representation independently, without attention or other token-mixing operations, and maps it directly to a distribution over the model vocabulary.
The memory is trained to imitate a token-level $k$NN retriever. A datastore is first constructed by running the pretrained language model over the training corpus and recording pairs of contextual hidden states and their next tokens. For each training position, nearest-neighbor search produces a soft next-token distribution. The query itself is excluded from its retrieved neighbor set to prevent trivial self-retrieval, and the resulting hidden-state and retrieval-distribution pairs are cached as offline supervision.
Given context$c_t$ , the memory predicts $p_{\mathrm{MLP}}(\cdot\mid c_t)$ from the backbone representation. Its training objective combines retrieval-distribution distillation with standard next-token supervision:
The KL term teaches the memory to reproduce the richer, multi-token distribution produced by retrieval, while the cross-entropy term anchors it to the observed next token and stabilizes language modeling. The default experiments use$\alpha=0.4$ . The backbone and the MLP Memory are pretrained separately with different pretraining tasks.
At inference time, the datastore and nearest-neighbor search are removed. The language model and MLP Memory produce next-token distributions from the same decoding state, which are combined through probability interpolation:
where$\lambda\in[0,1]$ controls the memory contribution and is selected on validation data. This preserves the fluent, general-purpose distribution of the backbone while injecting the sharper long-tail knowledge distribution learned from retrieval.
Preliminary Results (if any)
No response
Experiments Plan
We will evaluate MLP Memory as a retrieval-free parametric memory that learns to reproduce retrieval behavior during pretraining.
Research questions
Comparisons
Evaluation
Success criteria