Block MSAR: On-Demand Expansion of Block Summaries

Why People Wanted It

Multi-stream attention residuals need hidden states from earlier layers as depth-attention sources. Keeping every layer’s hidden state can create substantial memory and pipeline-communication overhead in large-scale training. The real question is how to preserve multi-stream retrieval without retaining and communicating every expanded stream at every layer.

Block MSAR answers by changing what persists: accumulate layer outputs into one summary per block, keep those summaries in single-stream form, and expand them only when a target layer performs depth attention.

The Core Move

Store completed block summaries as single streams, add the current block’s preceding outputs as a partial source, and expand these sources on demand for retrieval. The source set for a target layer consists of the embedding, completed block summaries, and—when present—the current partial summary; the paper gives this source structure in prose rather than a displayed formula here.

Quiz

Which sources can a target layer use in Block MSAR?

Expanded streams permanently stored for every previous layer.

Only the completed summary of the target’s own block.

Every individual layer output, whether or not its block has completed.

The embedding, completed block summaries, and the current partial when present.

$ check

Why the Current Block Has a Partial Source

Now the version for readers who want more of the machinery: a target layer inside a block cannot use that block’s completed summary yet, so its source is the sum accumulated from earlier outputs in the same block.

The paper specifies that this partial starts at zero and becomes the completed block summary when the block finishes. At the first layer of a block, there are no preceding outputs in that block, so the partial source is absent. It is omitted from both the retrieval sum and the softmax denominator—not treated as an ordinary zero-valued source.

For each target layer, the available sources are the embedding, completed summaries, and the current partial if one exists. Each stream’s depth attention normalizes over the available sources. The paper’s excerpt does not supply the summary or retrieval equations themselves, so the mechanism is stated here without inventing notation for them.

Why Single-Stream Storage Cuts Persistent Traffic

Here is the part that gets technical: expanded streams are deterministic functions of the single-stream summaries, so storing every expanded version is unnecessary for later retrieval.

Block MSAR accumulates outputs before expansion. It therefore retains and communicates single-stream sources, then applies lightweight causal convolutions when depth attention is computed. For $N$ summaries, the reported persistent source storage and pipeline payload per token change from $O(mNd)$ to $O(Nd)$, matching Block AttnRes under the same pipeline schedule. Here $m$ is the number of streams and $d$ is the hidden dimension. This is a claim about persistent sources and communication, not about eliminating every multi-stream intermediate: retrieved streams are still needed for fusion.

Why Retrieval Splits into Two Phases

Now the next layer of detail: targets in one block share the same completed sources, and their learned pseudo-queries are available before that block executes. That makes it possible to batch their completed-source retrievals.

In phase one, batch the target layers’ queries against completed sources, computing each target’s per-stream retrieval results and softmax normalization statistics. Shared source reads reduce repeated memory access across layers. In phase two, each layer evaluates its current partial source, if it has one, and merges that contribution with the precomputed completed-source result through online softmax. The first layer uses only the completed-source result. Fusion follows the merge, so phase one still retains the retrieved streams per target separately from the single-stream source store. The open tradeoff: how do you avoid turning on-demand expansion into another round of intermediate-memory traffic?

Why Fused Expansion Avoids Intermediate Writes

The deeper implementation point is that expanded streams have no other forward-pass consumer: depth attention consumes them immediately. A fused kernel can combine causal-convolution expansion with per-stream attention, avoiding writes of the intermediate expanded source streams to memory and their later reads.

This complements the two-phase strategy rather than replacing it: batching reduces repeated source reads across targets, while fusion removes intermediate expansion writes and reads. The retrieved multi-stream outputs are still written and later read for aggregation, so this does not eliminate all extra memory traffic relative to AttnRes.

Minimal Causal Code

This sketch makes the accumulation order visible. It represents one block and the single-stream summaries; it does not implement stream expansion or attention.

partial = torch.zeros_like(block_outputs[0])

for layer_output in block_outputs:
    # The target layer sees only outputs preceding it in this block.
    current_partial = partial
    retrieved_sources = [embedding, *completed_summaries]
    if has_preceding_output:
        retrieved_sources.append(current_partial)

    # The layer's residual output joins the block's single-stream summary.
    partial = partial + layer_output
    has_preceding_output = True

completed_summaries.append(partial)

What You Gain, and What You Lose

You avoid persistent storage and pipeline communication of expanded source streams, and batching completed-source retrievals can reduce repeated reads across target layers. Fusing expansion with attention avoids writing and rereading those intermediate expanded streams.

You still need causal-convolution expansion when retrieval runs, and the retrieved multi-stream outputs remain materialized for later aggregation. The approach also relies on block structure: targets within a block share completed sources, while each target’s current partial changes with depth. This is a different memory/computation tradeoff, not attention with its costs erased.

Quiz

What does fused stream expansion avoid, and what does it not avoid?

It removes the need to expand sources for depth attention.

It reduces the number of block summaries that must be retained.

It avoids intermediate expanded-source writes and reads, but retrieved multi-stream outputs still incur memory traffic for aggregation.

It eliminates all memory traffic for retrieved streams and aggregation.

$ check