ProxyKV: Cross-Model Proxy Pruning for Efficient Long-Context LLM Inference

# Cross-Model Proxy Pruning for Efficient Long-Context LLM Inference
Source: [https://arxiv.org/html/2605.16360](https://arxiv.org/html/2605.16360)
###### Abstract

Efficient long\-context inference in Large Language Models \(LLMs\) is severely constrained by the Key\-Value \(KV\) cache memory wall, yet existing pruning methods force a choice between low\-latency heuristics that sacrifice precision and high\-precision reconstruction methods that incur prohibitive prefilling overhead\. To bridge this scoring\-cost–accuracy gap, we propose ProxyKV, a cross\-model proxy pruning framework that offloads importance scoring to a lightweight intra\-family Small\-Model Proxy executed asynchronously to the Large\-Model Target\. To bridge the architectural gap between heterogeneous models, we design the HybridAxialMapper, which disentangles temporal feature extraction from cross\-head alignment, together with a Multi\-Granularity Hybrid Loss that shifts the learning objective from rigid regression to relative ranking consistency\. Across the Llama\-3\.1, Qwen\-2\.5, and Qwen\-3 families spanning targets from 7B up to 32B parameters on LongBench, SCBench, and RULER, ProxyKV matches KVZip on aggregate \(recovering∼\\sim98\.7%98\.7\\%of its mean accuracy\) while delivering up to a3\.21×3\.21\\timesprefilling speedup on Llama\-3\.1\-8B \(dual\-GPU;∼\\sim1\.5×1\.5\\timesshared single\-GPU\) and sustaining the speedup at contexts up to 170k tokens on Qwen\-2\.5\-7B\.

## 1Introduction

The capability to process and reason over long\-context sequences has become a critical requirement for modern Large Language Models \(LLMs\), enabling applications ranging from whole\-repository code analysis\(Jianget al\.,[2024](https://arxiv.org/html/2605.16360#bib.bib30)\)to complex multi\-hop document reasoning\(Minaeeet al\.,[2024](https://arxiv.org/html/2605.16360#bib.bib25)\)\. To facilitate efficient autoregressive generation, the Transformer architecture utilizes a Key\-Value \(KV\) cache to store past token activations, effectively eliminating redundant computations\(Kwonet al\.,[2023](https://arxiv.org/html/2605.16360#bib.bib31)\)\. However, as the sequence lengthNNincreases, this indispensable mechanism introduces a severe memory bottleneck\. The linear growth of the KV cache footprint often exceeds the physical memory capacity of individual GPUs, leading to frequent I/O overhead and a prohibitive increase in overall inference latency\(Liet al\.,[2024a](https://arxiv.org/html/2605.16360#bib.bib26)\)\.

KV cache pruning has emerged to mitigate this bottleneck, but existing query\-agnostic solutions face a trade\-off between scoring efficiency and pruning precision: heuristic methods like H2O\(Zhanget al\.,[2023](https://arxiv.org/html/2605.16360#bib.bib1)\)and SnapKV\(Liet al\.,[2024c](https://arxiv.org/html/2605.16360#bib.bib24)\)are fast but their local observation windows miss global semantic dependencies, while reconstruction\-based methods like KVZip\(Kimet al\.,[2025](https://arxiv.org/html/2605.16360#bib.bib2)\)achieve high precision through global context reconstruction at the cost of a prohibitive secondary prefilling pass on the target model\.

As illustrated in[Figure˜1](https://arxiv.org/html/2605.16360#S2.F1)\(a, b\), the two paradigms anchor opposite ends of the efficiency–precision spectrum\. We hypothesize that the high\-precision scoring of reconstruction methods can be achieved without the heavy compute on the target, motivated by intra\-family attention correlation observed by IAM\(Zhaoet al\.,[2025a](https://arxiv.org/html/2605.16360#bib.bib3)\)and SmallKV\(Zhaoet al\.,[2025b](https://arxiv.org/html/2605.16360#bib.bib4)\); however, existing static head\-to\-head alignment is too rigid to bridge structural gaps in head count and layer depth\.

We propose ProxyKV \([Figure˜1](https://arxiv.org/html/2605.16360#S2.F1)\(c\)\), a cross\-model proxy pruning framework that offloads scoring to a lightweight intra\-family Small\-Model Proxy executed asynchronously to the target’s critical path\. To handle architectural discrepancies, our HybridAxialMapper disentangles temporal feature extraction from cross\-head alignment, and a Multi\-Granularity Hybrid Loss shifts the learning objective from rigid value regression to relative ranking consistency—more directly aligned with Top\-KKpruning quality\. Our contributions: \(i\) an asynchronous proxy\-based pruning framework that removes the scoring bottleneck from the target’s critical path, delivering up to3\.21×3\.21\\timesprefilling speedup on Llama\-3\.1\-8B; \(ii\) the HybridAxialMapper and a five\-term Multi\-Granularity Hybrid Loss that together bridge structural gaps between heterogeneous models; and \(iii\) an empirical study across the Llama\-3\.1, Qwen\-2\.5, and Qwen\-3 families \(7B–32B targets\) showing ProxyKV recovers∼\\sim98\.7%98\.7\\%of the KVZip oracle and sustains the speedup at contexts up to 170k tokens\.

## 2Related work

![Refer to caption](https://arxiv.org/html/2605.16360v1/x1.png)Figure 1:Three KV\-cache pruning paradigms: SnapKV \(a\) heuristic, KVZip \(b\) reconstruction, ProxyKV \(c\) asynchronous proxy\.#### Heuristic and Architectural Pruning\.

Rule\-based methods identify non\-essential tokens via local patterns: StreamingLLM\(Xiaoet al\.,[2023](https://arxiv.org/html/2605.16360#bib.bib5)\)retains attention sinks; H2O\(Zhanget al\.,[2023](https://arxiv.org/html/2605.16360#bib.bib1)\), Scissorhands\(Liuet al\.,[2023](https://arxiv.org/html/2605.16360#bib.bib8)\), and AhaKV\(Guet al\.,[2025](https://arxiv.org/html/2605.16360#bib.bib9)\)use accumulated scores or recent attention patterns\. Architectural variants further exploit structure: FastGen\(Geet al\.,[2023](https://arxiv.org/html/2605.16360#bib.bib7)\)applies head\-wise differentiated policies, PyramidKV\(Caiet al\.,[2024](https://arxiv.org/html/2605.16360#bib.bib6)\)assigns layer\-wise budgets, and SnapKV\(Liet al\.,[2024c](https://arxiv.org/html/2605.16360#bib.bib24)\)clusters keys via prompt\-end observation windows\. ProxyKV instead distills global reconstruction signals, yielding finer\-grained pruning that is robust in query\-agnostic scenarios\.

#### Learned and Surrogate Pruning\.

Data\-driven policies pursue adaptive computation: CoT\-Influx\(Huanget al\.,[2024](https://arxiv.org/html/2605.16360#bib.bib10)\)uses RL for coarse\-to\-fine selection, while GateSkip\(Laitenbergeret al\.,[2025](https://arxiv.org/html/2605.16360#bib.bib11)\)and LTP\(Kimet al\.,[2022](https://arxiv.org/html/2605.16360#bib.bib14)\)introduce differentiable gating\. Most recently, KVzap\(Jegou and Jeblick,[2026](https://arxiv.org/html/2605.16360#bib.bib17)\)approximates reconstruction oracles via per\-layer surrogates on hidden states; we do not include it as a baseline because \(i\) its training cost is substantially higher than ProxyKV’s, \(ii\) its released implementation adopts a token\-selection criterion that differs from the KVzip oracle and is therefore not directly comparable, and \(iii\) it does not support evaluation at a fixed retention/compression ratio, which is the standard protocol on LongBench, SCBench, and RULER\. Unlike these, which run surrogates within the target’s execution, ProxyKV offloads scoring entirely off the critical inference path, with lower training overhead\.

#### Cross\-Model KV Cache Alignment\.

Leveraging intra\-family correlations, IAM\(Zhaoet al\.,[2025a](https://arxiv.org/html/2605.16360#bib.bib3)\)reuses attention similarity across scales, SmallKV\(Zhaoet al\.,[2025b](https://arxiv.org/html/2605.16360#bib.bib4)\)compensates pruning loss with smaller models, and SpeContext\(Xuet al\.,[2025](https://arxiv.org/html/2605.16360#bib.bib18)\)aligns retrieval for speculative prefetch\. These rely on rigid head\- or layer\-wise mappings between fixed pairs; ProxyKV trains a*learnable cross\-axial*mapping that jointly models temporal context and head\-axis alignment, transferring across architectures with disparate head counts and depths\. ProxyKV is also conceptually adjacent to speculative decoding\(Kwonet al\.,[2023](https://arxiv.org/html/2605.16360#bib.bib31)\), but speculates an*importance\-score distribution*for prefill rather than tokens for generation\.

## 3System description and problem formulation

#### KV cache pruning background\.

Transformer LLM decoding caches the Key/Value tensors of past tokens to avoid redundant compute\. As context lengthNNgrows the cache becomes a memory wall, motivating pruning that retains only the most informative KV pairs and balances*efficiency*\(low scoring latency\) against*precision*\(long\-context reasoning quality\)\.

#### System overview\.

ProxyKV is a deploy\-time pipeline of three components: the Large\-Model Targetℳl\\mathcal\{M\}\_\{l\}\(owns the prefill critical path and the entire decode phase\), an intra\-family Small\-Model Proxyℳs\\mathcal\{M\}\_\{s\}that runs asynchronously on a separate execution stream to extract cross\-head attention features𝐗\\mathbf\{X\}, and the HybridAxialMapperΦθ\\Phi\_\{\\theta\}that maps𝐗\\mathbf\{X\}to target\-aligned importance scores𝐘^\\hat\{\\mathbf\{Y\}\}\. Concretely,𝐗\\mathbf\{X\}is the post\-softmax attention probability accumulated over the query axis at each proxy layer, i\.e\.,𝐗b,ℓ,h,n=∑q=1NSoftmax​\(𝐐𝐊⊤/D\)b,ℓ,h,q,n\\mathbf\{X\}\_\{b,\\ell,h,n\}=\\sum\_\{q=1\}^\{N\}\\mathrm\{Softmax\}\(\\mathbf\{Q\}\\mathbf\{K\}^\{\\\!\\top\}\\\!/\\\!\\sqrt\{D\}\)\_\{b,\\ell,h,q,n\}, which compresses a quadraticN×NN\{\\times\}Nattention map into a single per\-key importance vector of lengthNNper \(batch, layer, head\)\. At inference, the input context is dispatched to both models concurrently;Φθ\\Phi\_\{\\theta\}ingests the proxy features and emits a target\-shaped score driving the Top\-KKpruning mask, after which the target performs decode*without*the proxy on its critical path\. The pipeline supports two deployment regimes \([Figure˜1](https://arxiv.org/html/2605.16360#S2.F1)\(c\)\): a dual\-GPU regime with parallel target/proxy on separate devices, and a single\-GPU regime where they share a device via independent CUDA streams; the proxy KV cache is released at the end of prefill so the proxy\-side memory premium is transient \([Figure˜12](https://arxiv.org/html/2605.16360#A2.F12)\)\.

#### Problem formulation\.

Letℳl\\mathcal\{M\}\_\{l\}haveLlL\_\{l\}layers andHlH\_\{l\}heads per layer with input lengthNNand per\-layer KV tensors𝐊,𝐕∈ℝB×Hl×N×D\\mathbf\{K\},\\mathbf\{V\}\\in\\mathbb\{R\}^\{B\\times H\_\{l\}\\times N\\times D\}, and letℳs\\mathcal\{M\}\_\{s\}haveLsL\_\{s\}layers andHs≪HlH\_\{s\}\\ll H\_\{l\}heads per layer\. Reconstruction\-based oracles \(e\.g\., KVZip\) derive the ground\-truth attention scores𝐘∈ℝB×Ll×Hl×N\\mathbf\{Y\}\\in\\mathbb\{R\}^\{B\\times L\_\{l\}\\times H\_\{l\}\\times N\}via a secondary teacher\-forced prefill onℳl\\mathcal\{M\}\_\{l\}\([Figure˜1](https://arxiv.org/html/2605.16360#S2.F1)\(b\)\), inflating TTFT\. ProxyKV instead approximates𝐘\\mathbf\{Y\}via a learnable cross\-axial mapping𝐘^=Φθ​\(𝐗\)≈𝐘\\hat\{\\mathbf\{Y\}\}=\\Phi\_\{\\theta\}\(\\mathbf\{X\}\)\\approx\\mathbf\{Y\},𝐗∈ℝB×Ls×Hs×N\\mathbf\{X\}\\in\\mathbb\{R\}^\{B\\times L\_\{s\}\\times H\_\{s\}\\times N\}, and produces a binary mask𝐌∈\{0,1\}B×Ll×Hl×N\\mathbf\{M\}\\in\\\{0,1\\\}^\{B\\times L\_\{l\}\\times H\_\{l\}\\times N\}by Top\-KKthresholding𝐘^\\hat\{\\mathbf\{Y\}\}at retention ratioρ∈\(0,1\]\\rho\\in\(0,1\]\. The mapperΦθ\\Phi\_\{\\theta\}is applied independently to each \(proxy\-layer, target\-layer\) pair according to the layer\-pairing schedule of[Section˜4\.1](https://arxiv.org/html/2605.16360#S4.SS1), so all per\-layer formulations below describe the operation on a single such pair; we drop the explicitLLindex from𝐗,𝐘,𝐌\\mathbf\{X\},\\mathbf\{Y\},\\mathbf\{M\}in the rest of[Section˜3](https://arxiv.org/html/2605.16360#S3)and[Section˜4\.2](https://arxiv.org/html/2605.16360#S4.SS2)to lighten notation\. We use*attention scores*for𝐘\\mathbf\{Y\}and*importance scores*for𝐘^\\hat\{\\mathbf\{Y\}\}\. The training objective is

θ⋆\\displaystyle\\theta^\{\\star\}=arg⁡minθ⁡𝔼\(𝐗,𝐘\)​\[ℒ​\(Φθ​\(𝐗\),𝐘\)\],\\displaystyle=\\arg\\min\_\{\\theta\}\\;\\mathbb\{E\}\_\{\(\\mathbf\{X\},\\mathbf\{Y\}\)\}\\big\[\\mathcal\{L\}\\\!\\left\(\\Phi\_\{\\theta\}\(\\mathbf\{X\}\),\\mathbf\{Y\}\\right\)\\big\],\(1\)s\.t\.​tscore​\(ℳs,Φθ\)≤tprefill​\(ℳl\),family​\(ℳs\)=family​\(ℳl\),\\displaystyle\\text\{s\.t\.\}\\;\\;t\_\{\\text\{score\}\}\(\\mathcal\{M\}\_\{s\},\\Phi\_\{\\theta\}\)\\leq t\_\{\\text\{prefill\}\}\(\\mathcal\{M\}\_\{l\}\),\\;\\;\\mathrm\{family\}\(\\mathcal\{M\}\_\{s\}\)=\\mathrm\{family\}\(\\mathcal\{M\}\_\{l\}\),wheretscore​\(ℳs,Φθ\)=tproxy​\(ℳs\)\+tΦθt\_\{\\text\{score\}\}\(\\mathcal\{M\}\_\{s\},\\Phi\_\{\\theta\}\)=t\_\{\\text\{proxy\}\}\(\\mathcal\{M\}\_\{s\}\)\+t\_\{\\Phi\_\{\\theta\}\}is the total time to run the proxy and the mapper end\-to\-end, and the async\-budget constraint keeps both off the target’s critical path \(otherwise the speedup vanishes\); the intra\-family constraint preserves the layer\-wise attention correlation thatΦθ\\Phi\_\{\\theta\}exploits\. The lossℒ\\mathcal\{L\}is the multi\-granularity hybrid loss formalized in[Section˜4\.2](https://arxiv.org/html/2605.16360#S4.SS2); the deployment metric is task\-accuracy recovery against the KVZip oracle at the sameρ\\rho\.

## 4Methodology

![Refer to caption](https://arxiv.org/html/2605.16360v1/x2.png)Figure 2:Overview of ProxyKV: an asynchronous Small\-Model Proxyℳs\\mathcal\{M\}\_\{s\}feeds the HybridAxialMapper, which produces target\-aligned importance scores𝐘^\\hat\{\\mathbf\{Y\}\}for the Large\-Model Targetℳl\\mathcal\{M\}\_\{l\}without a secondary prefilling pass\.Building on the system pipeline of[Section˜3](https://arxiv.org/html/2605.16360#S3.SS0.SSS0.Px2), the HybridAxialMapperΦθ\\Phi\_\{\\theta\}transforms the proxy features𝐗\\mathbf\{X\}through a three\-stage pipeline—\(1\)*Temporal Feature Extraction*via a 1D Conv Stem, \(2\)*Time\-Axis Context Encoder*via a Transformer encoder, \(3\)*Head\-Axis Cross\-Attention*via learnable target queries \([Figure˜2](https://arxiv.org/html/2605.16360#S4.F2)\)—and the predicted scores produce the binary mask𝐌\\mathbf\{M\}at the prescribed budgetℬ\\mathcal\{B\}\. The two subsections below detail the mapper architecture and the multi\-granularity hybrid loss\.

### 4\.1HybridAxialMapper architecture

#### Design rationale\.

Cross\-model alignment must reconcile two coupled axes: a*temporal*axis along which token saliency evolves over the sequence, and a*head*axis along which proxy and target attention patterns are partitioned at different granularities \(Hs≪HlH\_\{s\}\\ll H\_\{l\}\)\. Prior static head\-to\-head schemes\(Zhaoet al\.,[2025a](https://arxiv.org/html/2605.16360#bib.bib3),[b](https://arxiv.org/html/2605.16360#bib.bib4)\)entangle these axes by fixing a permutation between proxy and target heads, which prevents transfer to architectures with different head counts or layer depths\. Our design instead decomposes the mapping into a temporal stage that produces a head\-agnostic latent and a head\-axis stage that learns the proxy\-to\-target aggregation, so the same recipe applies to any intra\-family pair\. Across the layer axis, we adopt a depth\-proportional schedule that pairs target layerℓl∈\[1,Ll\]\\ell\_\{l\}\\in\[1,L\_\{l\}\]with proxy layerℓs=⌈ℓl⋅Ls/Ll⌉\\ell\_\{s\}=\\lceil\\ell\_\{l\}\\cdot L\_\{s\}/L\_\{l\}\\rceiland runs an independent forward pass ofΦθ\\Phi\_\{\\theta\}for each pair; mapper parameters are shared across allLlL\_\{l\}pairs, soΦθ\\Phi\_\{\\theta\}remains a single network whose forward is invokedLlL\_\{l\}times per request and whose output, stacked along the layer axis, recovers the𝐁×𝐋𝐥×𝐇𝐥×𝐍\\mathbf\{B\\times L\_\{l\}\\times H\_\{l\}\\times N\}score tensor of[Section˜3](https://arxiv.org/html/2605.16360#S3.SS0.SSS0.Px3)\.

The three\-stage pipeline disentangles temporal dependency extraction from head\-specific alignment\.

#### Stage 1: temporal feature extraction\.

A 1D Convolutional Stem \(two stacked Conv1D layers withk=3k\{=\}3,p=1p\{=\}1, each followed by BatchNorm and GELU\) projects the raw proxy features𝐗\\mathbf\{X\}into a model\-agnostic latent space along the temporal axis:

𝐙l​o​c​a​l=GELU​\(BN​\(Conv1D​\(𝐗\)\)\)∈ℝB×Dt​i​m​e×N,\\mathbf\{Z\}\_\{local\}=\\mathrm\{GELU\}\(\\mathrm\{BN\}\(\\mathrm\{Conv1D\}\(\\mathbf\{X\}\)\)\)\\in\\mathbb\{R\}^\{B\\times D\_\{time\}\\times N\},\(2\)whereDt​i​m​e=512D\_\{time\}\{=\}512acts as a bottleneck that normalizes the proxy’s head configuration\.

#### Stage 2: time\-axis context encoder\.

We transpose𝐙l​o​c​a​l\\mathbf\{Z\}\_\{local\}toℝB×N×Dt​i​m​e\\mathbb\{R\}^\{B\\times N\\times D\_\{time\}\}\(axis swap, not memory reshape\), add sinusoidal positional encodings𝐏p​o​s\\mathbf\{P\}\_\{pos\}, and process the sequence with a 6\-layer Transformer Encoder \(8 heads, FFN dim4​Dt​i​m​e4D\_\{time\}\) to produce a globally contextualized𝐙g​l​o​b​a​l\\mathbf\{Z\}\_\{global\}\. Trained on2,0482\{,\}048\-token crops, the encoder is applied at inference via a stride\-1,0241\{,\}024sliding window with overlap averaging, keeping positional indices in\-distribution and attention compute at𝒪​\(N\)\\mathcal\{O\}\(N\)even at170170k tokens\.

#### Stage 3: head\-axis cross\-attention\.

𝐙g​l​o​b​a​l\\mathbf\{Z\}\_\{global\}is projected and reshaped into latent keys/values𝐊,𝐕∈ℝ\(B⋅N\)×Hs×Dh​e​a​d\\mathbf\{K\},\\mathbf\{V\}\\in\\mathbb\{R\}^\{\(B\\cdot N\)\\times H\_\{s\}\\times D\_\{head\}\}, treatingHsH\_\{s\}as a synthetic head axis \(the Stage 111D Conv has already mixed the original proxy heads into theDt​i​m​eD\_\{time\}bottleneck, so theHsH\_\{s\}slots here are learned head\-axis tokens rather than the raw proxy heads\)\. A bank of Learnable Target Queries𝐐l∈ℝHl×Dh​e​a​d\\mathbf\{Q\}\_\{l\}\\in\\mathbb\{R\}^\{H\_\{l\}\\times D\_\{head\}\}then aggregates these synthetic head tokens into target heads independently per token,

𝐎=Softmax​\(𝐐l​𝐊⊤/Dh​e​a​d\)​𝐕,\\mathbf\{O\}=\\mathrm\{Softmax\}\\\!\\left\(\\mathbf\{Q\}\_\{l\}\\mathbf\{K\}^\{\\top\}/\\sqrt\{D\_\{head\}\}\\right\)\\mathbf\{V\},\(3\)followed by a learned linear projectionℝDh​e​a​d→ℝ\\mathbb\{R\}^\{D\_\{head\}\}\\\!\\to\\\!\\mathbb\{R\}to obtain final pruning scores𝐘^∈ℝB×Hl×N\\hat\{\\mathbf\{Y\}\}\\in\\mathbb\{R\}^\{B\\times H\_\{l\}\\times N\}\. The query bank is the onlyHlH\_\{l\}\-dependent parameter, so the mapper trivially scales to arbitrary target head counts\.

#### Parameter and FLOP budget\.

The complete mapper has∼\\sim1515M trainable parameters across the three stages—roughly0\.50\.5M in the convolutional stem,1313M in the six\-layer transformer encoder, and1\.51\.5M in the head\-axis cross\-attention block \(queries, KV projections, output linear\)—two\-to\-three orders of magnitude smaller than any target we evaluate\. Forward FLOPs scale as𝒪​\(N⋅Dt​i​m​e2\)\\mathcal\{O\}\(N\\cdot D\_\{time\}^\{2\}\)from the encoder and𝒪​\(N⋅Hl⋅Hs⋅Dh​e​a​d\)\\mathcal\{O\}\(N\\cdot H\_\{l\}\\cdot H\_\{s\}\\cdot D\_\{head\}\)from the cross\-attention, so the dominant term is linear in sequence length and independent of the target’s hidden width; this is what bounds the mapper’s wall\-clock share to a single\-digit percentage of prefill in[Figure˜8](https://arxiv.org/html/2605.16360#S5.F8)and ensures the share stays bounded as targets grow from 7B to 32B\.

### 4\.2Multi\-granularity hybrid loss

#### Why ranking\-consistent loss\.

The deployment metric is Top\-KKretention quality, not pointwise score reconstruction; under sparse, long\-tailed attention distributions a small absolute regression error at non\-Top\-KKtokens can flip the binary mask, so MSE alone trains the mapper to imitate magnitudes the target itself never uses\. We therefore decompose the objective into one term that supervises the binary decision boundary at multiple retention ratios and four auxiliaries that supervise complementary structural facets of the score distribution\.

We minimize a composite lossℒt​o​t​a​l=∑k∈𝒦λk​ℒk\\mathcal\{L\}\_\{total\}=\\sum\_\{k\\in\\mathcal\{K\}\}\\lambda\_\{k\}\\mathcal\{L\}\_\{k\}with𝒦=\{b​i​n,m​s​e,f​i​n​e,g​l​o​b​a​l,c​o​s\}\\mathcal\{K\}=\\\{bin,mse,fine,global,cos\\\}, centered on a Multi\-Ratio Binary term and complemented by auxiliaries covering numerical magnitude, intra\-Top\-KKorder, signal–noise boundary, and directional drift\.

#### Multi\-ratio binaryℒb​i​n\\mathcal\{L\}\_\{bin\}and value\-weighted MSEℒm​s​e\\mathcal\{L\}\_\{mse\}\.

The mapper’s final linear layer emits a single scalar logit per \(head, token\); this logit feeds two parallel heads with different output activations: a sigmoid headσ​\(𝐘^\)\\sigma\(\\hat\{\\mathbf\{Y\}\}\)that yields a probability for the binary objectiveℒb​i​n\\mathcal\{L\}\_\{bin\}, and a magnitude headσ​\(𝐘^\)⋅smax\\sigma\(\\hat\{\\mathbf\{Y\}\}\)\\cdot s\_\{\\max\}that rescales the same logit to the target’s attention\-score range for the regression objectiveℒm​s​e\\mathcal\{L\}\_\{mse\}\(wheresmax=maxi​j⁡𝐘i​js\_\{\\max\}=\\max\_\{ij\}\\mathbf\{Y\}\_\{ij\}is the per\-batch ground\-truth maximum\)\. This shared\-logit, dual\-head design eliminates the domain conflict between BCE \(which pushes the logit to±∞\\pm\\inftyfor confident decisions\) and MSE \(which would otherwise pull the same logit to a finite magnitude≥0\\geq 0\)\. For retention ratiosℛ=\{0\.05,0\.1,…,0\.5\}\\mathcal\{R\}=\\\{0\.05,0\.1,\\ldots,0\.5\\\}we materialize binary ground\-truth masks𝐌r∗\\mathbf\{M\}\_\{r\}^\{\\ast\}from the target’s Top\-KKthreshold and apply a power\-law weighted BCE that emphasizes aggressive compression;ℒm​s​e\\mathcal\{L\}\_\{mse\}is a value\-weighted MSE with weightswi​j=\(𝐘i​j\+ϵ\)1\.5w\_\{ij\}=\(\\mathbf\{Y\}\_\{ij\}\+\\epsilon\)^\{1\.5\}\(ϵ=0\.1\\epsilon\{=\}0\.1\) so salient tokens dominate the magnitude error:

ℒb​i​n=∑r∈ℛ\(rm​i​nr\)γ​BCE​\(σ​\(𝐘^\),𝐌r∗\),ℒm​s​e=1\|Ω\|​∑\(i,j\)∈Ωwi​j​\(σ​\(𝐘^i​j\)⋅smax−𝐘i​j\)2,\\mathcal\{L\}\_\{bin\}=\\\!\\\!\\sum\_\{r\\in\\mathcal\{R\}\}\\\!\\left\(\\\!\\frac\{r\_\{min\}\}\{r\}\\\!\\right\)^\{\\\!\\\!\\gamma\}\\\!\\mathrm\{BCE\}\\\!\\left\(\\sigma\(\\hat\{\\mathbf\{Y\}\}\),\\mathbf\{M\}\_\{r\}^\{\*\}\\right\),\\quad\\mathcal\{L\}\_\{mse\}=\\frac\{1\}\{\|\\Omega\|\}\\\!\\sum\_\{\(i,j\)\\in\\Omega\}\\\!w\_\{ij\}\\,\\big\(\\sigma\(\\hat\{\\mathbf\{Y\}\}\_\{ij\}\)\\cdot s\_\{\\max\}\\\!\-\\\!\\mathbf\{Y\}\_\{ij\}\\big\)^\{2\},\(4\)withγ=1\.0\\gamma\{=\}1\.0andrm​i​n=0\.05r\_\{min\}\{=\}0\.05\.

#### Intra\-Top\-KKrankℒf​i​n​e\\mathcal\{L\}\_\{fine\}, mass\-aware global rankℒg​l​o​b​a​l\\mathcal\{L\}\_\{global\}, cosineℒc​o​s\\mathcal\{L\}\_\{cos\}\.

ℒf​i​n​e\\mathcal\{L\}\_\{fine\}enforces relative order*within*the target’s Top\-KKset via a softplus pairwise loss weighted by\|𝐘i−𝐘j\|\|\\mathbf\{Y\}\_\{i\}\\\!\-\\\!\\mathbf\{Y\}\_\{j\}\|;ℒg​l​o​b​a​l\\mathcal\{L\}\_\{global\}is a margin\-based pairwise loss over \(positivei∈TopKi\\\!\\in\\\!\\mathrm\{TopK\}, negativej∉TopKj\\\!\\notin\\\!\\mathrm\{TopK\}\) samples with quality\-aware weightswi​jg​l​o​b​a​l=clip​\(1\+\|yi,norm−yj,norm\|,1,5\)w\_\{ij\}^\{global\}=\\mathrm\{clip\}\(1\+\|y\_\{i,\\mathrm\{norm\}\}\-y\_\{j,\\mathrm\{norm\}\}\|,1,5\)on the min\-max\-normalised scoresyi,normy\_\{i,\\mathrm\{norm\}\}, upweighting well\-separated pairs that are most informative for the signal\-to\-noise boundary;ℒc​o​s\\mathcal\{L\}\_\{cos\}aligns the global orientation of the predicted score vector to the target:

ℒf​i​n​e=𝔼i,j∈TopK​\[w¯i​j​ln⁡\(1\+e−sgn​\(𝐘i−𝐘j\)​\(𝐘^i−𝐘^j\)\)\],ℒc​o​s=1−𝐘^⋅𝐘‖𝐘^‖​‖𝐘‖,\\mathcal\{L\}\_\{fine\}=\\mathbb\{E\}\_\{i,j\\in\\mathrm\{TopK\}\}\\\!\\left\[\\bar\{w\}\_\{ij\}\\ln\\\!\\\!\\left\(1\+e^\{\-\\mathrm\{sgn\}\(\\mathbf\{Y\}\_\{i\}\-\\mathbf\{Y\}\_\{j\}\)\\,\(\\hat\{\\mathbf\{Y\}\}\_\{i\}\-\\hat\{\\mathbf\{Y\}\}\_\{j\}\)\}\\right\)\\right\],\\quad\\mathcal\{L\}\_\{cos\}=1\-\\frac\{\\hat\{\\mathbf\{Y\}\}\\\!\\cdot\\\!\\mathbf\{Y\}\}\{\\\|\\hat\{\\mathbf\{Y\}\}\\\|\\,\\\|\\mathbf\{Y\}\\\|\},\(5\)ℒg​l​o​b​a​l=𝔼i∈TopK,j∉TopK​\[wi​jg​l​o​b​a​l​\[1−\(𝐘^i−𝐘^j\)\]\+\],\\mathcal\{L\}\_\{global\}=\\mathbb\{E\}\_\{i\\in\\mathrm\{TopK\},\\,j\\notin\\mathrm\{TopK\}\}\\\!\\left\[w\_\{ij\}^\{global\}\\,\\big\[\\,1\-\(\\hat\{\\mathbf\{Y\}\}\_\{i\}\-\\hat\{\\mathbf\{Y\}\}\_\{j\}\)\\,\\big\]\_\{\+\}\\right\],\(6\)withw¯i​j=\|𝐘i−𝐘j\|/𝔼​\[\|𝐘i−𝐘j\|\]\\bar\{w\}\_\{ij\}=\|\\mathbf\{Y\}\_\{i\}\\\!\-\\\!\\mathbf\{Y\}\_\{j\}\|/\\mathbb\{E\}\[\|\\mathbf\{Y\}\_\{i\}\\\!\-\\\!\\mathbf\{Y\}\_\{j\}\|\]and\[⋅\]\+=max⁡\(0,⋅\)\[\\,\\cdot\\,\]\_\{\+\}=\\max\(0,\\cdot\)\. Pairs with score differences below1%1\\%of the max are filtered for numerical stability\. We setλm​s​e=20\\lambda\_\{mse\}\{=\}20,λb​i​n=10\\lambda\_\{bin\}\{=\}10,λf​i​n​e=3\\lambda\_\{fine\}\{=\}3,λg​l​o​b​a​l=2\\lambda\_\{global\}\{=\}2,λc​o​s=0\.5\\lambda\_\{cos\}\{=\}0\.5to balance gradient magnitudes\. The five terms supervise complementary failure modes \(ℒm​s​e\\mathcal\{L\}\_\{mse\}magnitude,ℒf​i​n​e\\mathcal\{L\}\_\{fine\}near\-tied Top\-KKorder atρ≤0\.2\\rho\\\!\\leq\\\!0\.2,ℒg​l​o​b​a​l\\mathcal\{L\}\_\{global\}signal\-to\-noise margin,ℒc​o​s\\mathcal\{L\}\_\{cos\}anti\-drift\) and the LOO study in[Section˜6](https://arxiv.org/html/2605.16360#S6)confirms they own nearly disjoint task families atρ=0\.1\\rho\{=\}0\.1, justifying targeted auxiliaries over interchangeable regularizers\.

### 4\.3Asynchronous scheduling and inference\-time pipeline

#### Stream\-level overlap\.

At inference,ℳl\\mathcal\{M\}\_\{l\}andℳs\\mathcal\{M\}\_\{s\}are dispatched to two independent CUDA streams \(separate GPUs in the dual\-GPU regime, distinct streams on the same device in the single\-GPU regime\);Φθ\\Phi\_\{\\theta\}is enqueued on the proxy stream as soon as𝐗\\mathbf\{X\}is ready, and the resulting𝐘^\\hat\{\\mathbf\{Y\}\}is consumed by a Top\-KKkernel on the target stream that materializes𝐌\\mathbf\{M\}in place\. The async\-budgettproxy\+tΦθ≤tprefill​\(ℳl\)t\_\{\\text\{proxy\}\}\+t\_\{\\Phi\_\{\\theta\}\}\\leq t\_\{\\text\{prefill\}\}\(\\mathcal\{M\}\_\{l\}\)holds empirically by an order of magnitude on every pair we evaluate; once𝐌\\mathbf\{M\}is applied the proxy KV cache is released and decode runs on the target alone \([Figure˜12](https://arxiv.org/html/2605.16360#A2.F12)\)\.

#### Training procedure\.

Φθ\\Phi\_\{\\theta\}is trained offline against a frozen oracle: we run one teacher\-forced KVZip pass per input to cache\(𝐗,𝐘\)\(\\mathbf\{X\},\\mathbf\{Y\}\), after which each gradient step costs a single mapper forward/backward\. We use AdamW \(lr2×10−42\{\\times\}10^\{\-4\}, weight decay10−410^\{\-4\}\) with a1,0001\{,\}000\-step linear warmup followed byReduceLROnPlateau\(factor0\.50\.5, patience33\), batch size88\(effective1212across44GPUs for the Qwen\-3\-32B mapper\), sliding\-window crops of2,0482\{,\}048tokens \(stride1,0241\{,\}024\), and grad\-clip1\.01\.0; loss coefficients are held fixed\. Each mapper trains for3030epochs without per\-pair tuning; per\-pair wall\-clock and full hyperparameters appear in[Appendix˜F](https://arxiv.org/html/2605.16360#A6)\.

## 5Experiments

We evaluate ProxyKV along three axes: \(i\) pruning accuracy across diverse benchmarks, \(ii\) prefilling latency reduction, and \(iii\) ablations of the mapper and loss design\. All experiments use PyTorch on two NVIDIA RTX PRO 6000 GPUs \(the 32B target additionally uses a multi\-GPU deployment for KVZip\)\.

### 5\.1Experiment setup

#### Model families\.

We use three intra\-family Target–Proxy pairs from 7B to 32B targets, training a dedicated HybridAxialMapper per pair:Qwen\-2\.5\(Yanget al\.,[2025b](https://arxiv.org/html/2605.16360#bib.bib19)\)\(7B / 1\.5B\),Llama\-3\.x\(Grattafioriet al\.,[2024](https://arxiv.org/html/2605.16360#bib.bib21)\)\(Llama\-3\.1\-8B / Llama\-3\.2\-1B; treated as one intra\-family pair\), andQwen\-3\(Yanget al\.,[2025a](https://arxiv.org/html/2605.16360#bib.bib20)\)\(32B / 4B\), which stresses ProxyKV at a∼\\sim8×8\\timestarget/proxy size ratio—the largest gap we evaluate\.

#### Baselines\.

We compare against three representative methods\.SnapKV\(Liet al\.,[2024c](https://arxiv.org/html/2605.16360#bib.bib24)\)is a training\-free heuristic that selects KV pairs from prompt\-end observation windows \(enhanced variant with global non\-uniform budget\)\.KVZip\+IAM\(Zhaoet al\.,[2025a](https://arxiv.org/html/2605.16360#bib.bib3)\)adapts IAM’s static layer\-to\-layer alignment to KVZip, isolating our learnable mapper’s contribution\.KVZip\(Kimet al\.,[2025](https://arxiv.org/html/2605.16360#bib.bib2)\)is the reconstruction\-based oracle, running a full secondary prefilling pass on the target to obtain precise attention scores\.

#### Data and metrics\.

The mapperΦθ\\Phi\_\{\\theta\}is trained on GSM8K\(Cobbeet al\.,[2021](https://arxiv.org/html/2605.16360#bib.bib28)\), SQuAD\(Rajpurkaret al\.,[2016](https://arxiv.org/html/2605.16360#bib.bib27)\), NIAH\(Kamradt,[2023](https://arxiv.org/html/2605.16360#bib.bib29)\), and the SCBench\(Liet al\.,[2024b](https://arxiv.org/html/2605.16360#bib.bib23)\)subsetsQA ENG,Summary,KV,Many Shot,Choice Eng, andPrefix Suffix; SCBench\.RepoQA is held out for zero\-shot evaluation, and we additionally evaluate on the full LongBench\(Baiet al\.,[2024](https://arxiv.org/html/2605.16360#bib.bib22)\)suite\. We report Accuracy, Exact Match \(EM\), and F1 across retention ratiosρ∈\{0\.1,0\.2,…,0\.9\}\\rho\\in\\\{0\.1,0\.2,\\ldots,0\.9\\\}, where smallerρ\\rhois more aggressive pruning\.

### 5\.2Main results

![Refer to caption](https://arxiv.org/html/2605.16360v1/x3.png)Figure 3:Aggregate accuracy on LongBench and SCBench for the Llama\-3\.1 and Qwen\-2\.5 families\. ProxyKV tracks the KVZip oracle within∼\\sim1\.51\.5pp atρ≥0\.5\\rho\\geq 0\.5\(the gap widens to∼\\sim55pp atρ≤0\.2\\rho\\leq 0\.2, where pruning bites hardest\) and outperforms heuristic SnapKV on SCBench\.Competitive performance across model families\.ProxyKV recovers∼\\sim98\.7%98\.7\\%of the KVZip oracle across all benchmarks and sparsity levels\.As shown in[Figure˜3](https://arxiv.org/html/2605.16360#S5.F3), ProxyKV \(blue\) tracks KVZip \(red\) on everyρ\\rho, and its margin over KVZip\+IAM \(green\) confirms that the learnable HybridAxialMapper is essential for bridging the head\-count / layer\-depth mismatches that static alignment cannot resolve\. SnapKV \(orange\) declines on SCBench because its local observation windows miss long\-range dependencies; it remains competitive only on tasks with strong syntactic or repetitive patterns \(RepoBench\-P,TREC\)\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x4.png)Figure 4:Zero\-shot transfer to held\-out SCBenchRepoQA\. ProxyKV \(blue\) tracks the KVZip oracle \(red\) within11–22pp on both targets\. Left: Qwen\-2\.5; right: Llama\-3\.1\.Robust zero\-shot generalization\.ProxyKV transfers zero\-shot to repository\-level reasoning, matching KVZip on the held\-out SCBench\.RepoQA task\.As shown in[Figure˜4](https://arxiv.org/html/2605.16360#S5.F4.5), ProxyKV maintains parity with KVZip on both Llama\-3\.1 and Qwen\-2\.5 despite RepoQA being excluded from the training mixture \(GSM8K,SQuAD,NIAH, and a subset of SCBench\)\. In theρ∈\[0\.2,0\.4\]\\rho\\in\[0\.2,0\.4\]band, ProxyKV outperforms SnapKV and KVZip\+IAM by a clear margin, evidencing the HybridAxialMapper’s structural alignment\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x5.png)Figure 5:Per\-dataset performance on the 16 English LongBench tasks \(Qwen\-2\.5\); the remaining 5 Chinese subsets are reported in[Appendix˜D](https://arxiv.org/html/2605.16360#A4)\. ProxyKV tracks the KVZip oracle and surpasses SnapKV on dense\-synthesis tasks\.Generalization to unseen datasets\.ProxyKV matches or exceeds the KVZip oracle on the 16 English LongBench tasks despite training on a different distribution\.As shown in[Figure˜5](https://arxiv.org/html/2605.16360#S5.F5), ProxyKV \(blue\) tracks KVZip \(red\) onNarrativeQAandMultiFieldQA\-Enwithout the secondary prefilling penalty, and maintains high precision at aggressive pruning ratios on2WikiMultiHopandQMSumwhere SnapKV \(orange\) saturates\. ProxyKV is near\-lossless on code\-intensiveLCCandRepoBench\-P; SnapKV remains competitive only on simple structured tasks likeTRECand fails on dense synthesis \(SAMSum\)\. The full 21\-subset breakdown including the 5 Chinese tasks is in[Appendix˜D](https://arxiv.org/html/2605.16360#A4)\.

### 5\.3Efficiency analysis

![Refer to caption](https://arxiv.org/html/2605.16360v1/x6.png)\(a\)Latency, Llama\-3\.1\-8B
![Refer to caption](https://arxiv.org/html/2605.16360v1/x7.png)\(b\)Latency, Qwen\-2\.5\-7B
![Refer to caption](https://arxiv.org/html/2605.16360v1/x8.png)\(c\)Memory, Llama\-3\.1\-8B
![Refer to caption](https://arxiv.org/html/2605.16360v1/x9.png)\(d\)Memory, Qwen\-2\.5\-7B

Figure 6:ProxyKV flattens the super\-linear latency curve of KVZip while paying a modest one\-time memory premium\. \(a–b\): prefilling latency; \(c–d\): peak GPU memory, across context length\.Latency scalability and memory trade\-off\.ProxyKV delivers up to3\.21×3\.21\\timesprefilling speedup at a modest16\.91%16\.91\\%–33\.72%33\.72\\%memory premium\.[Figure˜6](https://arxiv.org/html/2605.16360#S5.F6)\(a–b\) shows KVZip’s secondary prefilling pass scales super\-linearly up to 125k \(Llama\-3\.1\) and 170k \(Qwen\-2\.5\); ProxyKV flattens this curve by asynchronous offloading, yielding3\.21×3\.21\\timeson Llama\-3\.1\-8B and2\.53×2\.53\\timeson Qwen\-2\.5\-7B\. The cost is the joint target\+proxy footprint \([Figure˜6](https://arxiv.org/html/2605.16360#S5.F6)\(c–d\)\); the proxy KV cache is freed at end of prefill, so the premium is a transient prefill\-only peak \([Figure˜12](https://arxiv.org/html/2605.16360#A2.F12),[Appendix˜B](https://arxiv.org/html/2605.16360#A2)\)\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x10.png)Figure 7:Score–latency Pareto on LongBench \(5 representative tasks, total time = prefill \+ generation\)\. Numbers next to each marker indicate retention ratioρ\\rho\. ProxyKV \(blue\) dominates KVZip \(red\) on latency at everyρ\\rhoand dominates SnapKV \(orange\) on score above∼\\sim1\.31\.3s\.Score\-latency Pareto on real LongBench\.ProxyKV occupies the upper\-left of the score–latency frontier on every model\.On Llama\-3\.1\-8B atρ=0\.5\\rho\{=\}0\.5, ProxyKV reaches53\.953\.9at1\.511\.51s vs\. KVZip’s52\.252\.2at2\.792\.79s \(1\.85×1\.85\\timesat near\-equal accuracy\); SnapKV collapses to47\.947\.9at the same budget\. On Qwen\-3\-32B, ProxyKV holds55\.855\.8at7\.07\.0s vs\. KVZip’s11\.711\.7s \(1\.68×1\.68\\timesat equal score\)\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x11.png)Figure 8:ProxyKV prefill breakdown atρ=0\.3\\rho\{=\}0\.3\. Teacher and proxy run in parallel; the visible bottom block ismax⁡\(tteacher,tproxy\)\\max\(t\_\{\\text\{teacher\}\},t\_\{\\text\{proxy\}\}\), while the HybridAxialMapper accounts for4\.7%4\.7\\%–6\.8%6\.8\\%of wall time and remains a low\-order term as the target scales from 7B \(Qwen\-2\.5\) and 8B \(Llama\-3\.1\) up to 32B \(Qwen\-3\)\.GPU\-matched comparison and mapper overhead\.ProxyKV’s speedup is algorithmic, not hardware\-parallel\.Co\-locating proxy and teacher on a single GPU via CUDA streams still preserves∼\\sim1\.5×1\.5\\timesspeedup over KVZip on Llama\-3\.1\-8B and∼\\sim1\.4×1\.4\\timeson Qwen\-2\.5\-7B across 1\.6K–30K LongBench inputs \([Appendix˜A](https://arxiv.org/html/2605.16360#A1); 32B omitted as it must be sharded\)\. The HybridAxialMapper itself contributes only4\.7%4\.7\\%–6\.8%6\.8\\%of prefill wall time across 7B–32B targets \([Figure˜8](https://arxiv.org/html/2605.16360#S5.F8)\), and the ratio cannot grow with target size because mapper compute scales with proxy width while the dominant teacher–proxy block scales with target width\.

### 5\.4Stress test on RULER

![Refer to caption](https://arxiv.org/html/2605.16360v1/x12.png)Figure 9:ProxyKV tracks the KVZip oracle on RULER across all three target scales \(7B, 8B, 32B\)\. RULER 13\-task average score versus retention ratioρ\\rho\.ProxyKV preserves the KVZip\-oracle RULER accuracy from 7B to 32B targets\.[Figure˜9](https://arxiv.org/html/2605.16360#S5.F9)averages the 13 RULER\(Hsiehet al\.,[2024](https://arxiv.org/html/2605.16360#bib.bib32)\)subsets overρ∈\[0\.1,0\.9\]\\rho\\in\[0\.1,0\.9\]: ProxyKV tracks KVZip within11–22points on every target, including the∼\\sim8×8\\timesQwen\-3\-32B/Qwen\-3\-4B pair\. SnapKV collapses on multi\-needle/aggregation primitives atρ≤0\.3\\rho\\leq 0\.3, while KVZip\+IAM degrades*most*on Qwen\-3\-32B \(reaching0on*NIAH\-MK3*/*MV*/*FWE*atρ=0\.1\\rho\{=\}0\.1\), confirming that rigid layer\-to\-layer alignment binds harder as the target scales\. Per\-task breakdowns:[Appendix˜E](https://arxiv.org/html/2605.16360#A5)\.

## 6Ablation studies

We isolate the five loss terms, the three HybridAxialMapper stages, and the dominant loss coefficients\. Since[Figures˜3](https://arxiv.org/html/2605.16360#S5.F3)and[9](https://arxiv.org/html/2605.16360#S5.F9)cluster tightly atρ≥0\.5\\rho\\geq 0\.5, we focus onρ∈\{0\.1,0\.2\}\\rho\\in\\\{0\.1,0\.2\\\}\. Ablations use Llama\-3\.1\-8B / Llama\-3\.2\-1B \(each LOO variant retrained\); the cross\-family/cross\-scale check is the Qwen results in[Section˜5\.4](https://arxiv.org/html/2605.16360#S5.SS4)\. Component\- and hyperparameter\-LOOs and per\-task swing tables are in[Appendix˜C](https://arxiv.org/html/2605.16360#A3); we summarise the loss\-function LOO here\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x13.png)Figure 10:Loss LOO ablation, LongBench\-21 average\.ℒb​i​n\\mathcal\{L\}\_\{bin\}is the single most critical term at low retention, and the five loss swing\-supports are nearly disjoint\.Atρ∈\{0\.1,0\.2\}\\rho\\in\\\{0\.1,0\.2\\\}removingℒb​i​n\\mathcal\{L\}\_\{bin\}produces the largest drop; forρ≥0\.4\\rho\\geq 0\.4all six curves collapse into a11\-point band\. The per\-task LOO on LongBench\-21 \([Table˜2](https://arxiv.org/html/2605.16360#A3.T2)\) shows nearly disjoint ownership atρ=0\.1\\rho\{=\}0\.1:ℒb​i​n\\mathcal\{L\}\_\{bin\}for sharp\-cut \(PassageRetr\-En/Zh:−11\.0\-11\.0/−10\.5\-10\.5\),ℒm​s​e\\mathcal\{L\}\_\{mse\}for magnitude \(TREC:−15\.5\-15\.5\),ℒc​o​s\\mathcal\{L\}\_\{cos\}for anti\-drift \(DuReader:−3\.78\-3\.78\)\. The aggregate11\-point band atρ≥0\.4\\rho\\geq 0\.4hides2525\-point per\-task swings atρ≤0\.2\\rho\\leq 0\.2; component and coefficient LOOs reach the same conclusion \([Appendix˜C](https://arxiv.org/html/2605.16360#A3)\)\.

## 7Conclusion

ProxyKV offloads importance scoring to an asynchronous Small\-Model Proxy via the HybridAxialMapper and a ranking\-consistent hybrid loss, matching KVZip across Llama\-3\.1, Qwen\-2\.5, and Qwen\-3 \(∼\\sim98\.7%98\.7\\%mean recovery,77B–3232B targets\) at up to3\.21×3\.21\\timesprefilling speedup \(dual\-GPU;∼\\sim1\.5×1\.5\\timesshared single\-GPU\), sustained to170170k tokens; the ProxyKV–IAM gap*widens*with target scale, exposing rigid layer\-to\-layer alignment as the binding constraint that the learnable cross\-axial mapping removes\.

#### Limitations\.

ProxyKV requires intra\-family pairs and raises peak\-prefill GPU memory by16\.91%16\.91\\%–33\.72%33\.72\\%\(transient proxy activations released after prefill,[Appendix˜B](https://arxiv.org/html/2605.16360#A2)\); in vLLM/TGI continuous batching this peak still caps concurrent batch size\. Cross\-family transfer is future work\.

#### Broader Impact\.

ProxyKV lowers long\-context inference cost; potential misuse on sensitive corpora warrants safety monitoring and access controls\.

## References

- Y\. Bai, X\. Lv, J\. Zhang, H\. Lyu, J\. Tang, Z\. Huang, Z\. Du, X\. Liu, A\. Zeng, L\. Hou,et al\.\(2024\)Longbench: a bilingual, multitask benchmark for long context understanding\.InProceedings of the 62nd annual meeting of the association for computational linguistics \(volume 1: Long papers\),pp\. 3119–3137\.Cited by:[§5\.1](https://arxiv.org/html/2605.16360#S5.SS1.SSS0.Px3.p1.3)\.
- Z\. Cai, Y\. Zhang, B\. Gao, Y\. Liu, Y\. Li, T\. Liu, K\. Lu, W\. Xiong, Y\. Dong, J\. Hu,et al\.\(2024\)Pyramidkv: dynamic kv cache compression based on pyramidal information funneling\.arXiv preprint arXiv:2406\.02069\.Cited by:[§2](https://arxiv.org/html/2605.16360#S2.SS0.SSS0.Px1.p1.1)\.
- K\. Cobbe, V\. Kosaraju, M\. Bavarian, M\. Chen, H\. Jun, L\. Kaiser, M\. Plappert, J\. Tworek, J\. Hilton, R\. Nakano,et al\.\(2021\)Training verifiers to solve math word problems\.arXiv preprint arXiv:2110\.14168\.Cited by:[§5\.1](https://arxiv.org/html/2605.16360#S5.SS1.SSS0.Px3.p1.3)\.
- S\. Ge, Y\. Zhang, L\. Liu, M\. Zhang, J\. Han, and J\. Gao \(2023\)Model tells you what to discard: adaptive kv cache compression for llms\.arXiv preprint arXiv:2310\.01801\.Cited by:[§2](https://arxiv.org/html/2605.16360#S2.SS0.SSS0.Px1.p1.1)\.
- A\. Grattafiori, A\. Dubey, A\. Jauhri, A\. Pandey, A\. Kadian, A\. Al\-Dahle,et al\.\(2024\)The llama 3 herd of models\.External Links:2407\.21783,[Link](https://arxiv.org/abs/2407.21783)Cited by:[§5\.1](https://arxiv.org/html/2605.16360#S5.SS1.SSS0.Px1.p1.2)\.
- Y\. Gu, Z\. Jiang, J\. Jin, K\. Guo, Z\. Zhang, and X\. Xu \(2025\)AhaKV: adaptive holistic attention\-driven kv cache eviction for efficient inference of large language models\.arXiv preprint arXiv:2506\.03762\.Cited by:[§2](https://arxiv.org/html/2605.16360#S2.SS0.SSS0.Px1.p1.1)\.
- C\. Hsieh, S\. Sun, S\. Kriman, S\. Acharya, D\. Rekesh, F\. Jia, Y\. Zhang, and B\. Ginsburg \(2024\)RULER: what’s the real context size of your long\-context language models?\.arXiv preprint arXiv:2404\.06654\.Cited by:[§5\.4](https://arxiv.org/html/2605.16360#S5.SS4.p1.8)\.
- X\. Huang, L\. L\. Zhang, K\. Cheng, F\. Yang, and M\. Yang \(2024\)Fewer is more: boosting math reasoning with reinforced context pruning\.InProceedings of the 2024 Conference on Empirical Methods in Natural Language Processing,pp\. 13674–13695\.Cited by:[§2](https://arxiv.org/html/2605.16360#S2.SS0.SSS0.Px2.p1.1)\.
- S\. Jegou and M\. Jeblick \(2026\)KVzap: fast, adaptive, and faithful kv cache pruning\.External Links:2601\.07891,[Link](https://arxiv.org/abs/2601.07891)Cited by:[§2](https://arxiv.org/html/2605.16360#S2.SS0.SSS0.Px2.p1.1)\.
- J\. Jiang, F\. Wang, J\. Shen, S\. Kim, and S\. Kim \(2024\)A survey on large language models for code generation\.ACM Transactions on Software Engineering and Methodology\.Cited by:[§1](https://arxiv.org/html/2605.16360#S1.p1.1)\.
- G\. Kamradt \(2023\)Needle in a haystack \- pressure testing llms\.Note:GitHub repositoryExternal Links:[Link](https://github.com/gkamradt/LLMTest_NeedleInAHaystack)Cited by:[§5\.1](https://arxiv.org/html/2605.16360#S5.SS1.SSS0.Px3.p1.3)\.
- J\. Kim, J\. Kim, S\. Kwon, J\. W\. Lee, S\. Yun, and H\. O\. Song \(2025\)KVzip: query\-agnostic kv cache compression with context reconstruction\.arXiv preprint arXiv:2505\.23416\.Cited by:[§1](https://arxiv.org/html/2605.16360#S1.p2.1),[§5\.1](https://arxiv.org/html/2605.16360#S5.SS1.SSS0.Px2.p1.1)\.
- S\. Kim, S\. Shen, D\. Thorsley, A\. Gholami, W\. Kwon, J\. Hassoun, and K\. Keutzer \(2022\)Learned token pruning for transformers\.InProceedings of the 28th ACM SIGKDD conference on knowledge discovery and data mining,pp\. 784–794\.Cited by:[§2](https://arxiv.org/html/2605.16360#S2.SS0.SSS0.Px2.p1.1)\.
- W\. Kwon, Z\. Li, S\. Zhuang, Y\. Sheng, L\. Zheng, C\. H\. Yu, J\. Gonzalez, H\. Zhang, and I\. Stoica \(2023\)Efficient memory management for large language model serving with pagedattention\.InProceedings of the 29th symposium on operating systems principles,pp\. 611–626\.Cited by:[§1](https://arxiv.org/html/2605.16360#S1.p1.1),[§2](https://arxiv.org/html/2605.16360#S2.SS0.SSS0.Px3.p1.1)\.
- F\. Laitenberger, D\. Kopiczko, C\. G\. Snoek, and Y\. M\. Asano \(2025\)What layers when: learning to skip compute in llms with residual gates\.arXiv preprint arXiv:2510\.13876\.Cited by:[§2](https://arxiv.org/html/2605.16360#S2.SS0.SSS0.Px2.p1.1)\.
- H\. Li, Y\. Li, A\. Tian, T\. Tang, Z\. Xu, X\. Chen, N\. Hu, W\. Dong, Q\. Li, and L\. Chen \(2024a\)A survey on large language model acceleration based on kv cache management\.arXiv preprint arXiv:2412\.19442\.Cited by:[§1](https://arxiv.org/html/2605.16360#S1.p1.1)\.
- Y\. Li, H\. Jiang, Q\. Wu, X\. Luo, S\. Ahn, C\. Zhang, A\. H\. Abdi, D\. Li, J\. Gao, Y\. Yang,et al\.\(2024b\)Scbench: a kv cache\-centric analysis of long\-context methods\.arXiv preprint arXiv:2412\.10319\.Cited by:[§5\.1](https://arxiv.org/html/2605.16360#S5.SS1.SSS0.Px3.p1.3)\.
- Y\. Li, Y\. Huang, B\. Yang, B\. Venkitesh, A\. Locatelli, H\. Ye, T\. Cai, P\. Lewis, and D\. Chen \(2024c\)Snapkv: llm knows what you are looking for before generation\.Advances in Neural Information Processing Systems37,pp\. 22947–22970\.Cited by:[§1](https://arxiv.org/html/2605.16360#S1.p2.1),[§2](https://arxiv.org/html/2605.16360#S2.SS0.SSS0.Px1.p1.1),[§5\.1](https://arxiv.org/html/2605.16360#S5.SS1.SSS0.Px2.p1.1)\.
- Z\. Liu, A\. Desai, F\. Liao, W\. Wang, V\. Xie, Z\. Xu, A\. Kyrillidis, and A\. Shrivastava \(2023\)Scissorhands: exploiting the persistence of importance hypothesis for llm kv cache compression at test time\.Advances in Neural Information Processing Systems36,pp\. 52342–52364\.Cited by:[§2](https://arxiv.org/html/2605.16360#S2.SS0.SSS0.Px1.p1.1)\.
- S\. Minaee, T\. Mikolov, N\. Nikzad, M\. Chenaghlu, R\. Socher, X\. Amatriain, and J\. Gao \(2024\)Large language models: a survey\.arXiv preprint arXiv:2402\.06196\.Cited by:[§1](https://arxiv.org/html/2605.16360#S1.p1.1)\.
- P\. Rajpurkar, J\. Zhang, K\. Lopyrev, and P\. Liang \(2016\)Squad: 100,000\+ questions for machine comprehension of text\.arXiv preprint arXiv:1606\.05250\.Cited by:[§5\.1](https://arxiv.org/html/2605.16360#S5.SS1.SSS0.Px3.p1.3)\.
- G\. Xiao, Y\. Tian, B\. Chen, S\. Han, and M\. Lewis \(2023\)Efficient streaming language models with attention sinks\.arXiv preprint arXiv:2309\.17453\.Cited by:[§2](https://arxiv.org/html/2605.16360#S2.SS0.SSS0.Px1.p1.1)\.
- J\. Xu, J\. Pan, H\. Wang, Y\. Zhou, J\. Ye, Y\. Wang, and G\. Dai \(2025\)SpeContext: enabling efficient long\-context reasoning with speculative context sparsity in llms\.arXiv preprint arXiv:2512\.00722\.Cited by:[§2](https://arxiv.org/html/2605.16360#S2.SS0.SSS0.Px3.p1.1)\.
- A\. Yang, A\. Li, B\. Yang, B\. Zhang, B\. Hui, B\. Zheng, B\. Yu, C\. Gao, C\. Huang, C\. Lv,et al\.\(2025a\)Qwen3 technical report\.arXiv preprint arXiv:2505\.09388\.Cited by:[§5\.1](https://arxiv.org/html/2605.16360#S5.SS1.SSS0.Px1.p1.2)\.
- A\. Yang, B\. Yu, C\. Li, D\. Liu, F\. Huang, H\. Huang, J\. Jiang, J\. Tu, J\. Zhang, J\. Zhou,et al\.\(2025b\)Qwen2\. 5\-1m technical report\.arXiv preprint arXiv:2501\.15383\.Cited by:[§5\.1](https://arxiv.org/html/2605.16360#S5.SS1.SSS0.Px1.p1.2)\.
- Z\. Zhang, Y\. Sheng, T\. Zhou, T\. Chen, L\. Zheng, R\. Cai, Z\. Song, Y\. Tian, C\. Ré, C\. Barrett,et al\.\(2023\)H2o: heavy\-hitter oracle for efficient generative inference of large language models\.Advances in Neural Information Processing Systems36,pp\. 34661–34710\.Cited by:[§1](https://arxiv.org/html/2605.16360#S1.p2.1),[§2](https://arxiv.org/html/2605.16360#S2.SS0.SSS0.Px1.p1.1)\.
- Y\. Zhao, Z\. Li, and H\. Zhao \(2025a\)Iam: efficient inference through attention mapping between different\-scale llms\.InProceedings of the 63rd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),pp\. 19522–19533\.Cited by:[§1](https://arxiv.org/html/2605.16360#S1.p3.1),[§2](https://arxiv.org/html/2605.16360#S2.SS0.SSS0.Px3.p1.1),[§4\.1](https://arxiv.org/html/2605.16360#S4.SS1.SSS0.Px1.p1.8),[§5\.1](https://arxiv.org/html/2605.16360#S5.SS1.SSS0.Px2.p1.1)\.
- Y\. Zhao, Y\. Peng, C\. Nguyen, Z\. Li, X\. Wang, H\. Zhao, and X\. Fu \(2025b\)Smallkv: small model assisted compensation of kv cache compression for efficient llm inference\.arXiv preprint arXiv:2508\.02751\.Cited by:[§1](https://arxiv.org/html/2605.16360#S1.p3.1),[§2](https://arxiv.org/html/2605.16360#S2.SS0.SSS0.Px3.p1.1),[§4\.1](https://arxiv.org/html/2605.16360#S4.SS1.SSS0.Px1.p1.8)\.

## Appendix AReal\-world LongBench latency analysis

This appendix expands the score–latency Pareto plot of[Figure˜7](https://arxiv.org/html/2605.16360#S5.F7)\([Figure˜6](https://arxiv.org/html/2605.16360#S5.F6)\) with two robustness checks: a single\-GPU context\-length scan that constrains every method to the same hardware budget, and a per\-stage prefill breakdown isolating the HybridAxialMapper’s overhead\. SnapKV is included as the heuristic floor; KVZip\+IAM is omitted because its in\-place layer alignment offers no measurable runtime difference from KVZip on the prefill path\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x14.png)Figure 11:Single\-GPU prefill latency vs\. context length on real LongBench inputs atρ=0\.3\\rho\{=\}0\.3for Llama\-3\.1\-8B and Qwen\-2\.5\-7B\. The ProxyKV–KVZip gap widens monotonically with context length\.#### GPU\-count\-matched single\-GPU context scan\.

ProxyKV remains1\.3×1\.3\\times–1\.6×1\.6\\timesfaster than KVZip when every method is constrained to the same single GPU\.[Figure˜11](https://arxiv.org/html/2605.16360#A1.F11)controls for device count on the two targets that fit on a single GPU by running ProxyKV’s teacher and proxy concurrently via CUDA streams on the*same*GPU as KVZip and SnapKV, isolating the algorithmic speedup from hardware parallelism\. Across 22 LongBench samples per model spanning1\.61\.6K to3030K tokens—the upper bound at which all four methods fit on a single RTX PRO 6000 without OOM—ProxyKV preserves a∼\\sim1\.5×1\.5\\timesspeedup over KVZip on Llama\-3\.1\-8B and a∼\\sim1\.4×1\.4\\timesspeedup on Qwen\-2\.5\-7B across the entire context\-length range\. The matched\-device comparison is omitted for Qwen\-3\-32B because the 32B target must be sharded across multiple GPUs, making “1 GPU per method” undefined\. The per\-stage prefill breakdown that quantifies the mapper’s4\.7%4\.7\\%–6\.8%6\.8\\%wall\-time share is reproduced in[Figure˜8](https://arxiv.org/html/2605.16360#S5.F8)of the main text\.

## Appendix BMemory timeline of the dual\-GPU pipeline

This appendix complements the per\-context memory footprint of[Figure˜6](https://arxiv.org/html/2605.16360#S5.F6)\(c–d\) \([Figure˜6](https://arxiv.org/html/2605.16360#S5.F6)\) with a per\-GPU*memory–time*trace that locates the16\.91%16\.91\\%–33\.72%33\.72\\%memory premium in time\. The aggregate plots in the main text report the joint footprint at a single \(worst\-case\) instant; the trace below shows that this instant is confined to the prefill window and that the proxy\-side overhead vanishes during decode\. We sample per\-GPU usage every5050ms via NVML while running the Llama\-3\.1\-8B target on GPU 1 and the Llama\-3\.2\-1B proxy on GPU 2 over an88K\-token context, followed by3232decode steps\. The trace separates the pipeline into three phases: \(i\) weight loading, \(ii\) parallel prefill \(target and proxy concurrent\), and \(iii\) decode on the target only\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x15.png)Figure 12:Dual\-GPU memory timeline for ProxyKV \(Llama\-3\.1\-8B target on GPU 1, Llama\-3\.2\-1B proxy on GPU 2\)\. The orange band marks the prefill phase, the green band marks decode\. The proxy GPU jumps from3\.53\.5GB \(weights only\) to26\.726\.7GB at the prefill peak—driven almost entirely by the prefill\-time activation working set \(attention logits, intermediate projections, and short\-lived hidden states\), since the11B proxy’s KV cache itself is well under11GB at this context length—then drops back to3\.53\.5GB the instant the prefill activations and the proxy KV cache are released, and stays flat for the entire decode window\. The∼\\sim2323GB premium is therefore a transient prefill\-only peak, not a sustained cost\.#### Prefill\-only transient peak\.

The proxy\-side memory overhead is bounded to the prefill window and is released before decode starts\.As shown in[Figure˜12](https://arxiv.org/html/2605.16360#A2.F12), GPU 2 \(proxy\) climbs from3\.53\.5GB \(weights only\) to26\.726\.7GB during prefill\. The dominant contributor to this∼\\sim2323GB spike is*not*the proxy KV cache \(which, for a11B\-parameter model in1616\-bit precision, occupies under11GB even at this context length\) but the proxy’s prefill activation working set—attention logit matrices, intermediate Q/K/V/FFN projections, and the auxiliary tensors required by the HybridAxialMapper to extract the cross\-head feature𝐗\\mathbf\{X\}\. The instant prefill finishes, both the proxy KV cache and the activation working set are released, and GPU 2 returns to3\.53\.5GB for the entire decode phase, a∼\\sim7\.6×7\.6\\timesreduction\. GPU 1 \(target\) follows a parallel pattern: it peaks at39\.739\.7GB during prefill \(target weights\+\+target prefill activations\+\+full target KV cache\) and settles at20\.520\.5GB during decode after the pruning mask is applied\. To make this distinction explicit, the16\.91%16\.91\\%–33\.72%33\.72\\%memory premium reported in[Figure˜6](https://arxiv.org/html/2605.16360#S5.F6)is the*peak prefill*GPU footprint relative to KVZip, including activations; the corresponding theoretical\-only premium \(weights\+\+retained KV cache\) is much smaller\. We note that this peak is the load\-bearing quantity for serving deployments: in continuous\-batching engines such as vLLM and TGI, where prefill and decode share a common GPU memory pool, the prefill peak must be reservable for any concurrent request, so the maximum batch size is set by the peak rather than by the time\-averaged footprint\.

## Appendix CDetailed ablation analysis

This appendix complements the loss\-LOO summary in[Section˜6](https://arxiv.org/html/2605.16360#S6)with \(i\) the per\-task LOO swing tables across LongBench\-21 \([Tables˜1](https://arxiv.org/html/2605.16360#A3.T1)and[2](https://arxiv.org/html/2605.16360#A3.T2)\), \(ii\) the component LOO of the three HybridAxialMapper stages \([Figure˜13](https://arxiv.org/html/2605.16360#A3.F13),[Table˜3](https://arxiv.org/html/2605.16360#A3.T3)\), and \(iii\) the hyperparameter sensitivity sweep overλb​i​n,λm​s​e\\lambda\_\{bin\},\\lambda\_\{mse\}\([Figure˜14](https://arxiv.org/html/2605.16360#A3.F14)\)\.

Table 1:Qualitative overviewof each loss term’s intended role and the task families we hypothesised it would dominate prior to running the leave\-one\-out study\. The “Dominant tasks” column lists representative datasets where the term is expected to matter; the actual per\-task LOO winners at the diagnostic retention ratios are reported in[Table˜2](https://arxiv.org/html/2605.16360#A3.T2)and should be treated as the quantitative ground truth \(a few tasks change owner between the two tables, which is informative about how the loss landscape shifts withρ\\rhorather than a contradiction in our claims\)\. “Diag\.ρ\\rho” marks the retention range where the term measurably moves the score; entries outside that are within noise\. Dataset abbreviations: NQA = NarrativeQA, TQA = TriviaQA, HQA = HotpotQA, MFQA\-Zh = MultiField\-QA\-Zh, PR\-En/Zh = PassageRetr\-En/Zh, QMS = QMSum, GR = GovReport, SAM = SAMSum, MN = MultiNews, PC = PassageCount, MuS = MuSiQue, DR = DuReader\.Table 2:Per\-task most\-needed loss termon the LongBench\-21 suite\. For every task we report the LOO variant whose removal causes the largest score drop versus*Full*at the diagnostic retention ratiosρ=0\.1\\rho\{=\}0\.1andρ=0\.2\\rho\{=\}0\.2, with the magnitude of that drop \(Δ\\Delta, points\)\. Drops with magnitude≥3\\geq 3pts are bolded\. Tasks are grouped by which loss term owns them atρ=0\.1\\rho\{=\}0\.1, with a final*non\-diagnostic*block listing tasks whose five LOO variants all stay within±0\.5\\pm 0\.5pt of*Full*atρ=0\.1\\rho\{=\}0\.1\. Abbreviations follow[Table˜1](https://arxiv.org/html/2605.16360#A3.T1)\.TaskTop loss @ρ=0\.1\\rho\{=\}0\.1Δ0\.1\\Delta\_\{0\.1\}Top loss @ρ=0\.2\\rho\{=\}0\.2Δ0\.2\\Delta\_\{0\.2\}*ℒb​i​n\\mathcal\{L\}\_\{bin\}\-dominated \(sharp\-cut / single\-needle / multi\-hop chains\)*PassageRetr\-Enℒb​i​n\\mathcal\{L\}\_\{bin\}−11\.00\\mathbf\{\-11\.00\}ℒb​i​n\\mathcal\{L\}\_\{bin\}−2\.50\-2\.50PassageRetr\-Zhℒb​i​n\\mathcal\{L\}\_\{bin\}−10\.50\\mathbf\{\-10\.50\}ℒb​i​n\\mathcal\{L\}\_\{bin\}−3\.00\\mathbf\{\-3\.00\}MultiField\-QA\-Zhℒb​i​n\\mathcal\{L\}\_\{bin\}−7\.05\\mathbf\{\-7\.05\}ℒg​l​o​b​a​l\\mathcal\{L\}\_\{global\}−1\.52\-1\.52LSHTℒb​i​n\\mathcal\{L\}\_\{bin\}−5\.25\\mathbf\{\-5\.25\}ℒm​s​e\\mathcal\{L\}\_\{mse\}−6\.00\\mathbf\{\-6\.00\}2WikiMultiHopℒb​i​n\\mathcal\{L\}\_\{bin\}−4\.88\\mathbf\{\-4\.88\}ℒb​i​n\\mathcal\{L\}\_\{bin\}−3\.67\\mathbf\{\-3\.67\}MuSiQueℒb​i​n\\mathcal\{L\}\_\{bin\}−3\.62\\mathbf\{\-3\.62\}ℒm​s​e\\mathcal\{L\}\_\{mse\}−1\.33\-1\.33HotpotQAℒb​i​n\\mathcal\{L\}\_\{bin\}−2\.53\-2\.53ℒf​i​n​e\\mathcal\{L\}\_\{fine\}−0\.15\-0\.15LCC \(code\)ℒb​i​n\\mathcal\{L\}\_\{bin\}−2\.39\-2\.39ℒb​i​n\\mathcal\{L\}\_\{bin\}−2\.21\-2\.21GovReportℒb​i​n\\mathcal\{L\}\_\{bin\}−2\.13\-2\.13ℒg​l​o​b​a​l\\mathcal\{L\}\_\{global\}−0\.39\-0\.39MultiNewsℒb​i​n\\mathcal\{L\}\_\{bin\}−1\.32\-1\.32ℒg​l​o​b​a​l\\mathcal\{L\}\_\{global\}−0\.26\-0\.26*ℒm​s​e\\mathcal\{L\}\_\{mse\}\-dominated \(magnitudes drive aggregation\)*TRECℒm​s​e\\mathcal\{L\}\_\{mse\}−15\.50\\mathbf\{\-15\.50\}ℒb​i​n\\mathcal\{L\}\_\{bin\}−25\.50\\mathbf\{\-25\.50\}MultiField\-QA\-Enℒm​s​e\\mathcal\{L\}\_\{mse\}−1\.15\-1\.15ℒg​l​o​b​a​l\\mathcal\{L\}\_\{global\}−3\.48\\mathbf\{\-3\.48\}*ℒg​l​o​b​a​l\\mathcal\{L\}\_\{global\}\-dominated \(Top\-KK/ non\-Top\-KKboundary\)*Qasperℒg​l​o​b​a​l\\mathcal\{L\}\_\{global\}−1\.38\-1\.38ℒb​i​n\\mathcal\{L\}\_\{bin\}−3\.68\\mathbf\{\-3\.68\}NarrativeQAℒg​l​o​b​a​l\\mathcal\{L\}\_\{global\}−0\.55\-0\.55ℒm​s​e\\mathcal\{L\}\_\{mse\}−1\.12\-1\.12*ℒf​i​n​e\\mathcal\{L\}\_\{fine\}\-dominated \(intra\-Top\-KKorder; code completion\)*RepoBench\-Pℒf​i​n​e\\mathcal\{L\}\_\{fine\}−0\.67\-0\.67ℒg​l​o​b​a​l\\mathcal\{L\}\_\{global\}−1\.30\-1\.30*ℒc​o​s\\mathcal\{L\}\_\{cos\}\-dominated \(anti\-drift; long\-tail Chinese / multilingual\)*DuReaderℒc​o​s\\mathcal\{L\}\_\{cos\}−3\.78\\mathbf\{\-3\.78\}ℒb​i​n\\mathcal\{L\}\_\{bin\}−2\.90\-2\.90TriviaQAℒc​o​s\\mathcal\{L\}\_\{cos\}−1\.51\-1\.51ℒg​l​o​b​a​l\\mathcal\{L\}\_\{global\}−1\.33\-1\.33SAMSumℒc​o​s\\mathcal\{L\}\_\{cos\}−0\.55\-0\.55ℒg​l​o​b​a​l\\mathcal\{L\}\_\{global\}−0\.82\-0\.82*Non\-diagnostic atρ=0\.1\\rho\{=\}0\.1\(all five LOO variants within±0\.5\\pm 0\.5pt of*Full*\)*QMSum——ℒb​i​n\\mathcal\{L\}\_\{bin\}−0\.80\-0\.80PassageCount——ℒb​i​n\\mathcal\{L\}\_\{bin\}−2\.34\-2\.34VCSUM——ℒb​i​n\\mathcal\{L\}\_\{bin\}−0\.40\-0\.40#### Component ablation\.

The HybridAxialMapper is a three\-stage pipeline: \(i\) 1D Conv Stem, \(ii\) Time\-Axis Encoder, \(iii\) Head\-Axis Cross\-Attention\. We disable one stage at a time, replacing it with an identity \(Conv Stem, Cross\-Attention\) or a mean\-pool \(Time\-Axis Encoder\), keeping all other parameters unchanged\. The simplifying replacements preserve tensor shapes so that the loss objective is unmodified\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x16.png)Figure 13:Component leave\-one\-out on the six representative LongBench subsets, six\-task average vs\. retention ratioρ\\rho\. The discriminating regime isρ∈\{0\.1,0\.2\}\\rho\\in\\\{0\.1,0\.2\\\}, where “w/o Conv” incurs the largest drop \(−3\.67\-3\.67on the six\-task average\), followed by “w/o Time” \(−2\.38\-2\.38\) and “w/o Head” \(−1\.58\-1\.58\); all four configurations re\-converge onceρ≥0\.5\\rho\\geq 0\.5\. Per\-task drops are quantified in[Table˜3](https://arxiv.org/html/2605.16360#A3.T3)\.Table 3:Component LOO at the discriminating regimeρ∈\{0\.1,0\.2\}\\rho\\in\\\{0\.1,0\.2\\\}on the six representative LongBench subsets\.*Full*columns are absolute ProxyKV accuracy;*w/o*columns give the change \(Δ\\Delta\) when the corresponding HybridAxialMapper stage is disabled and re\-trained\. Drops with magnitude≥3\\geq 3pts are bolded\.The full design unlocks model usability atρ=0\.2\\rho\{=\}0\.2, where every stage is load\-bearing\.The complete HybridAxialMapper reaches33\.0333\.03on the selected\-six average atρ=0\.2\\rho\{=\}0\.2, recovering94\.5%94\.5\\%of the uncompressed reference \(34\.9434\.94atρ=1\.0\\rho\{=\}1\.0\) while discarding80%80\\%of the KV cache\. The per\-task evidence in[Table˜3](https://arxiv.org/html/2605.16360#A3.T3)concentrates the aggregate drop on TREC \(−19\.5\-19\.5/−15\.5\-15\.5/−8\.0\-8\.0atρ=0\.1\\rho\{=\}0\.1for w/o Conv/Time/Head\) and 2WikiMQA \(−4\.7\-4\.7w/o Conv\); the three\-stage architecture is most necessary on tasks that require sharp\-cut classification of vital tokens or multi\-hop alignment at aggressive compression\.

#### Hyperparameter sensitivity\.

We sweep the two dominant loss coefficientsλb​i​n\\lambda\_\{bin\}andλm​s​e\\lambda\_\{mse\}around their default values \(λb​i​n=10\\lambda\_\{bin\}\{=\}10,λm​s​e=20\\lambda\_\{mse\}\{=\}20\)\. For each sweep we hold all other coefficients fixed and retrain\. The remaining coefficients \(λf​i​n​e=3\\lambda\_\{fine\}\{=\}3,λg​l​o​b​a​l=2\\lambda\_\{global\}\{=\}2,λc​o​s=0\.5\\lambda\_\{cos\}\{=\}0\.5\) are an order of magnitude smaller and were not found to be performance\-critical in pilot studies\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x17.png)Figure 14:Hyperparameter sensitivity of the two dominant loss coefficients on the six representative LongBench subsets\. The default\(λb​i​n,λm​s​e\)=\(10,20\)\(\\lambda\_\{bin\},\\lambda\_\{mse\}\)=\(10,20\)leads atρ∈\{0\.1,0\.2\}\\rho\\in\\\{0\.1,0\.2\\\}, and all sweeps reconverge at high retention\.The default coefficients lead atρ∈\{0\.1,0\.2\}\\rho\\in\\\{0\.1,0\.2\\\}and the sweep is uninformative aboveρ=0\.4\\rho\{=\}0\.4\.In the high\-retention tailρ∈\[0\.4,0\.9\]\\rho\\in\[0\.4,0\.9\]the curves overlap within0\.60\.6average points and cannot rank the candidates\. Atρ∈\{0\.1,0\.2\}\\rho\\in\\\{0\.1,0\.2\\\}the spread widens \(up to1\.51\.5points\), with a∼\\sim2×2\\timesperturbation in either direction degrading by less than11point\. We attribute the residual robustness to the multi\-ratio nature ofℒb​i​n\\mathcal\{L\}\_\{bin\}: even atλb​i​n=5\\lambda\_\{bin\}\{=\}5the binary signal still receives roughly a quarter of the gradient budget\.

## Appendix DComplete LongBench results

[Figures˜15](https://arxiv.org/html/2605.16360#A4.F15),[16](https://arxiv.org/html/2605.16360#A4.F16),[17](https://arxiv.org/html/2605.16360#A4.F17),[18](https://arxiv.org/html/2605.16360#A4.F18),[19](https://arxiv.org/html/2605.16360#A4.F19)and[20](https://arxiv.org/html/2605.16360#A4.F20)report the full 21\-dataset LongBench evaluation for the Qwen\-2\.5\-7B, Llama\-3\.1\-8B, and Qwen\-3\-32B target families, using the same retention ratios and method ordering as in the main experiments\. The Qwen\-3\-32B panels \(paired with the dedicated Qwen\-3\-4B proxy\) span a∼\\sim8×8\\timestarget/proxy size ratio and complement the RULER stress test \([Section˜5\.4](https://arxiv.org/html/2605.16360#S5.SS4)\) and the per\-task RULER breakdown in[Appendix˜E](https://arxiv.org/html/2605.16360#A5), which together exercise the same retention rangeρ∈\{0\.1,…,0\.9\}\\rho\\in\\\{0\.1,\\ldots,0\.9\\\}as LongBench while emphasizing the long\-context primitives \(NIAH, multi\-key/value, variable tracking, word\-frequency aggregation\) most sensitive to scoring quality at aggressive compression\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x18.png)Figure 15:Complete LongBench results, Qwen\-2\.5, Part I\.#### Reading the panels\.

Each panel reports task accuracy as a function of the KV cache retention ratioρ∈\{0\.1,0\.2,…,0\.9\}\\rho\\in\\\{0\.1,0\.2,\\ldots,0\.9\\\}, with the four methods drawn in the same color and marker scheme as the main\-text figures:ProxyKV\(blue\),KVZip\(red\),KVZip\+IAM\(green\), andSnapKV\(orange\)\. Part I covers 11 single\-document QA, multi\-hop reasoning, and summarization datasets where global semantic dependencies dominate; Part II \([Figures˜16](https://arxiv.org/html/2605.16360#A4.F16)and[18](https://arxiv.org/html/2605.16360#A4.F18)\) covers the remaining 10 datasets, predominantly retrieval, classification, and code\-completion benchmarks where SnapKV’s local attention spikes are most competitive, together with two additional summarization tasks \(SAMSum,VCSum\) that we group with Part II to keep the per\-page panel layout balanced\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x19.png)Figure 16:Complete LongBench results, Qwen\-2\.5, Part II\.
#### Qwen\-2\.5\-7B target\.

On the 11 retrieval/classification/code panels in Part II, ProxyKV \(blue\) and KVZip \(red\) overlap on*TREC*,*TriviaQA*,*LSHT*, and*LCC*across all retention ratios, indicating that the proxy successfully recovers the layer\-wise key signal that drives label and span retrieval\. SnapKV \(orange\) is competitive only on*PassageRetrieval\-En/Zh*where the answer span is locally clustered, and KVZip\+IAM \(green\) trails on*PassageCount*and*SAMSum*where the static layer alignment cannot disambiguate near\-duplicate keys\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x20.png)Figure 17:Complete LongBench results, Llama\-3\.1, Part I\.
#### Llama\-3\.1\-8B target, Part I\.

The 11 single\-document QA, multi\-hop, and summarization panels reproduce the cross\-method ordering observed on Qwen\-2\.5: ProxyKV stays within11–22points of KVZip on every panel and the gap to SnapKV widens monotonically asρ\\rhodrops below0\.30\.3, with the largest separations on*HotpotQA*,*2WikiMultiHop*, and*MuSiQue*, the three multi\-hop tasks where multiple non\-adjacent keys must be retained simultaneously and SnapKV’s local\-window heuristic fails first\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x21.png)Figure 18:Complete LongBench results, Llama\-3\.1, Part II\.
#### Llama\-3\.1\-8B target, Part II\.

On the 10 retrieval/classification/code panels, ProxyKV again tracks KVZip within11–22points across allρ\\rho\. The most diagnostic panels are*PassageRetr\-En/Zh*and*LCC*: KVZip\+IAM \(green\) lags by2020–4040points atρ=0\.1\\rho\{=\}0\.1–0\.30\.3, confirming that rigid layer\-to\-layer alignment is the binding constraint when the target tokens are long\-tailed and unevenly distributed across layers, while ProxyKV’s hybrid axial mapping closes the gap entirely\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x22.png)Figure 19:Complete LongBench results, Qwen\-3\-32B, Part I\.
#### Qwen\-3\-32B target, Part I\.

The 11 single\-document QA, multi\-hop, and summarization panels test whether the HybridAxialMapper recipe holds when retrained on the Qwen\-3\-4B proxy paired with the much larger Qwen\-3\-32B target \(∼\\sim8×8\\timestarget/proxy size ratio, the largest in our setup\)\. The cross\-method ordering observed on Qwen\-2\.5 and Llama\-3\.1 reproduces here without exception: ProxyKV \(blue\) tracks KVZip \(red\) within11–22points across the entire retention rangeρ∈\[0\.3,0\.9\]\\rho\\in\[0\.3,0\.9\]on all four single\-doc QA panels \(*NarrativeQA*,*Qasper*,*MultiFieldQA\-En/Zh*\) and on every multi\-hop and summarization panel, with the two curves visually overlapping aboveρ=0\.5\\rho\{=\}0\.5\. The most diagnostic panels are the multi\-hop trio*HotpotQA*,*2WikiMultiHop*, and*MuSiQue*, where multiple non\-adjacent keys must be retained simultaneously: atρ=0\.3\\rho\{=\}0\.3ProxyKV matches KVZip to within11point \(58\.558\.5vs\.58\.958\.9,57\.257\.2vs\.55\.055\.0, and43\.343\.3vs\.43\.243\.2respectively\), while SnapKV \(orange\) trails by55–1111points and KVZip\+IAM \(green\) collapses by2020–2525points\. KVZip\+IAM exhibits the sharpest deterioration of the four methods on Qwen\-3\-32B as well: atρ=0\.1\\rho\{=\}0\.1it drops to11\.311\.3on*NarrativeQA*,13\.113\.1on*Qasper*,17\.217\.2on*MultiFieldQA\-En*, and14\.714\.7on*MultiNews*,1515–5050points below ProxyKV in each case, confirming that the rigid layer\-to\-layer alignment also fails to scale to the 32B target\. The summarization panels*GovReport*,*QMSum*, and*MultiNews*compress better than the QA panels: all four methods stay within a few points of the unpruned baseline down toρ=0\.3\\rho\{=\}0\.3, since extractive summary generation tolerates aggressive eviction better than span\-precise QA\.

#### Qwen\-3\-32B target, Part II\.

The 10 retrieval, classification, and code\-completion panels probe the regime where SnapKV’s local attention spikes are most competitive\. ProxyKV continues to track KVZip on the four classification and retrieval panels that drive the cross\-scale claim: on*TREC*and*LSHT*the two curves overlap within11point acrossρ∈\[0\.3,0\.9\]\\rho\\in\[0\.3,0\.9\], on*TriviaQA*all four methods saturate near9090% \(with even SnapKV atρ=0\.1\\rho\{=\}0\.1reaching91\.691\.6, indistinguishable from the oracle\), and*PassageRetrieval\-En/Zh*are the most diagnostic of all—ProxyKV preserves the unpruned99\.599\.5% accuracy down toρ=0\.3\\rho\{=\}0\.3on both languages while KVZip\+IAM collapses to10\.510\.5\(En\) and12\.012\.0\(Zh\) and SnapKV degrades to79\.579\.5and66\.066\.0, exposing the2020–9090\-point margin by which rigid layer alignment and local\-window heuristics fail when the answer span is sparsely distributed across non\-adjacent passages\. At the most aggressiveρ=0\.1\\rho\{=\}0\.1setting ProxyKV even surpasses the KVZip oracle on*PassageRetrieval\-En*\(92\.092\.0vs\.82\.582\.5\), suggesting that the proxy\-driven hybrid axial mapping captures a smoother score distribution than the target’s own attention under extreme compression\. The one panel where ProxyKV does not dominate is*LCC*: atρ=0\.1\\rho\{=\}0\.1SnapKV reaches58\.058\.0versus ProxyKV’s35\.635\.6and KVZip’s55\.955\.9, the only Part II task where local\-window scoring genuinely helps because code\-completion targets are locally clustered around the cursor;*RepoBench\-P*, by contrast, remains tightly clustered \(6262–6767\) for all four methods at everyρ\\rho, so the LCC effect does not generalize across the code suite\. The remaining panels \(*SAMSum*,*PassageCount*,*VCSum*\) are non\-discriminative—either saturated or near\-floor for every method—so the Part II reading at the 32B target scale extends the cross\-scale conclusion of the RULER stress test \([Section˜5\.4](https://arxiv.org/html/2605.16360#S5.SS4)\): ProxyKV recovers the KVZip oracle on the discriminative retrieval and classification primitives and never falls below the strongest baseline on any Part II task except*LCC*at the most aggressive retention\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x23.png)Figure 20:Complete LongBench results, Qwen\-3\-32B, Part II\.

## Appendix EPer\-task RULER breakdown

[Figures˜21](https://arxiv.org/html/2605.16360#A5.F21),[22](https://arxiv.org/html/2605.16360#A5.F22)and[23](https://arxiv.org/html/2605.16360#A5.F23)report the full per\-task RULER breakdown for each of the three target models, expanding the aggregate curves of[Figure˜9](https://arxiv.org/html/2605.16360#S5.F9)into 13 individual panels per target\. The 13 RULER subsets group into four primitive families: NIAH retrieval \(single S1\-S3, multi\-key MK1\-MK3, multi\-value MV, multi\-query MQ\), Variable Tracking \(VT\), word\-frequency aggregation \(Common Word Extraction CWE, Frequent Word Extraction FWE\), and Question Answering \(QA\-1, QA\-2\)\. Method coloring matches the main text:ProxyKV\(blue\),KVZip\(red\),KVZip\+IAM\(green\),SnapKV\(orange\);xx\-axis is the retention ratioρ∈\{0\.1,…,0\.9\}\\rho\\in\\\{0\.1,\\ldots,0\.9\\\}andyy\-axis is task\-specific accuracy\.

#### Reading the panels\.

Each row collects subsets that probe a single primitive: row 1 stress\-tests retrieval at increasing key\-needle counts, row 2 stresses multi\-value/multi\-query aggregation, row 3 stress\-tests cross\-document tracking and word\-frequency reasoning, and the final QA pair tests retrieval\-conditional generation\. We highlight three reading angles aligned with the main\-text claims in[Section˜5\.4](https://arxiv.org/html/2605.16360#S5.SS4): \(i\)NIAH multi\-key/multi\-value—SnapKV \(orange\) collapses belowρ=0\.3\\rho\{=\}0\.3where its observation\-window heuristic loses the global key signal, while ProxyKV \(blue\) tracks the KVZip oracle \(red\) within11–22points; \(ii\)CWE/FWE—KVZip\+IAM \(green\) trails by a wide margin, exposing the limit of static layer\-to\-layer alignment under long\-tail token statistics; \(iii\)QA\-1/QA\-2—all four methods cluster near saturation, confirming that the aggregate gap on[Figure˜9](https://arxiv.org/html/2605.16360#S5.F9)is driven by the retrieval and aggregation subsets above and not by QA\. Across all 13 subsets, ProxyKV and KVZip remain visually indistinguishable on Llama\-3\.1\-8B \(Llama\-3\.2\-1B proxy\), Qwen\-2\.5\-7B \(Qwen\-2\.5\-1\.5B proxy\),*and*Qwen\-3\-32B \(Qwen\-3\-4B proxy\), demonstrating that the same HybridAxialMapper recipe holds when retrained on each intra\-family pair and is not an artifact of any specific architecture or scale\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x24.png)Figure 21:Complete per\-task RULER results, Llama\-3\.1\-8B target\. 13 subsets×\\times4 methods×\\times9 retention ratios\. Companion to the aggregate curve in[Figure˜9](https://arxiv.org/html/2605.16360#S5.F9)\.
#### Llama\-3\.1\-8B target\.

The diagnostic panel on Llama is*NIAH\-MK3*\(multi\-key with three needles\): SnapKV \(orange\) collapses belowρ=0\.3\\rho\{=\}0\.3because its observation\-window scoring loses the global key signal once aggressive eviction begins, while ProxyKV \(blue\) tracks KVZip \(red\) to within11–22points down toρ=0\.1\\rho\{=\}0\.1\.*CWE*and*FWE*surface the limit of static layer alignment: KVZip\+IAM \(green\) drops by2020–3030points relative to ProxyKV across the entireρ\\rhorange, confirming that the long\-tail token distribution requires the per\-layer adaptive routing learned by the HybridAxialMapper\. The*QA\-1*/*QA\-2*panels saturate near the upper bound for all four methods, so the aggregate gap reported in[Figure˜9](https://arxiv.org/html/2605.16360#S5.F9)is driven almost entirely by retrieval and aggregation primitives, not by question answering\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x25.png)Figure 22:Complete per\-task RULER results, Qwen\-2\.5\-7B target with a Qwen\-2\.5\-1\.5B proxy\. Same axes and method ordering as[Figure˜21](https://arxiv.org/html/2605.16360#A5.F21)\.
#### Qwen\-2\.5\-7B target\.

On the Qwen\-2\.5\-7B target paired with the Qwen\-2\.5\-1\.5B proxy, the two curves are visually indistinguishable on 11 of the 13 subsets, with the only marginal separation on*NIAH\-S3*atρ=0\.1\\rho\{=\}0\.1\. The relative ordering of methods is otherwise identical to Llama: SnapKV remains the weakest baseline on every multi\-needle and multi\-value subset, and KVZip\+IAM exhibits the same long\-tail collapse on*CWE*/*FWE*\. The fact that the per\-family Qwen\-2\.5 result is not visibly better than the per\-family Llama result rules out the concern that ProxyKV’s accuracy is dominated by any one family’s structural prior — the same HybridAxialMapper recipe, separately retrained per intra\-family pair, recovers the KVZip oracle on both architectures\.

#### Qwen\-3\-32B target \(largest target/proxy size ratio\)\.

Qwen\-3\-32B is the most stringent target in our setup: paired with the Qwen\-3\-4B proxy, the target/proxy size ratio reaches∼\\sim8×8\{\\times\}, the widest gap we evaluate\. Despite this, ProxyKV \(blue\) overlays the KVZip oracle \(red\) on every one of the 13 panels with no visible separation at anyρ\\rho, including the previously discriminating*NIAH\-MK3*,*NIAH\-MV*, and*CWE*/*FWE*primitives\. The relative ordering of the two baselines, however,*flips*at this scale: KVZip\+IAM \(green\) is now the weakest method on the aggregate \([Figure˜9](https://arxiv.org/html/2605.16360#S5.F9)\) and exhibits the most severe collapse of any panel in our entire study—dropping to0\.00\.0on*NIAH\-MK3*,*NIAH\-MV*, and*FWE*atρ=0\.1\\rho\{=\}0\.1and recovering only pastρ=0\.5\\rho\{=\}0\.5on the multi\-needle and aggregation primitives—so the static layer\-to\-layer alignment becomes*less*, not more, viable at the 32B target\. SnapKV \(orange\) is correspondingly the stronger of the two baselines here: atρ=0\.3\\rho\{=\}0\.3it stays at8484–9696on the multi\-key panels \(*MK1*/*MK2*/*MK3*\) where KVZip\+IAM is at2929/5050/33, and on the 13\-task average atρ=0\.3\\rho\{=\}0\.3SnapKV reaches≈85\{\\approx\}85versus KVZip\+IAM’s≈43\{\\approx\}43\. The takeaway is that the per\-layer adaptive routing learned by the HybridAxialMapper does*not*commute with target scale—the same intra\-family pair retraining recipe that makes ProxyKV scale\-invariant also exposes the rigidity of static alignment more sharply at 32B\. This figure is the central empirical evidence for the scale\-up claim of[Section˜5\.4](https://arxiv.org/html/2605.16360#S5.SS4): ProxyKV is the only method that recovers the oracle at the 32B target scale\.

![Refer to caption](https://arxiv.org/html/2605.16360v1/x26.png)Figure 23:Complete per\-task RULER results, Qwen\-3\-32B target paired with a dedicated Qwen\-3\-4B proxy \(∼\\sim8×8\{\\times\}target/proxy size ratio, the largest in our setup\); ProxyKV stays within11–22points of KVZip on every subset\.

## Appendix FImplementation details and training protocols

To map the KV cache saliency from the target model to the proxy model, we employ a multi\-objective training protocol\. The specific configurations are detailed below\.

#### Optimization Strategy\.

The HybridAxialMapper is trained for 30 epochs using the AdamW optimizer with a learning rate of2×10−42\\times 10^\{\-4\}and a weight decay of1×10−41\\times 10^\{\-4\}\. We implement a linear warmup for the first1,0001\{,\}000steps, followed by aReduceLROnPlateauscheduler with a patience of 3 epochs and a decay factor of0\.50\.5\. We use a batch size of 8 and reserve5%5\\%of the samples for validation\. The two mappers for the Llama\-3\.1\-8B and Qwen\-2\.5\-7B targets are each trained on a single NVIDIA RTX PRO 6000 GPU \(the same hardware used for evaluation in[Figure˜6](https://arxiv.org/html/2605.16360#S5.F6)\) for approximately 14 hours; the mapper for the Qwen\-3\-32B target is trained on 4×\\timesNVIDIA A100 GPUs for approximately 10 hours\. Total mapper\-training compute is approximately6868GPU\-hours; counting preliminary and ablation runs, the full study consumed on the order of a few hundred GPU\-hours\. The trained mapper checkpoints are lightweight: approximately100100MB on disk for the Llama\-3\.1\-8B and Qwen\-2\.5\-7B mappers, and approximately150150MB for the Qwen\-3\-32B mapper, several orders of magnitude smaller than their respective target models\.

#### Multi\-Objective Loss\.

The training objective is the composite lossℒ=∑iλi​ℒi\\mathcal\{L\}=\\sum\_\{i\}\\lambda\_\{i\}\\mathcal\{L\}\_\{i\}\. The loss coefficients are empirically set toλb​i​n=10\.0\\lambda\_\{bin\}\{=\}10\.0,λm​s​e=20\.0\\lambda\_\{mse\}\{=\}20\.0,λf​i​n​e=3\.0\\lambda\_\{fine\}\{=\}3\.0,λg​l​o​b​a​l=2\.0\\lambda\_\{global\}\{=\}2\.0, andλc​o​s=0\.5\\lambda\_\{cos\}\{=\}0\.5\. The Multi\-Ratio Binary Loss \(ℒb​i​n\\mathcal\{L\}\_\{bin\}\) evaluates retention ratiosρ∈\{0\.05,0\.1,0\.15,0\.2,0\.3,0\.4,0\.5\}\\rho\\in\\\{0\.05,0\.1,0\.15,0\.2,0\.3,0\.4,0\.5\\\}using a power\-law decay to emphasize high\-compression regimes\. We additionally use Value\-Weighted MSE \(α=1\.5\\alpha\{=\}1\.5\) and Hierarchical Ranking losses to ensure numerical alignment and preserve the relative importance of attention scores\.

#### Data Mixture and Augmentation\.

Training data is curated from a diverse mixture of benchmarks, including GSM8K, SQuAD, NIAH, and multiple SCBench subsets \(QA ENG,Summary,KV,Many Shot,Choice Eng, andPrefix Suffix\)\. To accommodate long\-context inputs, we process sequences using a sliding window strategy with a crop length of2,0482\{,\}048tokens and a stride of1,0241\{,\}024\.

#### Training Dynamics\.

[Figure˜24](https://arxiv.org/html/2605.16360#A6.F24)illustrates the convergence of the total loss \(red\) alongside the Mass Reconstruction ratio \(blue\) over100,000100\{,\}000global training steps\. The consistent decline in total loss and the rapid stabilization of the Mass Reconstruction ratio near0\.950\.95demonstrate the effectiveness of our Multi\-Granularity Hybrid Loss in distilling high\-fidelity importance scores\. The smooth convergence further confirms that the mapper successfully captures structural attention invariants without overfitting to the specific linguistic patterns of the training mixture\.*Note that the Mass Reconstruction ratio \(≈0\.95\{\\approx\}0\.95\) is an intra\-training metric measuring how much of the target’s attention probability mass is captured by the predicted Top\-KKset, and is not directly comparable to the end\-to\-end task\-accuracy recovery \(≈98\.7%\{\\approx\}98\.7\\%\) reported in the main experiments\.*

![Refer to caption](https://arxiv.org/html/2605.16360v1/x27.png)Figure 24:Mass Reconstruction stabilizes near0\.950\.95within the first quarter of training, confirming that the Multi\-Granularity Hybrid Loss converges smoothly without overfitting\. Total loss \(red\) and Mass Reconstruction ratio \(blue\) over100,000100\{,\}000training steps\.