BatchParScale: Sparse Batch-Mixing for Sequence Models

Hans Peter Lyngsøe Raaschou-jensen; M. Onurcan Kaya; Desmond Elliott

cs.LGOct 18, 2026

BatchParScale: Sparse Batch-Mixing for Sequence Models

Authors: Hans Peter Lyngsøe Raaschou-jensen, M. Onurcan Kaya, Desmond Elliott

Organizations: University of Copenhagen · Technical University of Denmark

Abstract

Recent test-time scaling methods improve model performance by increasing inference-time compute. One practical form of this is parallel scaling, where a model samples multiple independent reasoning streams and aggregates their answers using methods like voting, reranking, or verification. While effective, this approach leaves the streams uncoordinated, and a discovery made in one rollout cannot influence any other. Multi-Sequence Verifiers jointly score candidate answers after generation, but the reasoning streams remain independent during inference. ParScale takes a step toward model-internal parallel scaling by learning to aggregate multiple forward passes, but its streams still do not exchange information while processing a sequence. We introduce BatchParScale, an attention mechanism that allows parallel streams to exchange information through a small set of communication positions. We evaluate BatchParScale on a new Parallel Needle-in-a-Haystack retrieval benchmark and in pretraining, fine-tuning, and cooperative multi-agent reinforcement learning. Across these settings, BatchParScale outperforms ParScale on language-model loss and remains competitive on coordination-sensitive MARL tasks. Masked pretraining runs beat an unmasked single-sequence baseline despite local sequence corruption, and LoRA fine-tuning shows the mechanism can be added to a pretrained Transformer. At the same time, fact retrieval performance for pretrained and finetuned BatchParScale models remains unclear.

Figures & tables

Figure 1: Overview of BatchParScale. (A) BatchParScale lets parallel streams share selected hidden states through sparse communication positions. (B) ParScale streams interact only at final aggregation.

Figure 2: PNIAH holder and non-holder accuracy for $B=4$ and $B=8$ . Rows correspond to hidden sizes 32, 64, and 128, columns correspond to $B=4$ and $B=8$ , and the $x$ -axis is log-spaced haystack length. The gray $\times$ marks a configuration that did not complete.

Figure 3: Masked-context recovery on TinyStories. Top: masked-token argmax accuracy with cross-stream attention enabled. Bottom: accuracy gain over the same diagnostic with cross-stream attention disabled.

Model	Mask	$B$	Eval loss	PQA acc.	PQA rank	PQA (Concat) acc.
Qwen3 baseline	-	-	1.711	-	-	-
ParScale	-	4	1.023	-	-	-
ParScale	-	8	1.268	-	-	-
Static-BPS	0%	4	1.014	0.375	2.81	1.000
Static-BPS	30%	4	1.485	0.250	2.75	1.000
Dynamic-BPS	0%	4	1.010	0.250	3.00	0.250

Table 1: TinyStories pretraining results. Eval loss is aggregate validation loss. PQA accuracy/rank evaluate a split source-query probe with

n=8

examples and five candidates, so chance accuracy is

0.2

and chance rank is about 3. PQA Concatenated accuracy uses the same probe with source and query in one ordinary sequence.

Figure 4: Router diagnostics for the Dynamic-BPS $B=4$ , $30\%$ masked TinyStories model over 500 validation chunks. Top: $P(\mathrm{mask}\mid\mathrm{virtual})$ (bars) and $P(\mathrm{virtual}\mid\mathrm{mask})$ / $P(\mathrm{virtual}\mid\mathrm{unmasked})$ (lines). Middle: cross-stream attention mass from masked versus unmasked virtual queries. Bottom: full attention-budget shares for virtual queries (cross-stream virtual, same-stream virtual, same-stream non-virtual). Layer 0 routes no mask tokens as virtual.

Layout	Source	CE on	CE off	$\Delta$ CE	Mass R0	Mass R1	Mass R2	Mass R3
$AAAA$	A	2.517	3.327	-0.811	0.248	0.258	0.252	0.241
$AABB$	A	2.557	3.311	-0.754	0.366	0.368	0.134	0.131
$AABB$	B	2.711	3.487	-0.777	0.125	0.128	0.375	0.372
$BBBB$	B	2.743	3.444	-0.700	0.248	0.244	0.256	0.253
$AAUU$	A	2.507	3.319	-0.812	0.367	0.366	0.132	0.135
$AAUU$	U	12.917	12.739	0.178	0.116	0.113	0.387	0.384

Table 2: Source-mismatch masked evaluation for the Dynamic-BPS

B=4

30\%

masked TinyStories model over 100 validation batches. Row indices follow the layout string left to right (e.g. AABB assigns rows 0–1 to stream A and rows 2–3 to stream B; AAUU assigns rows 0–1 to A and rows 2–3 to U). Mass R

i

is the fraction of attention from masked router-selected queries in the source stream to router-selected keys in batch row

i

, renormalized over all router-selected keys in the batch.

Model	Mask	Train	Eval	PN on	PN off	PN cat	PM on	PM off	PM cat	WH on	WH off	WH cat
Qwen3 base	-	-	-	-	-	1.000	-	-	1.000	-	-	.360
Qwen3 LoRA	-	3.079	3.121	-	-	1.000	-	-	1.000	-	-	.350
ParScale $B=4$	-	2.890	2.912	-	-	-	-	-	-	-	-	-
BatchParScale $B=4$	0%	2.889	2.853	.094	.107	-	.260	.350	-	.230	.060	.200
BatchParScale $B=4$	30%	3.879	3.932	.112	.117	-	.220	.340	-	.270	.130	.200
BatchParScale $B=4$	50%	4.682	4.777	.102	.133	-	.300	.290	-	.210	.110	.210

Table 3: FineWeb-Edu LoRA continued-pretraining and QA evaluation results. Train and eval are ordinary FineWeb-Edu next-token losses after 100M unique training tokens. PN is Parallel-NIAH non-holder accuracy, PM is four-way Parallel-MHQA accuracy, and WH is WikiHop candidate accuracy. On/off denote cross-stream attention enabled and disabled for BatchParScale rows. Cat(concatenated) puts all evidence in one ordinary sequence. All candidate-scoring columns use mean token log-probability, which avoids the length bias of sequence-sum scoring. The BatchParScale

B=4

WikiHop row reports the direct split-evidence fine-tuning control from Appendix H .

Task	Policy	Final	Std.	Best	Last 10
Spread	MLP	-22.7	0.4	-22.3	-22.7
Spread	RNN	-21.0	0.2	-20.7	-21.3
Spread	Agent	-19.3	0.4	-18.8	-19.2
Spread	Temporal	-19.2	0.2	-18.8	-19.2
Spread	Temporal-Agent	-19.1	0.4	-18.3	-18.8
Spread	BatchParScale	-18.9	0.3	-18.5	-19.0

Table 4: MPE Simple Spread and Simple Reference results, 500K timesteps per run, seeds 0–2. Higher return is better. RNN is the official recurrent baseline; Agent, Temporal, and Temporal-Agent are transformer policies.

Policy	Return	Std.	Win rate	Time
MLP	1.5	0.7	0.65	212s
RNN	1.2	0.6	0.35	342s
Agent Transformer	1.1	0.5	0.30	223s
Temporal Transformer	0.8	0.0	0.09	645s
Temporal-Agent Transformer	0.8	0.1	0.09	1231s
BatchParScale	0.8	0.0	0.07	1307s

Table 5: JaxMARL SMAX 5m_vs_6m , seeds 0–2, about 10M environment steps per seed. Higher return and win rate are better. Time is mean wall-clock time per seed. Win-rate spread was not logged for this summary.

Policy	Final	Std.	Best	Last 10
MLP	74	2	79	74
RNN	97	11	113	96
Agent Transformer	110	2	114	107
Temporal Transformer	95	5	101	92
Temporal-Agent Transformer	96	7	99	93
BatchParScale	85	3	89	84

Table 6: MABrax HalfCheetah 6x1, seeds 0–2, 500K timesteps per run. Metrics are rollout returns; higher is better.

Policy	Found	Std.	Remaining	Time
MLP	18.4	0.4	0.08	44.5s
RNN	18.1	0.9	0.10	144.0s
Agent Transformer	19.2	0.6	0.04	67.6s
Temporal Transformer	18.4	1.3	0.08	110.5s
Temporal-Agent Transformer	19.7	0.2	0.02	99.3s
BatchParScale	19.4	0.4	0.03	111.8s

Table 7: Jumanji SearchAndRescue-v0, seeds 0–4, 500K timesteps per run with 4 searchers, 20 targets, time limit 200, and 64 vision rays. Higher found and lower remaining are better.

Figure 5: Single-layer prefill/block overhead on an A100 for $B=4,8$ at $K/n=1/32$ , normalized by local causal attention. Crosses mark out-of-memory runs. The direct masked implementation can erase the expected advantage over dense cross-stream attention.

Figure 6: Single-layer timing normalized by local causal attention. Dense cross-stream attention grows expensive with sequence length and batch count; exact two-pass attention stays close to local causal attention while matching the selected-token output.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	Vision	Agent	Temporal-Agent	BatchParScale	$\Delta$
Sparse	4	17.29	16.87	17.18	+0.31
Sparse	8	15.69	16.23	15.49	-0.74
Sparse	16	15.84	15.52	16.18	+0.66
Sparse	24	15.78	16.01	16.53	+0.52
Large	12	12.39	10.55	11.87	+1.32
Large	16	10.22	10.16	10.58	+0.41

Appendix

Table 8: SearchAndRescue visibility sweep, reporting final targets found over five seeds. Sparse settings use env size 2.0, 4 searchers, 25 targets, and time limit 300; large settings use env size 4.0 and time limit 800.

\Delta

is BatchParScale minus Temporal-Agent. Higher is better.

Figure 7: Zoomed TinyStories training loss over unique tokens for the same runs summarized in Table 1 . The plot omits the first 25M unique tokens and applies a 9-point median smoothing window so that late-training differences are visible. The top panel shows no-mask Static-BPS and Dynamic-BPS runs together with Qwen3 and ParScale baselines; the BPS curves cluster below ParScale $B=4$ , while the y-axis zoom clips the highest-loss baseline segments. The bottom panel shows the 30% masked duplicate runs.

Run	$\Delta=0$ oracle	$\Delta=0$ non-oracle	$\|\Delta\|\leq 4$
$B=4$ , mask 30%	2.3	2.8	37.2
$B=4$ , mask 50%	1.1	7.7	40.6
$B=8$ , mask 30%	0.1	7.2	40.9
$B=8$ , mask 50%	0.2	4.8	42.4

Appendix

Table 9: Relay-source offset diagnostics. Values are percentages of top-source relay mass. Oracle means the aligned position in another stream contains the visible copy of the masked token. Non-oracle means the aligned position is also masked, so the token identity is not available at that location.

Figure 8: FineWeb-Edu LoRA continued-pretraining training loss over unique tokens for the runs summarized in Table 3 . The top panel shows all five runs. The bottom panel zooms in on the Qwen3, ParScale, and no-mask BatchParScale runs. The masked BatchParScale variants have substantially higher language-model loss, even though the no-mask BatchParScale run is competitive with the regular ParScale baseline. Periodic oscillations reflect the fixed ordering of the training stream.

Figure 9: Masked-context recovery diagnostics for the masked FineWeb-Edu LoRA continued-pretraining runs summarized in Table 3 . The top panel evaluates masked-token argmax accuracy with cross-stream attention enabled. The bottom panel reports accuracy gain over the same diagnostic with cross-stream attention disabled.

Stream	Text used while scoring one candidate
0	Source record only
1	Masked query followed by $\langle$ candidate $\rangle$
2	Neutral background note about the school yard.
3	Neutral background note about the paper crown.

Appendix

Table 10: Stream layout for the cross-stream candidate probe. The callback repeats this layout once for each candidate and scores the candidate tokens on stream 1.

Stream	Prompt	Greedy continuation
0	In the toy shop, Anna picked up a tiny drum.	She was so excited! She took it home and started to play. Anna’s mom saw her playing and said, “That drum looks so cute!”
1	In the toy shop, Ben looked for a yellow train.	He found one and said, “I want this one!” The shopkeeper said, “That’s a nice train. It will fit in your hand”
2	In the toy shop, the shopkeeper smiled and opened a box.	Inside the box was a toy car. The car was red and shiny. The shopkeeper said, “This car is for you!” The little boy
3	In the toy shop, a kitten slept under the shelf.	The kitten was very cute and soft. Lily wanted to take the kitten home, but her mom said no. Lily was sad, but she listened

Appendix

Table 11: Generated multi-stream TinyStories sample from the completed Static-BPS

B=4

, no-mask run udilkiu0 . Continuations are greedy decodes with max_new_tokens=32 .

Stream	Masked prefix tail	Sampled continuation
0	Once $\langle$ mask $\rangle$ a time, in $\langle$ mask $\rangle$ big $\langle$ mask $\rangle$ , there lived a rhinoceros named Roxy. Roxy loved to climb. She climbed $\langle$ mask $\rangle$ $\langle$ mask $\rangle$ , rocks, and hills $\langle$ mask $\rangle$ One day, Roxy found an icy $\langle$ mask $\rangle$	. She thought it would be fun to climb it. Roxy climbed and climbed, but
1	Once upon a $\langle$ mask $\rangle$ , in a big forest, $\langle$ mask $\rangle$ lived a rhinoceros $\langle$ mask $\rangle$ R $\langle$ mask $\rangle$ . R $\langle$ mask $\rangle$ loved $\langle$ mask $\rangle$ climb. She climbed trees, $\langle$ mask $\rangle$ $\langle$ mask $\rangle$ , and hills. One $\langle$ mask $\rangle$ , R $\langle$ mask $\rangle$ found an icy hill	. It was so cold, she thought she would freeze! Remy wanted to see how
2	Once upon $\langle$ mask $\rangle$ time, $\langle$ mask $\rangle$ a big forest $\langle$ mask $\rangle$ , there $\langle$ mask $\rangle$ a rhin $\langle$ mask $\rangle$ eros named $\langle$ mask $\rangle$ oxy. Roxy loved to climb. She climbed trees, rocks, $\langle$ mask $\rangle$ $\langle$ mask $\rangle$ . $\langle$ mask $\rangle$ day, Roxy found an icy hill	. She thought it would be fun to climb the icy hill. Roxy started to climb
3	$\langle$ mask $\rangle$ upon a time $\langle$ mask $\rangle$ , in a $\langle$ mask $\rangle$ forest, there lived $\langle$ mask $\rangle$ $\langle$ mask $\rangle$ $\langle$ mask $\rangle$ oc $\langle$ mask $\rangle$ named Roxy $\langle$ mask $\rangle$ Roxy $\langle$ mask $\rangle$ to $\langle$ mask $\rangle$ $\langle$ mask $\rangle$ $\langle$ mask $\rangle$ $\langle$ mask $\rangle$ trees, rocks, and hills. One day $\langle$ mask $\rangle$ , Roxy $\langle$ mask $\rangle$ $\langle$ mask $\rangle$ $\langle$ mask $\rangle$ hill	and decided to climb it. As Roxy climbed, she met a tiny bird named Ch

Appendix

Table 12: Masked-prefix multi-stream TinyStories sample from the completed dynamic local-causal

B=4

, 30% mask Hugging Face model. Each row is one masked copy of the same validation prefix.

Task	Train	Eval	Ctrl.	Cross
WikiHop	800	100	.230	.300
QASC	800	926	.454/.185	.472
Guarded QA	800	926	.125	.146
QA-v4	800	128	.117	.164

Appendix

Table 13: Direct split-evidence QA fine-tuning probes. Values are accuracies. Control is isolated split accuracy, except QASC reports question/options-only followed by isolated split accuracy. Cross-stream is the best split-evidence fine-tuning result. WikiHop uses candidate scoring with chance about

0.047

; the other multiple-choice tasks have chance

0.125

. QA-v4 has 926 validation examples; for comparability across these runs we report the fixed 128-example evaluation.

Figure 10: SearchAndRescue env-size 4.0 visibility sweeps ordered by target count. Rows scale the number of searchers and targets together: 4 searchers with 25 targets, 8 searchers with 50 targets, and 16 searchers with 100 targets. Points report mean final targets found over five seeds, with error bars showing standard deviation. Higher is better.

Benchmark	Policy	Final	Std. / Win	Best	Last 10 / Time
SMAX	Perceiver Transformer	1.2	0.6 / 0.38	–	1580s
SMAX	Agent-Query Perceiver	1.6	0.6 / 0.67	–	1418s
SMAX	Agent-Query Sink Perceiver	1.2	0.6 / 0.39	–	1806s
SMAX	Agent-Query Sink Perceiver + RoPE	2.0	0.0 / 0.96	–	1892s
SMAX	Sparse Cross-Batch Transformer	0.7	0.0 / 0.02	–	1565s
MABrax	Last-token Agent Transformer	90	13	106	94

Appendix

Table 14: Additional MARL architecture probes excluded from the main tables. SMAX entries use seeds 0, 1, and 2 from the A100 num_envs=64 rerun; their second metric reports return standard deviation / mean win rate, and the final column is mean wall-clock time per seed. MABrax entries report rollout-return standard deviation, best rollout return, and mean last-10 rollout return.

--

indicates that a baseline was not run for that suite. RoPE denotes rotary position embeddings ( Su et al., 2024 ) .

Parallel Reasoning in Language Models

See all 18 papers

Towards Better Exploration in Sequential Test-Time Scaling
Sep 30, 2026Joseph Rance, Fabio Pizzati, Juil Sock +4Test-Time Scaling
University of Oxford · MBZUAI · BBC R&D
Message Passing Enables Efficient Reasoning
Jul 1, 2026Xuecheng Liu, Daman Arora, Gokul Swamy +1Parallel Reasoning in Language Models Inference-Time Scaling
Carnegie Mellon University
Parallel Test-Time Scaling with Multi-Sequence Verifiers
Mar 3, 2026Yegon Kim, Seungyoo Lee, Chaeyun Jang +2LLM Answer Verification Test-Time Scaling
Graduate School of AI, KAIST, Seoul, South Korea

Sep 30, 2026J/K move · Enter open · S save

Towards Better Exploration in Sequential Test-Time Scaling

Joseph Rance, Fabio Pizzati, Juil Sock, Woody Bayliss+3

University of Oxford · MBZUAI · BBC R&D

Test-Time Scaling

Open paper AskPDF arXiv

Test-time scaling improves language model reasoning by spending additional compute at inference. However, both classes of existing methods often fail to continue improving over long timescales. Parallel methods repeatedly sample independent answers from the model, scaling poorly on problems the model is unlikely to solve in a single attempt. In contrast, sequential methods build on previous answers to access new ideas, yet so far have not been shown to reach answers beyond those found by parallel scaling. First, we show that sequential scaling often stops improving because it becomes prematurely trapped in an attractor: a set of answers that prevents exploration of different answers once entered. Across 27 combinations of scaling methods, models, and benchmarks, we find that 53.8% of sequential scaling trajectories enter an attractor within four iterations. Second, we show that a simple model-mixing intervention helps escape attractors. This reduces the attractor hit rate by 21.2 percentage points on average, expands solution coverage beyond a compute-matched parallel baseline, and improves accuracy of recursive self-aggregation by at least 2.2 percentage points. Our results motivate refocusing long-horizon test-time scaling from parallel methods to sequential methods that improve previous answers.

Abstract

Parallel Reasoning in Language Models

Explore similar work

Towards Better Exploration in Sequential Test-Time Scaling

Message Passing Enables Efficient Reasoning

Parallel Test-Time Scaling with Multi-Sequence Verifiers