PAIR: Memory-Efficient and Fast Inference with Reading--Reasoning Fusion

Sejun Park; Haocheng Xia; Yongjoo Park

cs.CLOct 12, 2026

PAIR: Memory-Efficient and Fast Inference with Reading--Reasoning Fusion

Authors: Sejun Park, Haocheng Xia, Yongjoo Park

Organizations: Siebel School of Computing and Data Science, University of Illinois Urbana-Champaign, Urbana, IL 61801, USA

Abstract

For context-based question answering, using one large model for both evidence extraction and answer reasoning requires substantial memory and computation. Existing approaches trade accuracy for efficiency within a single model only. We propose a multi-model approach for efficient inference: our method, PAIR, employs two differently sized models for evidence extraction and answer reasoning, respectively. A smaller Reader model processes and retains the full context, while a larger Reasoner model sees only the question and generated sequence without directly attending to the context. These capabilities are fused at the logit level: at each token, both Reader and Reasoner produce a prediction, and their logits are combined into one, behaving as a single next-token-prediction model. Fine-tuning further boosts their performance. On HotpotQA, PAIR with a 32B Reasoner retains 96% of that same 32B model's full-context accuracy while using roughly one quarter of its KV memory. At 45.7K-token contexts, PAIR achieves up to 3.2 $\times$ higher end-to-end throughput than the full-context 32B model on the same GPU. PAIR also improves the accuracy-memory and accuracy-throughput trade-offs over evaluated token-reduction baselines. Its accuracy-preserving behavior generalizes across Qwen2.5, OLMo-3, and Gemma-3 model families.

Figures & tables

Panel of Figure 1 — Figure 1: A HotpotQA example motivating PAIR. The example uses Qwen2.5 7B and 32B as the Reader and Reasoner, respectively. (a) The small full-context model extracts both birth dates but reasons over them incorrectly, while the large no-context model lacks the relevant evidence and cannot determine the answer. PAIR correctly combines evidence extraction and answer reasoning to produce the answer. (b) Teacher-forced next-token distributions along the PAIR-generated sequence. The top row evaluates an evidence-extraction position whose continuation depends on context-specific information, while the bottom row evaluates an answer-reasoning position after the relevant evidence has appeared in the generated sequence.

Figure 2: Overview of PAIR. The Reader prefills the full context and retains its KV cache, while the Reasoner does not attend to the full context. At each decoding step, the Reader and Reasoner produce next-token predictions from the generated sequence, PAIR combines their logits, and the resulting token is appended to the sequence.

Figure 3: Accuracy–efficiency trade-offs across the four benchmarks. (a) F1 against KV memory normalized by the full-context 32B model. (b) F1 against time per completed answer, the inverse of aggregate throughput, where lower values indicate higher throughput. Arrows indicate the direction of improvement. Regions marked blue contain operating points that are worse than PAIR on both axes. Horizontal dashed lines indicate the full-context 32B model’s accuracy.

Figure 4: Accuracy generalization across model families and Reader sizes on HotpotQA. Each plot reports F1 for the full-context large model, PAIR, and the corresponding Reader at varying context lengths. Blue percentages below the PAIR curve report its KV memory relative to the corresponding full-context large model. For OLMo-3, the ratios exceed 100% because the 7B model uses MHA with 32 KV heads, while the 32B model uses GQA with 8 KV heads ( Ainslie et al., 2023 ) .

Figure 5: End-to-end throughput across batch sizes at four context lengths on the same 95-GiB GPU used for the main throughput experiments. F1 annotations report the accuracy of each method at the corresponding context length. PAIR remains close to the full-context 32B model in F1 while moving toward the 7B Reader in throughput as batch size increases.

Figure 6: Context-length scaling on HotpotQA. The left plot reports F1, the center plot reports end-to-end throughput at the largest tested feasible batch size, and the right plot reports time to first token at batch size one. PAIR remains close to the full-context 32B model in accuracy while its throughput advantage increases with context length.

Benchmark	Full-context 7B	Plain fusion	Reader FT	Reader + Reasoner FT	Full-context 32B
HotpotQA	0.5939	0.6720	0.6633	0.7178	0.7492
MuSiQue	0.3323	0.3716	0.3939	0.4380	0.4695
LooGLE	0.3878	0.3654	0.3658	0.4438	0.4876
LoCoMo	0.4243	0.4255	0.4916	0.4992	0.5583

Table 1: Fine-tuning ablation on the four primary benchmarks. Reader FT adapts only the Reader, while Reader + Reasoner FT adapts both models and corresponds to PAIR. Shaded columns show the full-context reference baselines; bold indicates the best result among the fusion variants.

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Position-wise KL divergence to the full-context 32B Teacher for the motivating HotpotQA example. The Reader, Reasoner, and PAIR are evaluated on the same Teacher-generated trajectory. The highlighted position corresponds to the decisive token, “Baz” shown in Section 3.1 .

Configuration	KL to Teacher	Gap closed
Reader 7B	0.2885	–
Reasoner 32B	0.6869	–
Plain fusion	0.2428	45%
PAIR 32B + 7B (Ours)	0.2338	53%
Oracle	0.1861	100%

Appendix

Table 2: KL divergence to the full-context 32B Teacher on the main HotpotQA configuration. Gap closed measures progress from the Reader toward the token-wise oracle.

Figure 8: Branch complementarity across Qwen2.5 model scales under teacher forcing. Rows vary the Reader size and columns the Reasoner size. The left plot reports token-wise oracle divergence $D_{\mathrm{oracle}}$ . The right plot reports $\min(D_{\mathrm{Read}},D_{\mathrm{Reason}})/D_{\mathrm{oracle}}$ . The model-size sweep includes the Teacher’s terminal token.

Benchmark	$\lambda=0.70$	$\lambda=0.75$	$\lambda=0.80$	$\lambda=0.85$
LoCoMo	0.461	0.486	0.495	0.499
HotpotQA	0.718	0.710	0.712	0.710

Appendix

Table 3: F1 under matched fusion-weight sweeps.

Benchmark	$D_{\mathrm{Read}}$	$D_{\mathrm{Reason}}$	$\hat{\lambda}$	KL optimum	Deployed $\lambda$
HotpotQA	0.385	0.882	0.70	0.75	0.70
MuSiQue	0.476	0.963	0.67	0.70	0.70
LooGLE	0.410	1.255	0.75	0.75	0.70
LoCoMo	0.273	1.325	0.83	0.85	0.85

Appendix

Table 4: Reliability estimates and teacher-forced fusion diagnostics. The teacher-KL optimum is the tested value of

\lambda

that minimizes divergence from the full-context distribution.

Method	Evaluated settings
Full-context 32B	full context
Reader 7B	full context
PAIR (32B + 7B)	full Reader context
SnapKV	15, 21.875, 30, 40%
ExpectedAttention	15, 21.875, 30, 40%
SpecPrefill	15, 21.875, 30, 40%

Appendix

Table 5: Configurations evaluated in the accuracy–memory comparison.

		HotpotQA		MuSiQue		LooGLE		LoCoMo
Method	Setting	KV / 32B	F1	KV / 32B	F1	KV / 32B	F1	KV / 32B	F1
Full-context 32B	full	100.0%	0.7492	100.0%	0.4695	100.0%	0.4876	100.0%	0.5583
Reader 7B	full	21.9%	0.5939	21.9%	0.3323	21.9%	0.3878	21.9%	0.4243
PAIR (32B + 7B)	–	24.8%	0.7178	25.2%	0.4380	22.7%	0.4438	23.2%	0.4992
SnapKV	5%	6.7%	0.6163	6.9%	0.2913	6.6%	0.2836	13.0%	0.2756
SnapKV	10%	11.6%	0.6564	11.8%	0.3436	11.6%	0.3226	18.0%	0.3157

Appendix

Table 6: Raw accuracy–memory operating points across the four primary benchmarks. All non-quantized configurations use fp16 KV.

Method	First	Later	Later $-$ First
Full-context 32B	0.471	0.495	+0.042
Full-context 7B	0.372	0.392	+0.007
PAIR	0.420	0.453	+0.043
SnapKV	0.455	0.341	-0.115
SpecPrefill	0.422	0.368	-0.051
ExpectedAttention	0.295	0.321	+0.013

Appendix

Table 7: First-question and later-question accuracy on the 65 LooGLE conversations with at least two questions. Differences are computed within conversation before averaging.

Method	Setting	Batch	All	Complete	Batch	All	Complete	Batch	All	Complete	Batch	All	Complete
		HotpotQA			MuSiQue			LooGLE			LoCoMo
Full-context 32B	full	12	0.560	0.560	10	0.542	0.547	3	0.154	0.154	$\geq 4$	0.347	0.347
Full-context 7B	full	80	1.712	1.760	48	2.302	2.302	16	0.458	0.458	$\geq 8$	1.159	1.159
PAIR (32B + 7B)	–	20	1.003	0.993	16	1.041	1.041	6	0.238	0.238	$\geq 8$	0.614	0.614
SnapKV	5%	$\geq 96$	0.880	0.880	$\geq 96$	0.963	0.963	24	0.443	0.443	$\geq 8$	0.874	0.874
SnapKV	10%	$\geq 80$	0.749	0.810	$\geq 80$	0.856	0.863	16	0.375	0.375	$\geq 8$	0.884	0.884

Appendix

Table 8: Raw accuracy–throughput operating points across the four primary benchmarks. Throughput is measured in completed answers per second.

Configuration	GPUs	Batch	Answers/s
Ours, one GPU	1	4	0.5896
Two independent one-GPU Ours replicas	2	4 each	1.1793
Ours, branch-parallel	2	16	1.2195
Full-context large model, one GPU	1	2	0.3868
Two independent full-context large model replicas	2	2 each	0.7736

Appendix

Table 9: Aggregate throughput for one- and two-GPU execution on LooGLE.

Configuration	Batch	First answer	Later turn
Full-context large model, one GPU	2	12.10 s	2.47 s
PAIR, one GPU	4	9.24 s	3.14 s
PAIR, two GPUs	4	8.31 s	2.34 s

Appendix

Table 10: Per-request latency on LooGLE. The first answer includes long-document prefill; later turns reuse the accumulated context state.

Model pair	Context	Full-context Large model	PAIR (ours)	Full-context Reader model
Qwen2.5 32B+7B	6.1k	0.7492	0.7178	0.5939
	11.8k	0.7232	0.6885	0.5969
	18.0k	0.7029	0.6828	0.5916
	23.3k	0.6893	0.6718	0.5554
Qwen2.5 32B+3B	6.1k	0.7492	0.6708	0.5432
	11.8k	0.7232	0.6550	0.5057

Appendix

Table 11: Raw HotpotQA F1 for the cross-family generalization experiment. All configurations use the same 600 questions and context construction.

Context	Positional setting	Reference(32B) $B$	Reference ans/s	PAIR $B$	PAIR ans/s	Reader(7B) $B$	Reader ans/s
6.1k	native	10	0.815	16	1.973	64	3.141
11.8k	native	6	0.432	10	1.116	32	1.729
18.0k	native	4	0.279	6	0.745	24	1.129
23.3k	native	3	0.197	6	0.592	24	0.814
29.2k	YaRN $\times 4$	2	0.145	4	0.415	16	0.621
34.8k	YaRN $\times 4$	2	0.119	4	0.367	16	0.491

Appendix

Table 12: Raw operating points for the context-length throughput experiment. Each method is evaluated at its largest tested feasible batch size using FKV.

Benchmark	32B fp16	32B int8	32B int4	32B int2	PAIR fp16
HotpotQA	0.7492 (100%)	0.7625 (53.9%)	0.7386 (28.6%)	0.3994 (16.8%)	0.7178 (24.8%)
MuSiQue	0.4695 (100%)	0.4793 (53.9%)	0.4571 (28.5%)	0.1761 (17.0%)	0.4380 (25.2%)
LooGLE	0.4876 (100%)	0.4927 (53.5%)	0.4828 (28.5%)	0.1538 (16.0%)	0.4438 (22.7%)
LoCoMo	0.5583 (100%)	0.5631 (53.6%)	0.5584 (28.6%)	0.1729 (16.0%)	0.4992 (23.2%)

Appendix

Table 13: F1 and resident KV memory for KV quantization of the full-context 32B model. Parentheses show KV memory relative to the fp16 full-context model. PAIR is included as a reference operating point.

Benchmark	PAIR fp16	PAIR + Reader-int8	Resident KV
HotpotQA	0.7178	0.7173	–
MuSiQue	0.4380	0.4411	–
LoCoMo	0.4992	0.4900	–
LooGLE	0.4438	0.4364	14.3%

Appendix

Table 14: Effect of applying int8 KV quantization to the Reader in PAIR. Resident KV is reported only where the canonical resident probe was run for the composed configuration.

Configuration	KiB / original context token	F1
SnapKV 15% + int8	30.30	0.6839
ExpectedAttention 15% + int8	31.11	0.5521
SnapKV 21.9% + int8	39.62	0.6870
ExpectedAttention 21.9% + int8	40.45	0.6026
PAIR + Reader-int8	46.79	0.7173
SnapKV 30% + int8	49.92	0.7023

Appendix

Table 15: Matched-byte comparison on HotpotQA. Cache cost is measured from batched accuracy runs and reported as KiB per original context token.

Configuration	KiB / original context token	F1
SnapKV 15% + int8	4.56	0.3261
ExpectedAttention 15% + int8	4.60	0.2992
SnapKV 21.9% + int8	6.46	0.3589
ExpectedAttention 21.9% + int8	6.47	0.3063
PAIR + Reader-int8	6.90	0.4364
SnapKV 30% + int8	8.65	0.3557

Appendix

Table 16: Matched-byte comparison on LooGLE. Cache cost is measured from batched accuracy runs and reported as KiB per original context token.

Memory-Efficient Inference

Memory-Efficient Inference reduces the memory needed to run neural networks while balancing accuracy, latency, throughput, and energy use.

See all 167 papers

PI-Mem: Pushing Long-Context Reasoning to 3.6M Tokens with Parallel-Iterative Memory
Aug 4, 2026Dawei Liu, Haixu Song, Shuang Cheng +9Memory-Augmented Language Models Long-Context QA
Shanghai Jiao Tong University · Shanghai Artificial Intelligence Laboratory · Tsinghua University +3
PARSER: Read in Parallel, Reason in Depth for Long-Context LLM Agents
Sep 6, 2026Kun Li, Zexuan Qiu, Tianhua Zhang +2Long-Context Inference Multi-Agent LLMs
The Chinese University of Hong Kong
Cache2LoRA: Understanding and Improving Context Internalization
Oct 16, 2026Yi Fang, Gaurav Srivastava, Zhenyu Bi +2LLM Context Compression Knowledge Augmentation for Language Models
Department of Computer Science, Virginia Tech, Blacksburg, VA, USA · Department of Biomedical Data Science, Stanford University, Stanford, CA, USA

Aug 4, 2026J/K move · Enter open · S save

PI-Mem: Pushing Long-Context Reasoning to 3.6M Tokens with Parallel-Iterative Memory

Dawei Liu, Haixu Song, Shuang Cheng, Shijie Wang+8

Shanghai Jiao Tong University · Shanghai Artificial Intelligence Laboratory · Tsinghua University+3

Memory-Augmented Language Models Long-Context QA LLM Inference Acceleration

Open paper AskPDF arXiv

Long-context reasoning remains a critical bottleneck for large language models, as recent recurrent-memory approaches face two inherent challenges: sequential chunk-wise updates can overwrite early critical evidence with later irrelevant content, and serial inter-chunk dependencies limit parallelism and cause latency to increase with context length. To address these issues, we propose PI-Mem (Parallel-Iterative Memory), a mechanism that processes all chunks in parallel and iteratively refines a shared memory over a bounded number of turns. In each turn, PI-Mem reads all chunks in parallel conditioned on the current memory, selects new or complementary evidence from each chunk, and merges the selected evidence into a compact shared memory for the next turn. To discourage redundant turns, we optimize the workflow through reinforcement learning with an auxiliary turn-efficiency reward, enabling the model to adaptively exit once sufficient evidence has been accumulated. We evaluate PI-Mem with Qwen3.5-35B-A3B and Qwen2.5-7B on the HotpotQA benchmark across context lengths up to 3.6 million tokens and find that it outperforms the recurrent-memory baseline by +6.25 and +7.81 absolute points while achieving 6.1

\times

and 2.1

\times

inference speedups, respectively. These results demonstrate that PI-Mem breaks the accuracy--efficiency trade-off in long-context reasoning and provides a scalable approach to complex multi-hop question answering over extremely long documents.

Abstract

Memory-Efficient Inference

Explore similar work

PI-Mem: Pushing Long-Context Reasoning to 3.6M Tokens with Parallel-Iterative Memory

PARSER: Read in Parallel, Reason in Depth for Long-Context LLM Agents

Cache2LoRA: Understanding and Improving Context Internalization