RFC: Engram-style Sparse Memory Retrieval in OLMo-core

## Summary

I’ve been experimenting with integrating DeepSeek’s **Engram** module into **OLMo-core** as an optional architectural component.

Engram adds a conditional memory path based on n-gram lookup. The high-level idea is to let the model retrieve static lexical / factual patterns through sparse memory, instead of spending dense FFN or MoE compute reconstructing them.

Current status:

- Forward pass verified
- Autograd verified
- Test training run completed
- Integration tested across:
- Attention + Dense FFN
- Attention + MoE
- GDN / Linear RNN + Dense FFN
- GDN / Linear RNN + MoE

I’m opening this RFC to get feedback on whether this is an interesting direction for OLMo-core, and what the right integration / evaluation path should look like.

A working proof-of-concept exists in my fork. If there is interest in reviewing the implementation, I’m happy to share the branch or open a draft PR.

## Motivation

The core question is not simply:

> Does adding Engram improve the model?

Adding more parameters often helps.

The more useful question is:

> Is Engram an efficient use of parameter budget and compute compared with allocating the same capacity to FFN or MoE components?

This seems especially interesting for dense FFN models, where every token pays the full FFN compute cost. Engram may provide a complementary sparsity path for static memory retrieval.

## Proposed Experiment

I’m proposing a 2×2 ablation across OLMo-core model families:

| | Dense FFN | MoE |
|---|---|---|
| Attention | OLMo-3 style dense baseline | OLMo-2 / OLMoE style sparse FFN |
| GDN / Linear RNN | OLMo Hybrid + Dense FFN | OLMo Hybrid + MoE |

Primary hypothesis:
> Engram’s relative benefit should be largest in the Attention + Dense FFN setting, where the baseline has no existing sparsity mechanism.

Secondary hypothesis:
> In GDN / Linear RNN models, Engram and recurrent state represent different memory mechanisms: static associative memory vs. dynamic in-context memory. The ablation tests whether these mechanisms are complementary.

## Systems Implication

Beyond the modeling question, Engram also raises an interesting systems question: it decouples static memory capacity from dense GPU compute.

The current proof of concept stores the Engram table on GPU for simplicity. Longer term, the memory table could potentially be treated differently from the main transformer weights, for example through sharding, offload, or async prefetch.

I’m not proposing a specific systems design here, but I think this decoupling is part of what makes the architecture interesting. If static memory can scale differently from active compute, Engram may offer a different memory / compute tradeoff than simply increasing FFN or MoE capacity.

Some open questions:

- Should Engram initially be evaluated as a GPU-resident module only?
- Is CPU offload or async prefetch worth exploring later?
- Should lookup be part of the model forward path, or staged earlier in the input pipeline?
- What profiling would be needed to understand whether lookup cost dominates?

## Current Implementation

The current proof of concept includes:

- `EngramConfig`
- Engram module integration
- n-gram hashing / compressed token mapping
- multi-head embedding lookup
- configurable layer injection
- forward / backward tests
- test training script

The current implementation is intended as a correctness / integration prototype, not a final production design.

## Scope

This RFC does not claim benchmark improvements yet. The current status is proof-of-concept integration and training stability validation.

The goal is to discuss whether this direction is worth exploring in OLMo-core, and what the cleanest next step should be.

## Requested Feedback

I’d appreciate feedback on:

1. Whether this direction is interesting for OLMo-core
2. Whether the current integration point makes sense
3. Whether the proposed 2×2 ablation is the right evaluation path
4. Whether Engram should be treated primarily as a model component, a memory-system component, or both
5. Whether a smaller PR would be useful first, such as:
- standalone Engram module
- config scaffolding
- integration test
- training script