PaperlayerView exact revision on arXiv

Source-linked AI summary

LatentIndex: Cross-Layer Sharing with Layer-Specific Selection for Sparse Attention

Zhaohui Wang, Zhixin Pan, Fanxu Meng, Muhan Zhang

arXiv:2610.04635v1cs.AI

TL;DR

Sparse-attention indexers still incur repeated selection work and per-layer key-cache storage, while reusing selections constrains layers to a shared token set. LatentIndex shares continuous representations while retaining layer-specific selection, and its evaluations report improved recall over index reuse with downstream performance close to native DSA.

  • Problem

    Sparse-attention indexers still score the full prefix at each layer and store per-layer keys, while reusing selected indices constrains followers to the anchor’s token set.

    Source: §1 · p. 2

  • Method

    LatentIndex shares an anchor-derived latent cache across indexer layers and uses layer-specific scoring, with training-free calibration, a training-aware instantiation, and hierarchical candidate refinement.

    Source: §1 · p. 2

  • Results

    Across DeepSeek-V3.2 and GLM-5, training-free LatentIndex improves head-wise attention-mass recall over IndexCache by up to 3.28 percentage points and keeps aggregate RULER and LongBench scores close to native DSA; HS achieves 2.30–2.72× decode indexer speedups over DSA.

    Source: §1 · p. 2

  • Takeaways & Limitations

    LatentIndex shares continuous representations rather than discrete selections, enabling cross-layer reuse while preserving layer-specific token selection.

    Source: §5 · p. 10

Abstract

from arXiv · show

Sparse attention reduces core-attention computation, but its indexers still incur repeated selection work and per-layer key-cache storage. Reusing selected indices across layers reduces this overhead but constrains multiple layers to the same token set. We introduce LatentIndex, which extends the latent-sharing principle of Multi-head Latent Attention across indexer layers. Each layer group constructs a shared latent cache from its first layer's hidden states, while layer-specific scoring enables independent token selection. Absorbing key decoders into queries enables direct scoring of the shared cache without reconstructing historical per-layer keys. We develop training-free calibration and investigate a training-aware instantiation of this principle. To balance quality and computation, a hierarchical selection (HS) variant lets followers independently refine a shared candidate set proposed by the anchor. With four-layer sharing, LatentIndex reduces logical indexer-cache storage by 61.1% on DeepSeek-V3.2. Across DeepSeek-V3.2 and GLM-5, training-free LatentIndex improves head-wise attention-mass recall over IndexCache by up to 3.28 percentage points while maintaining RULER and LongBench performance close to native DSA. HS further achieves 2.30-2.72 times decode indexer speedups over DSA across 8K-128K contexts, retaining most of LatentIndex's recall. LatentIndex offers a new perspective on cross-layer indexing: sharing continuous representations rather than discrete selections enables efficient reuse while preserving layer-specific token selection.

1 INTRODUCTION

Long-context attention makes indexing overhead consequential: sparse attention reduces core-attention cost, but its indexers still repeatedly score the prefix and store per-layer keys. LatentIndex shares continuous representations while retaining layer-specific selection, and its evaluations report improved selection quality over index reuse with downstream performance close to native DSA.

1 INTRODUCTION

  • Indexers still scan the full prefix at each layer and maintain separate key caches, leaving repeated selection computation and per-layer key storage as overheads.Sparse core attention reduces core-attention complexity to O(Lk), while dense attention requires O(L2).Source: §1 · pp. 1–2
  • IndexCache reuses nearby layers’ top-k selections, but forces followers to use the same token set rather than adapt to their own queries and relevance distributions.This motivates cross-layer reuse that preserves layer-specific token selection.Source: §1 · p. 2
  • LatentIndex shares a latent cache derived from each group’s anchor hidden states, while layer-specific decoders and query-side absorption support direct scoring and independent token selection.The main-attention KV caches remain unchanged.Source: §1 · p. 2
  • Hierarchical selection has the anchor propose candidates that followers independently rescore, reducing repeated full-prefix scoring while retaining layer-specific refinement.The paper also develops training-free calibration and investigates training-aware adaptation.Source: §1 · p. 2
  • Training-free LatentIndex improves selection quality over index reuse and keeps downstream performance close to native DSA; hierarchical selection trades some recall for lower latency.Training experiments also show shared latent indexers can be learned while maintaining language-model loss close to native DSA.Source: §1 · p. 2

2 RELATED WORK

Related work addresses sparse attention and its indexing costs through selective access, cross-layer reuse, and representation sharing. LatentIndex explores continuous cross-layer representation sharing for indexers while preserving layer-specific token selection.

Sparse Attention

Efficient Indexing for Sparse Attention

  • Indexing methods reduce cost through cross-layer reuse, restricted search, amortized scoring, compressed representations, or head-cost reduction; LatentIndex instead shares representations across layers.IndexCache reuses anchor-layer top-k indices, while HISA selects candidate blocks before token-level refinement.Source: §2 · p. 3

Representation Sharing and Compression

  • Prior representation-sharing work covers main-attention KV caches, low-rank factorization, and compressed keys; LatentIndex extends sharing to indexer key representations.Examples include CLA, MLKV, YOCO, xKV, CommonKV, MLA, and Loki.Source: §2 · p. 3

3 METHOD: LATENTINDEX

LatentIndex shares a continuous key representation across indexer layers while preserving layer-specific scoring and token selection. Its method combines a shared cache, query-side decoder absorption, calibration or training, and hierarchical candidate refinement.

3 METHOD: LATENTINDEX

  • LatentIndex replaces separate indexer key caches with one shared latent cache per layer group, while each layer retains its own decoder, scoring, and top-k selection.The shared cache is derived from the anchor layer’s hidden states.Source: §3 · p. 4, §3.1 · p. 4

3.1 SHARED KEY REPRESENTATIONS WITH LAYER-SPECIFIC INDEXERS

  • Each group’s anchor hidden states produce a shared RoPE and NoPE representation, which every group member uses through its own decoder to define effective indexer keys.The two branches are concatenated and cached once per group.Source: §3.1 · p. 4
  • At inference, decoders are absorbed into queries so layers score the shared cache directly without reconstructing historical keys, then independently select tokens for their own main-attention KV caches.The reformulation is exact for the effective keys in exact arithmetic.Source: §3.1 · pp. 4–5

3.2 TRAINING-FREE LATENTINDEX

  • Offline calibration jointly fits shared projections and layer-specific decoders with ridge reduced-rank regression; its closed-form fitted parameters remain fixed during inference.The RoPE branch is fit by frequency to preserve positional-rotation compatibility.Source: §3.2 · p. 5
  • Hierarchical selection has the anchor score the full prefix and propose top-B candidates, then lets each follower rescore those candidates and choose its own top-k tokens.Each follower scores at most B positions per query.Source: §3.2 · p. 5

3.3 TRAINING-AWARE LATENTINDEX

  • Training-aware LatentIndex uses dense-attention warmup to train indexers with a frozen backbone, followed by sparse language-model training that adapts the backbone and indexers.During warmup, each layer’s dense-attention distribution supervises its own LatentIndex scoring distribution.Source: §3.3 · p. 6

3.4 STORAGE AND COMPUTATION

  • With four-layer sharing on DeepSeek-V3.2, logical indexer-cache storage falls by 61.1%, while the main-attention KV cache remains unchanged.Hierarchical selection reduces per-query scored positions from mL to L + (m −1) min(B, L) for a group of m layers.Source: §3.4 · p. 6

4 EXPERIMENTS

Across the evaluated models and settings, LatentIndex improves selection quality over index reuse while keeping downstream performance close to native DSA. Hierarchical selection adds decode speedups, and training-aware experiments show loss close to DSA.

4.1 EXPERIMENTAL SETUP

  • The evaluation compares native DSA, IndexCache, and training-free LatentIndex on DeepSeek-V3.2 and GLM-5 using attention-mass recall, RULER, and LongBench.Grouped methods use G = 4 unless stated otherwise, and select up to 2048 tokens per query.Source: §4.1 · p. 6

4.2 TRAINING-FREE TOKEN-SELECTION QUALITY

  • At 128K, LatentIndex improves attention-mass recall over IndexCache by 3.28 percentage points on DeepSeek-V3.2 and 1.57 points on GLM-5.It outperforms IndexCache at every evaluated length on both models, with gains increasing toward longer contexts.Source: Table 1 · p. 7, §4.2 · p. 7

4.3 TRAINING-FREE DOWNSTREAM PERFORMANCE

  • LatentIndex stays close to native DSA on RULER and LongBench while improving aggregate scores over IndexCache on both models.On RULER, average scores are 0.05 and 0.03 points below DSA on DeepSeek-V3.2 and GLM-5; LongBench averages improve over IndexCache by 0.73 and 0.30 points, respectively.Source: Table 2 · p. 7, §4.3 · p. 7, Table 3 · p. 8

4.4 ANALYSIS OF CROSS-LAYER SHARING

  • LatentIndex reduces follower recall loss at every offset and context length; at 128K, average follower loss falls 73%, from 6.69 to 1.81 percentage points.With G = 8 at 128K, LatentIndex reaches 86.21 recall, 4.07 points above IndexCache with G = 8 and 1.15 points above IndexCache with G = 4.Source: Figure 2 · p. 8, §4.4 · p. 8

4.5 HIERARCHICAL SELECTION: QUALITY–EFFICIENCY TRADE-OFF

  • HS achieves 2.30–2.72× decode indexer speedups over native DSA across 8K–128K while exceeding IndexCache on all three aggregate quality metrics.The evaluation uses a fixed candidate budget of 8192; HS has Rhead = 93.19 versus 93.48 for full-prefix LatentIndex.Source: §4.5 · pp. 8–9, Table 4 · p. 9

4.6 TRAINING-AWARE EVALUATION

  • During later warmup steps, training-aware LatentIndex remains about 1.3–1.5 percentage points below DSA in Recall@2048, versus 4.4–4.7 for IndexCache.During sparse training, its LM loss fluctuates close to DSA and is generally lower than IndexCache.Source: Figure 3 · p. 9, §4.6 · p. 9

5 CONCLUSION

LatentIndex shares continuous key representations across layers while preserving layer-specific token selection, and the paper’s results support reducing indexer storage while maintaining selection quality. Its training-free construction jointly fits a shared projection and layer-specific decoders to native indexer keys.

5 CONCLUSION

  • LatentIndex shares continuous key representations across layers while retaining layer-specific selection, with results supporting reduced storage and preserved selection quality.Hierarchical selection adds computational savings, and training experiments provide complementary evidence with language-model loss close to DSA.Source: §5 · p. 10

A OFFLINE CONSTRUCTION OF LATENTINDEX

  • Training-free LatentIndex jointly fits a shared projection and layer-specific decoders to each group’s native indexer keys, including the anchor’s.The NoPE and RoPE branches are fitted separately using ridge reduced-rank regression.Source: A OFFLINE CONSTRUCTION OF LATENTINDEX · p. 11

A.1 CALIBRATION DATA AND TARGETS

The calibration pipeline collects anchor hidden states and group members’ native keys, then fits shared latent sources and layer-specific decoders. It handles rotary keys frequency-wise so the shared latent can be rotated before caching while preserving layer-specific decoding.

A.1 CALIBRATION DATA AND TARGETS

  • Calibration records anchor hidden states after input RMSNorm and group members’ native keys after key normalization, before RoPE, Hadamard rotation, and quantization.Five data sources span general text, code, mathematics, dialogue, and long documents; the protocol uses 50 training and 15 validation requests and keeps calibration data separate from downstream evaluation.Source: §A.1 · p. 12

A.2 ROPE-COMPATIBLE DECODERS

  • Frequency-specific rotary sources share the native rotary phase, allowing the shared latent to be rotated before caching while each layer retains its own decoder.The decoder is absorbed into post-RoPE, pre-Hadamard queries, which score the latent directly without reconstructing historical per-layer keys.Source: §A.2 · p. 12

A.3 JOINT NOPE REGRESSION

  • Joint NoPE regression fits a shared projection from anchor inputs to concatenated group targets, then separates the fitted outputs into layer-specific decoders and biases.Inference projects uncentered inputs; the output biases incorporate the mean correction.Source: §A.3 · pp. 12–13

A.4 FREQUENCY-WISE COMPLEX ROPE REGRESSION

  • Frequency-wise complex RoPE regression fits each rotary frequency independently, with shared complex sources and layer-specific decoder coefficients.The fitted projections generate the shared sources, while each layer’s decoder coefficients reconstruct its target; the final GLM-5 model uses two sources per frequency.Source: §A.4 · p. 13

A.5 REGULARIZATION SELECTION AND NUMERICAL IMPLEMENTATION

The calibration procedure selects regularization strengths using validation reconstruction error while retaining parameters fitted on training data. It computes sufficient statistics separately for training and validation, then solves and verifies the fitted models at high precision.

A.5 REGULARIZATION SELECTION AND NUMERICAL IMPLEMENTATION

  • Both regression branches scale regularization using centered training variance, and all rotary frequencies within a group share one rotary regularization strength.The selection rule retains the strongest regularization within 0.5% of minimum validation error.Source: §A.5 · p. 13
  • For DeepSeek-V3.2, all 15 groups in the G = 4 calibration select αN = αR = 10−2; GLM-5 retains NoPE parameters while refitting rotary projections and decoders.For DeepSeek-V3.2 G = 8, groups are refit using the same calibration and selection protocol.Source: §A.5 · p. 13
  • Training and validation sufficient statistics are accumulated separately in FP64, and rotary statistics are solved in complex128 without storing token-level features.The implementation checks statistics-based validation errors against explicit predictions and exported FP32 parameters against the high-precision solution.Source: §A.5 · p. 13

B STORAGE AND INFERENCE IMPLEMENTATION

LatentIndex reduces logical indexer-cache storage by sharing one latent across each group and absorbs layer-specific decoders into queries for direct latent scoring. Hierarchical selection reduces follower scoring to shared candidates, while the evaluation compares the methods on DeepSeek-V3.2 and GLM-5.

B.1 CACHE STORAGE AND FIXED PARAMETERS

  • Logical indexer-cache storage is reduced by 61.1% with G = 4, from 8052 bytes/token for mDSA to 3132 bytes/token for mG=4; accounting excludes allocation overhead, temporary workspaces, HS candidate indices, fixed parameters, and main-attention KV caches.Source: §B.1 · p. 14

B.2 QUERY ABSORPTION AND FP8 EXECUTION

  • Absorbing each layer’s decoder into its query enables direct scoring of the rotated shared latent, with segmented FP8 execution for DeepSeek-V3.2 and GLM-5.The latent is cached once per group; each layer retains native query and head-weight computations.Source: §B.2 · p. 14

B.3 COMPUTATION WITH AND WITHOUT HIERARCHICAL SELECTION

  • For L > 8192, HS scores 16L + 45 × 8192 positions instead of 61L for native DSA in the stated DeepSeek-V3.2 configuration.The counts cover scoring positions only; complete indexer forward-call measurements include candidate construction and TopK but exclude main attention and other model components.Source: §B.3 · p. 15

C.1 MODELS AND COMPARISON SETTINGS

  • The comparison uses native DSA, IndexCache, and LatentIndex on DeepSeek-V3.2 and GLM-5, with G = 4 for grouped methods and up to 2048 selected tokens per query.Main results use fixed fitted parameters, offline calibration, and segmented FP8 inference.Source: §C.1 · p. 15

C.2 DOWNSTREAM EVALUATION AND RECALL PROTOCOL

The section specifies how downstream benchmark scores and selector recall are evaluated and aggregated. It also identifies the group-offset recall comparison shown in Figure 4.

C.2 DOWNSTREAM EVALUATION AND RECALL PROTOCOL

  • RULER uses 13 tasks at six context lengths, while LongBench covers six categories with scores averaged within and across categories.Each method runs end to end with its own token selections.Source: §C.2 · p. 15
  • Recall compares selectors at the first question token on the native DSA trajectory, using the same dense core-attention probabilities over the causal prefix.The recall protocol samples 100 examples from each RULER task at each nominal context length.Source: §C.2 · p. 15
  • Recall is averaged across heads, indexed layers, tasks, and six equally weighted lengths; DeepSeek-V3.2 uses all 128 core-attention heads.The aggregation includes 61 indexed layers for DeepSeek-V3.2 and 78 for GLM-5.Source: §C.2 · p. 15
  • Figure 4 reports additional Rmean loss relative to native DSA by group offset on DeepSeek-V3.2 with G = 4.It uses the same trajectory, sampling, and aggregation as Figure 2a.Source: Figure 4 · p. 16

C.3 CROSS-LAYER SHARING ANALYSIS

The appendix examines recall by layer offset and group size, then describes hierarchical candidate selection and the decode-indexer timing protocol.

C.3 CROSS-LAYER SHARING ANALYSIS

  • For G = 4, offset analysis measures recall loss relative to each layer’s native DSA selection, averaging corresponding offsets across 15 groups.Offset 0 denotes anchors and offsets 1–3 followers; independent layer 0 is excluded from offset analysis but included in overall recall.Source: §C.3 · p. 16
  • For G = 8, layer 0 stays independent, layers 1–56 form seven groups, and layers 57–60 form a four-layer group.The RoPE64/NoPE128 layout is retained, and projections and decoders are refit; IndexCache uses identical boundaries.Source: §C.3 · p. 16
  • Mean-attention recall tests whether shared token sets preserve head-averaged attention mass relative to the optimal shared size-K set.Figure 4 reports additional Rmean loss relative to native DSA using the same offset analysis.Source: §C.3 · p. 16

C.4 HIERARCHICAL SELECTION AND DECODE INDEXER TIMING

  • Hierarchical selection uses 8,192 anchor-proposed candidates and a 2,048-token final budget; each follower independently rescores candidates and selects its own top-2,048.Budgets are capped by valid prefix length, and candidates are rebuilt for each group at every decode step.Source: §C.4 · p. 16
  • Timing measures complete indexer.forward calls during batch-size-one decoding on eight GPUs with tensor parallelism 8, reporting milliseconds per decode token.Three requests per method and length each generate 65 tokens; timing covers the 64 steps after prefill.Source: §C.4 · pp. 16–17
  • Table 6 states that all methods completed both training schedules and distinguishes global indexer-loss coefficients from method-specific layer weights.The table caption identifies these as training settings.Source: Table 6 · p. 17

D.1 TRAINING CONFIGURATION

The training setup compares DSA, IndexCache, and LatentIndex under matched data and training schedules. The indexers use shared latent representations with layer-specific projections and supervision.

D.1 TRAINING CONFIGURATION

  • Training compares DSA, IndexCache, and LatentIndex on DeepSeek-V2-Lite with nine groups of three layers and matched data order, seed, and sampler.Dense warmup freezes the backbone; sparse training activates each method’s routes and updates the backbone with next-token language-model loss while retaining indexer supervision.Source: §D.1 · p. 17

D.2 INDEXER CONSTRUCTION AND TRAINING OBJECTIVES

  • Each group projects its anchor hidden state into a 128-dimensional NoPE latent and 64-dimensional RoPE key, with layer-specific decoders producing 64-dimensional NoPE keys.LayerNorm applies only to the NoPE latent, while the RoPE key is reused directly.Source: §D.2 · p. 17
  • Queries and head weights come from the anchor hidden state through separate layer-specific projections, and each layer makes an independent top-2,048 selection.During warmup, each layer’s indexer output is evaluated against its own dense-attention teacher using the cached anchor input.Source: §D.2 · p. 17
  • Attention supervision compares each layer’s teacher distribution, formed from core-attention heads, with its indexer softmax distribution over the full causal prefix during warmup.The teacher averages 16 core-attention heads’ probabilities using the runtime attention scale, including YaRN scaling.Source: §D.2 · p. 17
  • During sparse training, indexer inputs are detached from the backbone, while the backbone is optimized with next-token language modeling.The stop-gradient specification is given for the training objective.Source: §D.2 · p. 18

D.3 RECALL DIAGNOSTICS AND PLOTTED TRAINING CURVES

The appendix defines a warmup recall diagnostic, specifies its sampling and aggregation, and describes the reported training-loss curves.

D.3 RECALL DIAGNOSTICS AND PLOTTED TRAINING CURVES

  • Warmup Recall@2048 measures selected attention mass relative to each head’s top-K oracle, with K = min(2048, N) for causal prefix length N.The diagnostic uses each layer’s dense-attention distribution.Source: §D.3 · p. 18
  • Recall samples at most two evenly spaced local queries per layer and context-parallel rank, then averages query recall across ranks and 27 layers.IndexCache followers are scored using the anchor’s queries, keys, and head weights against each follower’s own dense-attention reference.Source: §D.3 · p. 18
  • The diagnostic uses attention scale 1/576 for the absorbed MLA query width, while plotted LM loss is next-token cross-entropy over valid prediction tokens.LM-loss differences are computed at matched steps and smoothed with bias-corrected causal EMA.Source: §D.3 · p. 18
Loading 2610.04635v1…