Compute Has Shape: Self-Morphing the Internal Computation of Frozen LLMs

Junxi Wu; Yiwen Guo; Shu-Tao Xia

cs.AIOct 28, 2026

Compute Has Shape: Self-Morphing the Internal Computation of Frozen LLMs

Authors: Junxi Wu, Yiwen Guo, Shu-Tao Xia

Organizations: School of Statistics and Data Science, Nankai University · Tsinghua Shenzhen International Graduate School, Tsinghua University · Independent Researcher

Abstract

Test-time scaling improves large language models through additional computation, but how to adapt the topology of internal computation to individual inputs remains underexplored. Existing approaches largely fix recurrent or parallel architectures during training, or reconfigure frozen models along a single structural axis. We introduce SelfMorph, which makes computational shape an input-dependent inference-time decision. It unifies the original forward path (BASE), serial reuse of a layer span (REPEAT), and local branching and merging of perturbed hidden states (BRANCH) in a unified space parameterized by type, location, and intensity. Since they offer complementary benefits rather than uniform dominance, we use a lightweight selector that reads the backbone's depth-wise hidden state trajectory on the prompt to choose a morphology before response generation. We train the selector by distilling utility-based distribution from offline morphing that jointly account for task performance and computational cost. Evaluated on three LLMs spanning 4B-27B, SelfMorph achieves a 8.06% mean relative improvement on five in-distribution benchmarks and 4.45% on four unseen benchmarks. It achieves the strongest performance in 24 of 27 settings. These findings identify input-adaptive computational topology as a complementary axis of test-time scaling beyond generating more tokens. Our code is available at https://github.com/creator-xi/SelfMorph.

Figures & tables

Figure 1: Overview of SelfMorph . (a) Probe: During prompt prefill along the original forward path, we collect a depth-wise trajectory. (b) Select: A lightweight selector encodes the trajectory and aggregates them through bidirectional attention with a learned token. Before response generation, it selects the highest-probability action. (c) Morph: Three morphology types.

Backbone	Benchmark	Base	Repeat	Branch	Bench Prior	SelfMorph	Gain
Qwen3.5-27B	BBH	0.5105	0.6632	0.4474	0.6632	0.6895	+35.06%
	GSM8K	0.4120	0.4280	0.2720	0.4120	0.4360	+5.83%
	MATH-500	0.5400	0.4800	0.4600	0.5300	0.5600	+3.70%
	ARC-Challenge	0.6200	0.7667	0.6733	0.7733	0.7800	+25.81%
	OpenBookQA	0.8267	0.8333	0.7400	0.8333	0.8400	+1.61%
Qwen3.5-4B	BBH	0.4105	0.4421	0.3789	0.4526	0.5211	+26.94%

Table 1: Scores on five benchmarks, evaluated on completed responses. Bold marks the highest score per row. Gain is relative improvement of SelfMorph over Base.

Benchmark	Base	Repeat	Branch	Bench Prior	SelfMorph	Gain
BBH	0.2000	0.5105	0.1684	0.5105	0.6053	+202.65%
ARC-Challenge	0.3667	0.7333	0.5000	0.7333	0.7467	+103.63%
OpenBookQA	0.5400	0.7867	0.6267	0.7867	0.7933	+46.91%

Table 2: Short-budget scores on Qwen3.5-27B across three benchmarks. Generation is limited to 32 tokens, so a correct answer must be produced and extracted within this window. This supplementary evaluation characterizes performance under an early-answer constraint.

Backbone	Benchmark	Base	Repeat	Branch	SelfMorph	Gain
Qwen3.5-27B	2WikiMultiHop	0.7463	0.7483	0.6624	0.7708	+3.28%
	HellaSwag	0.9100	0.9500	0.8700	0.9300	+2.20%
	TruthfulQA-MC1	0.7900	0.8000	0.7300	0.8000	+1.27%
	GPQA-Diamond	0.5800	0.5500	0.4000	0.6000	+3.45%
Qwen3.5-4B	2WikiMultiHop	0.5643	0.5527	0.5757	0.5833	+3.37%
	HellaSwag	0.8500	0.8800	0.8500	0.8900	+4.71%

Table 3: Scores on four unseen benchmarks. Bold marks the highest score per row.

Table 5

Figure 2: The trajectory features before morphing. Four groups are defined by which family the per-sample oracle selects: base-hopeless (oracle is Base but sample is unsolved), repeat-best, base-enough ( Base and already correct), branch-best. (a) Relative hidden-state update near the readout point. (b) Minimum cosine similarity across the trajectory window.

Figure 3: Oracle action distribution. Left: morphology type breakdown. Center: intensity distribution within Repeat . Right: intensity distribution within Branch .

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Score gain $\Delta(i,j)$ on the numerical estimation probe. Top: Repeat ( $r{=}1$ ). Bottom: Branch ( $K{=}4$ ). Columns left to right: Qwen3.5-27B ( $N{=}64$ ), Qwen3.5-4B ( $N{=}32$ ), Gemma4-12B ( $N{=}48$ ). The $x$ -axis is span start $i$ , the $y$ -axis is span end $j$ , and only the upper triangle ( $i<j$ ) is valid.

Figure 5: Score gain $\Delta(i,j)$ over Base on the emotion-intensity probe. Top: Repeat ( $r{=}1$ ). Bottom: Branch ( $K{=}4$ ). Columns left to right: Qwen3.5-27B ( $N{=}64$ ), Qwen3.5-4B ( $N{=}32$ ), Gemma4-12B ( $N{=}48$ ). The $x$ -axis is span start $i$ , the $y$ -axis is span end $j$ , and only the upper triangle ( $i<j$ ) is valid.

Type		Qwen3.5-27B	Qwen3.5-4B	Gemma4-12B
Repeat	Span 1	$(33,35)$	$(30,31)$	$(36,37)$
	Span 2	$(31,35)$	$(28,31)$	$(26,27)$
	Span 3	$(30,36)$	$(24,31)$	$(15,16)$
	Span 4	$(26,35)$	$(24,30)$	$(25,27)$
Branch	Span 1	$(1,10)$	$(23,28)$	$(0,9)$
	Span 2	$(4,7)$	$(3,8)$	$(0,10)$

Appendix

Table 6: Layer spans in the final action space

\mathcal{A}

for each backbone. Each span is paired with three intensity levels (

r\in\{1,2,3\}

for Repeat ,

K\in\{2,4,8\}

for Branch ).

Benchmark	Train	Validation	Test	Total
BBH	2619	266	190	3075
ARC-Challenge	779	200	150	1129
OpenBookQA	623	200	150	973
GSM8K	2922	500	250	3672
MATH500	350	50	100	500
Total	7293	1216	840	9349

Appendix

Table 7: Sample counts per benchmark and split. All three backbones use the same samples, drawn from a single frozen sample manifest with a fixed random seed.

Figure 6: Selector architecture. (a) Encoder: per-layer hidden states are pre-normalized and projected to a shared low-dimensional space, then combined with learned depth embeddings. (b) Trunk: a learned query token $[\mathrm{MORPH}]$ and the $L$ layer tokens attend to one another in a bidirectional transformer block. The output at the query position serves as the decision vector. (c) Head: the decision vector is mapped to a type distribution and two joint span-intensity distributions, composed into the final action distribution.

	Qwen3.5-27B	Qwen3.5-4B	Gemma4-12B
Backbone $d_{h}$	5120	2560	3840
Encoder output $d$	192	192	384
Trunk input	$d$	$d$	$2d$
Trunk width $d_{z}$	64	64	128
Attention heads $n_{h}$	2	2	4
Blocks $n_{\mathrm{blk}}$	1	1	1

Appendix

Table 8: Selector dimensions and parameter counts. All parameters are trainable. The pre-norm gains (

\gamma_{l}

) are frozen buffers and not included.

Optimizer	AdamW
Peak learning rate	$8\times 10^{-4}$ / $8\times 10^{-4}$ / $7\times 10^{-4}$
Schedule	Cosine decay, 5% linear warmup, min = 2% of peak
Gradient clipping	Global norm 5.0
Effective batch / micro-batch	256 / 128 (accumulation = 2)
Max epochs	50 / 50 / 50
Early stopping patience (epochs)	8 / 8 / 8

Appendix

Table 9: Training hyperparameters. Values that differ across backbones are separated by slashes (27B / 4B / Gemma).

Figure 7: Generation throughput on 840 inputs per backbone, with branch-cache deduplication enabled. Top: decode-phase rate, excluding prefill and its first output token. Bottom: generation-call throughput, including prompt prefill. Bars show sample means; annotations compare each mean with BASE. Strategies can produce different output lengths, so throughput and per-query latency need not have the same ordering.

Figure 8: Average resident GPU memory during generation on the same 840 inputs. Top: frozen weights and runtime overhead. Bottom: runtime overhead on a separate scale. Eligible branch caches share physical storage; recurrent states with incompatible in-place updates remain replicated. Values are resident averages rather than peaks.

Figure 9: Hidden state L2 norm (log scale) across executed depth for one ARC-Challenge context (frozen Qwen3.5-27B). Gray: unmodified Base trajectory. Red solid/dashed: Repeat with $r{=}1$ and $r{=}2$ extra passes through span $[31,35)$ (shaded red). Green: Branch with $K{=}8$ replicas through span $[1,10)$ (shaded green). The eight thin lines fan out from the entry and collapse into one merged state (dark dot) at the exit.

Figure 10: Per-task dissociation between Repeat and Branch on the training split. Each point is one BBH subtask ( $\circ$ ) or one of five sample sub-blocks of ARC-Challenge ( $\square$ ), OpenBookQA ( $\triangle$ ), GSM8K ( $\diamond$ ), or MATH-500 ( $\triangledown$ ), colored by task domain (blue = reasoning, orange = science MC, green = mathematics). Axes show the net effect (help% $-$ hurt%) relative to Base . Points in the shaded off-diagonal quadrants indicate that one morphology actively hurts while the other helps on the same task group.

Variant	Mean Score	$\Delta$
Default	0.6083
w/o PreNorm	0.6036	$-0.0047$
w/o Attention	0.6012	$-0.0071$

Appendix

Table 10: Architecture ablation.

\Delta

is the absolute difference in mean score from the default architecture.

Table 20

Figure 11: Sensitivity to the cost coefficient $\lambda$ across three backbones on the training split. Top: mean score. Bottom: mean analytic cost. Each point represents a separately trained SelfMorph .

Inference-Time Adaptation of Language Models

See all 36 papers

Compute Aligned Training: Optimizing for Test Time Inference
Apr 27, 2026Adam Ousherovitch, Ambuj TewariTest-Time Optimization Inference-Time Scaling
Department of Statistics University of Michigan
Test-Time Augmentation for LLMs: Where to Inject Diversity at Inference Time
Aug 10, 2026Nikita Kozodoi, Zainab Afolabi, Jack ButlerLLM Inference Efficiency Self-Consistency Decoding
Amazon Web Services
Query-Conditioned Test-Time Self-Training for Large Language Models
May 13, 2026Chaehee Song, Minseok Seo, Yeeun Seong +2Test-Time Adaptation Test-Time Optimization
School of Electrical Engineering, KAIST · 2Graduate School of Green Growth and Sustainability, KAIST

Apr 27, 2026J/K move · Enter open · S save

Compute Aligned Training: Optimizing for Test Time Inference

Adam Ousherovitch, Ambuj Tewari

Department of Statistics University of Michigan

Test-Time Optimization Inference-Time Scaling RL Post-Training LLM Post-Training

Open paper AskPDF arXiv

Scaling test-time compute has emerged as a powerful mechanism for enhancing Large Language Model (LLM) performance. However, standard post-training paradigms, Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL), optimize the likelihood of individual samples under a base policy, creating a misalignment with test time procedures that rely on aggregated or filtered outputs. In this work, we propose Compute Aligned Training, which aligns training objectives with test-time strategies. By conceptualizing inference strategies as operators on the base policy, we derive new loss functions that maximize performance when said strategies are applied. We instantiate such loss functions for SFT and RL across common test time strategies. Finally, we provide empirical evidence that this training method substantially improves test time scaling over standard training.

Abstract

Inference-Time Adaptation of Language Models

Explore similar work

Compute Aligned Training: Optimizing for Test Time Inference

Test-Time Augmentation for LLMs: Where to Inject Diversity at Inference Time

Query-Conditioned Test-Time Self-Training for Large Language Models