Task-Oriented Key-Layer KV Communication for Efficient Latent Multi-Agent Collaboration

1 Introduction

• Text-based collaboration repeatedly converts internal reasoning to text and received text back into model-internal representations.Multi-agent collaboration is used to combine complementary capabilities across tasks such as science reasoning, code generation, and embodied intelligence.Source: §1 · p. 1
• Existing latent communication methods focus on sender-side state fidelity, which need not align with the receiver-side task completion that collaboration seeks.The paper reframes the objective as receiver-side task sufficiency.Source: §1 · p. 1
• Task-oriented communication should preserve receiver performance with minimal communication volume and make transmitted states directly reusable to avoid redundant computation.These are the paper’s three stated design properties: task effectiveness, communication efficiency, and computational efficiency.Source: §1 · pp. 1–2
• KITE selects a model-specific key layer using receiver trajectory distortion, then communicates only its evolving latent working memory and reuses that layer for autoregressive latent reasoning.This avoids recomputing preceding Transformer layers.Source: §1 · p. 2
• 28–36× lower communication volume, up to 3× faster end-to-end inference, and up to 23.3 percentage points higher accuracy than full-layer KV communication.Experiments span seven benchmarks, two model families, and three model scales.Source: §1 · p. 2

2.1 LLM Communication in Latent Space

• KITE differs from latent-space methods such as CIPHER, AC, ThoughtComm, Interlat, C2C, and LatentMAS by transmitting only one task-effective key layer’s latent working memory.LatentMAS communicates layer-wise KV cache accumulated during prompt prefill and latent reasoning.Source: §2.1 · p. 2

2.2 KV Cache Optimization

• KV cache optimization methods address sharing, layer reuse, compression, or explicit shared-text reuse; these approaches are described as orthogonal to KITE or distinct from its evolving latent-memory communication.The paper says single-model cache optimizations can be combined with KITE.Source: §2.2 · p. 2
• KVComm selects layers by attention during receiver prefill, whereas KITE supports sender latent reasoning, targets receiver-side task utility, and uses the communication layer as the latent-reasoning entry point.This alignment is intended to reduce intra-agent computation.Source: §2.2 · p. 2

2.3 Layer Skipping in Inference

• Layer-skipping methods reduce inference cost by avoiding layers, while KITE’s latent injection skips layers not directly involved in inter-agent communication.The related work includes static pruning, layer collapse, and input-adaptive routers or execution programs.Source: §2.3 · p. 3

3.1 Sequential Multi-Agent Collaboration

• Four same-LLM agents collaborate sequentially as planner, critic, refiner, and solver; intermediate agents exchange KV cache, and only the solver decodes a natural-language output.The order is A1 → A2 → A3 → A4.Source: §3.1 · p. 3

3.2 KV Cache as Latent Working Memory

• Each non-terminal agent performs latent autoregressive reasoning without explicit decoding, and its prefill and reasoning states accumulate in the KV cache as evolving latent working memory.The reasoning steps generate last-layer hidden representations.Source: §3.2 · p. 3
• An agent’s passed memory concatenates inherited memory with its prefill and reasoning KV cache, and is forwarded to the next agent.The concatenation is along the sequence dimension; the planner starts with empty inherited memory.Source: §3.2 · p. 3

4 Method

• KITE transmits only a selected key layer’s KV cache and skips preceding layers during latent reasoning.This localizes both inter-agent communication and intra-agent computation.Source: §4 · p. 3

4.1 Key-layer KV Selection and Communication

• Receiver trajectory distortion serves as a label-free proxy for layer utility, comparing single-layer KV trajectories with the full-KV reference under teacher forcing.Distortion is measured by mean squared error over decoding steps and hidden dimensions; lower distortion indicates better preservation of receiver behavior.Source: §4.1 · pp. 3–4
• KITE ranks layers by average trajectory distortion across a two-domain calibration set, selects the highest-ranked layer, and reuses it downstream with subset size |S| = 1.The calibration set uses medical reasoning and mathematical reasoning datasets.Source: §4.1 · p. 4
• Each sender transmits only the selected layer’s latent working memory, making it the single communication path between agents.The selection objective balances receiver task effectiveness against minimizing communication volume, though direct optimization is intractable because it can require task feedback and combinatorial subset evaluation.Source: §4.1 · pp. 3–4

4.2 Layer Aligned Latent Injection

• In layer-aligned latent reasoning, agents inject the last hidden states into the key layer and propagate forward only from that layer to the final layer.This supports autoregressive reasoning without decoding intermediate latent states into explicit tokens.Source: §4.2 · pp. 4–5

4.3 Complexity Analysis

• KITE communicates 1/L of LatentMAS’s volume and reduces computation in prefill and generation stages, with savings dependent on sequence lengths, hidden dimension, and key-layer position.The passage gives per-agent prefill savings of O((L−1)cnd) and separate generation-stage savings for intermediate and solver agents.Source: §4.3 · p. 5

5.1 Experimental Setup

• The evaluation spans seven benchmarks, three models from two LLM families, and accuracy, communication volume, inference time, and output-token usage.Calibration uses 300 MedQA and 50 GSM8K samples; latent reasoning steps and maximum output lengths vary by task difficulty.Source: §5.1 · p. 5

5.3 Selection Strategy Effectiveness

• Trajectory-distortion scores correlated with single-layer accuracy at −0.780 for Qwen3-4B and −0.855 for Qwen3-8B; key-layer communication also outperformed attention-selected and random layers.The correlations remained strong in- and out-of-domain, and the performance gaps were larger on challenging tasks than on easier ones.Source: §5.3 · pp. 5–6, Table 2 · p. 6, Figure 3 · p. 6

5.2 Main Results

• 28×–36× lower communication volume than LMAS, up to 3× inference speedup, and up to 23.3 percentage points higher accuracy than the stronger latent baseline were reported for KITE.KITE consistently outperformed latent communication baselines in accuracy across model–benchmark combinations; the comparison covers seven benchmarks and three models.Source: §5.2 · p. 5, Table 1 · p. 6

5.4 Optimal Number of Communicated KV Layers

• Increasing the number of communicated top-ranked KV layers generally degrades task performance; communicating only the top-ranked key layer achieves the best overall accuracy.KITE therefore uses a selected subset size of one to jointly optimize task performance and communication efficiency.Source: §5.4 · pp. 6–7, Figure 4 · p. 7

5.5 Effect of Latent Injection Position

• As injection moves deeper, latent reasoning time decreases while accuracy fluctuates; injecting at key layer 17 gives competitive accuracy at the lowest reasoning time.Middle-layer injection generally yields higher accuracy than shallow-layer injection.Source: Figure 5 · p. 7, §5.5 · p. 7

5.6 Analysis of the Performance Gains

• During HumanEval+ decoding, KITE maintains top-1 token probability around 90% and cumulative top-5 probability near 1, while full-layer communication falls below 50% and 70%, respectively.The authors interpret the more dispersed full-layer output distribution as greater uncertainty and suggest key-layer communication may suppress redundant or conflicting information.Source: Figure 6 · p. 7, §5.6 · p. 7

6 Conclusion

• KITE selects a task-effective key layer using receiver trajectory distortion, communicates its latent working memory, and uses that layer to begin latent reasoning.The conclusion describes KITE as training-free.Source: §6 · p. 7
• Across diverse benchmarks and models, KITE reduces communication volume by 28–36× and accelerates end-to-end inference by up to 3×..Source: §6 · p. 7
• KITE improves accuracy by up to 23.3 percentage points over full-layer KV communication..Source: §6 · p. 7

A Experiment Setups

• The section introduces additional experimental, dataset, and implementation details.Source: A Experiment Setups · p. 10

• MedQA tests biomedical knowledge, clinical reasoning, and diagnostic decision-making through medical licensing-exam questions.Source: §A.1 · p. 10
• GSM8K, AIME24, and AIME25 assess multi-step mathematical reasoning, from grade-school word problems to competition- and Olympiad-style problems.Source: §A.1 · p. 10
• GPQA-Diamond contains 198 expert-written physics, biology, and chemistry questions designed to require graduate-level knowledge and multi-step scientific reasoning.Source: §A.1 · p. 10
• ARC-Easy and ARC-Challenge cover school-level science questions, with ARC-Challenge generally requiring more reasoning and knowledge integration.Source: §A.1 · p. 10
• MBPP-Plus pairs short programming tasks with function signatures and test cases, adding more extensive and rigorous tests to MBPP.Source: §A.1 · p. 10
• HumanEval-Plus evaluates Python implementations against strengthened tests, including additional tests that can expose solutions passing the original HumanEval tests.Source: §A.1 · p. 10

A.2 Implementation Details

• The experiments use each model’s official chat template and special tokens and run on an NVIDIA RTX 5880 Ada Generation GPU with 48 GB of memory.Source: §A.2 · p. 10
• Noncoding benchmark accuracy is determined by normalized answer matching, with exact option-letter matching for multiple-choice tasks and numerical equality for numerical tasks.Unparseable numerical predictions count as incorrect.Source: §A.2 · p. 10
• Code-generation correctness is tested by executing generated code against benchmark tests in a sandbox with a 10-second timeout; all tests must pass without runtime errors.Source: §A.2 · p. 10

B Positional Embedding Coherence

• Receiver positions are shifted by the incoming memory length at every layer, keeping a consistent positional frame while only the selected key layer contains sender KV states.Receiver tokens begin at position c, and subsequent generated states continue from c + n.Source: B Positional Embedding Coherence · p. 10

• For each calibration size, the study uses proportional stratified sampling, repeats sampling three times, ranks layers, and evaluates the highest-ranked layer on MedQA.Source: §C.1 · pp. 10–11
• All evaluated calibration sample sizes achieve the same MedQA accuracy across three backbone models, suggesting key-layer selection is robust to calibration-set size.Source: §C.1 · p. 11, Figure 7 · p. 11

C.2 Effect of Trajectory Distortion Metrics

• MSE provides the strongest average correlations between layer-ranking scores and task accuracy across the three models, while KL divergence gives conflicting positive correlations on Qwen3-8B.The comparison evaluates KL divergence, cosine similarity, and MSE across model–task combinations.Source: §C.2 · p. 11, Table 3 · p. 11
• MSE is reported as a stable and effective metric for measuring receiver trajectory distortion and identifying task-effective KV layers.Source: §C.2 · p. 11

D Prompts

• Planner, critic, and refiner prompts remain the same across tasks, whereas solver prompts are adapted to each task’s required output format.Source: D Prompts · p. 11, Figure 8 · p. 11, Figure 9 · p. 11