ArXiv

Role-Decoupled Attention Residuals: Separating Matching and Content Retrieval Across Depth

Authors
Kehan Wang
Categories
cs.AI
arXiv
https://arxiv.org/abs/2608.01075v1
PDF
https://arxiv.org/pdf/2608.01075v1

Brief

Role-Decoupled Attention Residuals (RD-AttnRes) modifies Block Attention Residuals to let queries/keys and values independently route over residual depths, costing one model-width vector per layer and no extra token-token attention. In frozen paired pretraining on FineWeb-Edu (2.0B tokens, five seeds) for 120M/343M models, RD-AttnRes reduced validation NLL by 0.0301/0.0247 (≈2.97%/2.43% perplexity). Controls and diagnostics indicate the improvement arises from decoupled routing rather than extra parameters. Full text was not available; summary is based on the abstract.

Why it matters

RD-AttnRes decouples depth routing for queries/keys versus values in Block Attention Residuals, adding one model-width vector per layer and no extra token-to-token attention; tying the routes exactly recovers the parent AttnRes architecture.

Key details

  • On FineWeb-Edu with frozen paired pretraining (2.0B-token budget, five matched seeds) for 120M and 343M models, RD-AttnRes improved validation NLL in all 10 matched comparisons: mean reductions of 0.0301 (120M) and 0.0247 (343M), corresponding to perplexity drops of ≈2.97% and ≈2.43%.
  • Early-budget controls show the gains are not explained by the extra parameter count, duplicated routing execution, or a fixed value route; routing diagnostics reveal persistent divergence between query-key and value depth distributions, supporting separate depth reads for matching vs. content retrieval.
Source evidence

Abstract

Depth-routing residual architectures allow Transformer layers to retrieve earlier representations instead of inheriting only the immediately preceding state. Existing Block Attention Residuals, however, use a single content-dependent depth mixture to construct the inputs to queries, keys, and values. This design couples two functionally different decisions: queries and keys determine where attention matches, whereas values determine what content is retrieved. We therefore ask whether matching and content retrieval should be forced to read from the same depth. We introduce Role-Decoupled Attention Residuals (RD-AttnRes), a minimal extension that shares one depth route between queries and keys while learning an independent value route over the same residual sources. Tying the two routing queries exactly recovers the parent architecture, while decoupling them adds only one model-width vector per layer and introduces no additional token-to-token attention operation. We evaluate RD-AttnRes using a frozen, paired pretraining protocol on FineWeb-Edu with five matched seeds for both 120M- and 343M-parameter models and a 2.0B-token training budget. RD-AttnRes improves validation negative log-likelihood in all 10 matched comparisons. The mean reductions are 0.0301 and 0.0247, corresponding to perplexity reductions of 2.97 percent and 2.43 percent at 120M and 343M parameters, respectively. Early-budget controls indicate that neither the additional parameter count, duplicated routing execution, nor a fixed value route reproduces the improvement. Routing diagnostics further reveal persistent divergence between the query-key and value depth distributions. These results suggest that, within the evaluated training regime, attention matching and content retrieval benefit from distinct reads over the residual hierarchy.

Comment: a improvement of attnres