Multi-Stream Attention Residuals

Why People Wanted It

Depth attention makes earlier layer outputs retrievable instead of forcing every layer to inherit only their running sum. But one representation of each source may not offer enough distinct views for multiple attention heads: making those views with separate dense projections would add parameters and computation. The real question is how to give depth retrieval several useful perspectives on each source without paying for a stack of expensive projections.

The Core Move

MSAR makes several context-scale views of every residual source, retrieves across depth separately in each view, then adaptively mixes the retrieved results into one layer input. If $x_i$ is source $i$, $E_r$ its stream-specific causal operator, and $z_l^{(r)}$ the retrieval for target layer $l$ in stream $r$, then

$x_i^{(r)}=E_r(x_i),\qquad z_l^{(r)}=\sum_{i<l}\alpha_{i\to l}^{(r)}x_i^{(r)},\qquad h_l=\sum_r g_l^{(r)}z_l^{(r)}.$

The causal operators provide context-diverse source views; the two sets of weights answer different questions: which depths matter within a view, and which views matter for this target layer?

Quiz

In MSAR, what are the two distinct roles of the retrieval weights and the stream-fusion gates?

Retrieval weights mix streams; gates select source layers across depth.

Retrieval weights expand sources; gates perform causal convolution.

Both are one shared normalization over all layers and streams.

Retrieval weights mix source layers within each stream; gates mix the resulting streams for the target layer.

$ check

Why Causal Convolutions Make Contextual Views

Now the version for readers who want more of the machinery: for stream zero, $E_0=I$, preserving the source as-is; for each additional stream, $E_r$ is a layer-specific causal depthwise 1D convolution. At token position $t$, its channel-wise form is

$x_{i,t}^{(r)}=\sum_{\delta=0}^{\kappa_r-1}c_{i,r,\delta}\odot x_{i,t-\delta},$

with positions before the sequence start zero-padded. In words, each channel combines the source at the current token with a learned, bounded history, and no term reaches into a future token. Different kernel sizes give views with different local receptive fields; the identity stream keeps the unfiltered source. The described main setting uses the identity plus convolutions with kernel sizes 4, 8, and 12. This is contextual expansion, not independent dense projection: each convolution is depthwise and causal.

Why Each Stream Retrieves Across Depth Separately

Here is the part that gets technical: after expansion, each stream has its own learned pseudo-query for each target layer. Within stream $r$, source keys and values are the same representation, $k_i^{(r)}=v_i^{(r)}=x_i^{(r)}$, and the normalized weights are

$\alpha_{i\to l}^{(r)}=\frac{\phi(q_l^{(r)},k_i^{(r)})}{\sum_{j<l}\phi(q_l^{(r)},k_j^{(r)})},\qquad z_l^{(r)}=\sum_{i<l}\alpha_{i\to l}^{(r)}v_i^{(r)}.$

In words, each stream scores the available earlier sources against its own query, normalizes those scores over depth, and forms a weighted retrieval. The kernel $\phi$ is the RMS-normalized exponential kernel used by the attention-residual formulation. Normalization is separate per stream, so each stream produces its own depth mixture rather than competing for weight in one shared softmax.

Why Fusion Uses Input-Dependent Scalar Weights

Now the deeper design choice: concatenating all $m$ retrieved vectors and applying a full output projection would require an $md\times d$ matrix. Instead, a lightweight projection from the flattened retrieved streams produces one scalar gate per stream:

$g_l=\sigma\!\left(\alpha_l^{\mathrm{pre}}\,\mathrm{flat}(Z_l)W_l^{\mathrm{pre}}+b_l^{\mathrm{pre}}\right),\qquad h_l=\sum_{r=0}^{m-1}g_l^{(r)}z_l^{(r)}.$

In words, the target layer's retrieved content determines a sigmoid weight for each stream, and the weighted vectors are summed to recover one layer-width input. The gate is scalar per stream, not a learned channel-by-channel output projection. This keeps fusion lightweight while letting each target layer adapt how much it uses each contextual scale.

Minimal Causal Code

This loop makes the three operations visible: causal source expansion, per-stream retrieval over earlier layers, and scalar stream fusion. It omits the paper's particular kernel $\phi$ and uses ordinary dot-product softmax as a readable stand-in.

import torch
import torch.nn.functional as F

# sources: [depth, batch, time, width]
# queries: [target_depth, streams, width]
# conv: one causal depthwise Conv1d per non-identity stream
# fusion: maps flattened [streams * width] to [streams]
expanded = []
for source in sources:
    views = [source]  # stream 0 is the identity
    for conv in convs:
        x = source.transpose(1, 2)  # [batch, width, time]
        left_pad = conv.kernel_size[0] - 1
        x = F.pad(x, (left_pad, 0))
        views.append(conv(x).transpose(1, 2))
    expanded.append(torch.stack(views, dim=0))  # [streams, batch, time, width]

layer_inputs = []
for l in range(len(sources)):
    retrieved = []
    for r in range(len(queries[l])):
        if l == 0:
            retrieved.append(torch.zeros_like(sources[0]))
            continue
        keys = torch.stack([expanded[i][r] for i in range(l)], dim=0)
        scores = torch.einsum("btd,ibtd->ibt", queries[l, r], keys)
        weights = scores.softmax(dim=0)  # normalize over earlier source layers
        retrieved.append((weights.unsqueeze(-1) * keys).sum(dim=0))
    z = torch.stack(retrieved, dim=0)  # [streams, batch, time, width]
    gates = torch.sigmoid(fusion(z.permute(1, 2, 0, 3).flatten(-2)))
    h = (z * gates.permute(2, 0, 1).unsqueeze(-1)).sum(dim=0)
    layer_inputs.append(h)

What You Gain, and What You Lose

You gain depth retrieval with several context scales: the identity view preserves the source, while causal convolutional views add bounded preceding-token context. Separate queries and normalizations let each stream retrieve a different mixture of earlier layers, and adaptive scalar gates decide how those mixtures contribute to the layer input.

You do not get those views for free. Expansion adds convolutional work; retrieval still maintains separate stream results; and fusion requires a projection and a read of those retrieved vectors. The design trades the parameter and compute cost of separate dense source projections for lightweight causal depthwise convolutions and scalar stream fusion. It also gives the layer an extra choice to learn: which contextual views to trust, rather than only which depth sources to retrieve.

Quiz

Why use causal depthwise convolutions to create MSAR's additional source views?

They let each stream inspect future tokens to improve its query.

They replace stream fusion with a full-width output projection.

They add context-scale views using lightweight causal channel-wise operations instead of separate dense source projections.

They remove the need to retrieve historical layers.

$ check