Mixture-of-Storage Capacity Scaling

Why People Wanted It

A native memory block has a capacity problem: if you want more online memory state and more trainable memory parameters, simply asking one block to do more does not create independent storage. The real question is how to expand memory without making each token choose only one expert or forcing reads and writes to use the same routing decision.

Mixture-of-Storage (MoS) answers by expanding a native memory block into $E$ parallel memory experts. Every expert participates, while separate dense routers decide how information is written and how existing memory is read.

The Core Move

Give each expert its own memory state, give the two routers a shared set of expert addresses, and let their distinct routing probabilities control storage and utilization. The paper gives the architecture and its equations in the surrounding text, but the displayed formulas are missing from the available page text; so no more detailed formula is reproduced here.

Quiz

What is the defining routing arrangement in Mixture-of-Storage?

A single dense router uses unrelated addresses for each expert.

One sparse router chooses a single expert for both reading and writing.

Separate dense storage and utilization routers share expert addresses but use different projections.

$ check

Why Experts Expand the Memory State

Now the version for readers who want more of the machinery: the state is not just a wider single memory; it gains an explicit expert axis. In MoS layer $l$, expert $e$ has dynamic state $M_{t,e}^{(l)} \in \mathbb{R}^{H_{kv} \times d_k \times d_v}$. Each expert also has independent hyper-memory parameters. The same pre-attention hidden states are supplied to all experts and to both routers, so specialization comes from expert-specific memory and routing rather than different input streams.

That distinction matters: an expert is an independently maintained memory participant, not merely a label attached to one shared state. What determines which parts of the input each independent state receives?

Why Shared Addresses Still Allow Different Routes

Here is the part that gets technical: the routers share a trainable matrix of expert addresses, but use separate projection matrices for utilization and storage. Each router projects the common input into the address dimension and computes probabilities with softmax along the expert dimension. The addresses therefore give both routers a common semantic set of experts, while their separate projections let them assign different probabilities.

This is dense routing, not sparse top-choice dispatch: all experts participate in memory computation. The probabilities determine how information is read from and written into the experts. Sharing addresses does not mean sharing the routing decision. Why should reading from memory use a different mixture from writing to it?

Why Reads Fuse While Writes Stay Expert-Specific

Now the deeper operational detail: each expert independently performs the native memory-utilization procedure and produces its own memory-attention output. MoS forms the final token-level memory attention by weighted read fusion using the utilization routing probabilities. On the storage side, each expert independently produces its memory key, value, erase gate, write gate, and log-decay using expert-specific parameters; storage routing controls the distribution of storage across experts.

So the design keeps the read combination explicit while preserving separate expert states and storage controls. The source pages identify these quantities and operations, but their displayed equations are absent in the available text, so the exact fusion and update equations are not stated here. What does this buy—and what does it cost to keep every expert involved?

Minimal Causal Code

This sketch shows the routing separation and state ownership, not the paper's omitted equations or a complete native-memory update. The update placeholder is deliberately visible rather than pretending to specify details absent from the page text.

# h: shared pre-attention hidden state
# addresses: trainable expert address matrix
# route_u and route_s: separate router projections
u = softmax((h @ route_u) @ addresses.T, dim=-1)  # utilization
s = softmax((h @ route_s) @ addresses.T, dim=-1)  # storage

read = 0
for e in range(num_experts):
    expert_read = native_memory_read(h, state[e])
    read = read + u[..., e:e+1] * expert_read

for e in range(num_experts):
    expert_controls = expert_specific_storage_controls(h, e)
    state[e] = native_memory_update(
        state[e], expert_controls, storage_weight=s[..., e:e+1]
    )

What You Gain, and What You Lose

You gain multiple independent memory states and expert-specific hyper-memory parameters, with dense routing that lets both storage and utilization draw on every expert. Shared addresses provide a common expert representation, while separate router projections let writing and reading distribute differently.

You also pay for expanding the memory state and running every expert in the memory computation: unlike sparse MoE routing, MoS does not skip non-selected experts. And the two-router design introduces distinct storage and utilization routing decisions rather than one shared route. It is a capacity-and-control tradeoff, not a free increase in memory.

Quiz

What is an important cost of MoS's dense routing?

Every expert participates, so the design does not skip non-selected experts as sparse routing would.

Reads and writes must use exactly the same probabilities.

Some experts are never maintained and cannot be read.

$ check