PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Junxi Wu, Yiwen Guo, Shu-Tao Xia

arXiv:2610.38382v1cs.AI

TL;DR

How frozen LLMs can adapt the topology of internal computation to individual inputs remains underexplored. SelfMorph selects among the original path, serial layer reuse, and local branching from the prompt’s hidden-state trajectory. It improves mean relative scores on both in-distribution and unseen benchmarks.

  • Problem

    How to adapt the topology of internal computation in frozen LLMs to individual inputs remains underexplored, beyond approaches that fix architectures during training or reconfigure models along a single structural axis.

    Source: §1 · p. 1

  • Method

    SelfMorph selects BASE, REPEAT, or BRANCH, parameterized by type, span, and intensity, using a prompt-trajectory selector trained by distilling utility-based offline evaluations.

    Source: §1 · p. 2

  • Results

    8.06% mean relative score improvement on five in-distribution benchmarks and 4.45% on four unseen benchmarks are reported over the unmodified backbones.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Across three backbones and nine benchmarks, the findings support computational shape as a useful dimension of inference-time adaptation.

    Source: §5 · p. 10

    The cost model uses an analytic proxy rather than wall-clock measurements, since runtime depends on hardware, generation length, and batching strategy.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper asks whether frozen LLMs can adapt internal computation to each input, beyond scaling computation through generated sequences. SelfMorph offers a shared space of computational morphologies and selects among them from the model’s prompt trajectory.

1 INTRODUCTION

  • BASE, serial REPEAT, and local BRANCH provide complementary effects, with neither intervention uniformly subsuming the other across tasks and backbones.This heterogeneity motivates choosing a morphology for each input rather than applying one fixed intervention.Source: §1 · p. 2
  • SelfMorph predicts morphology type, span, and intensity from the backbone’s prompt trajectory before the first response token is generated.It distills a utility-based teacher distribution from offline candidate evaluations that combine task score and analytic cost.Source: §1 · p. 2
  • The framework places the original path, serial layer reuse, and local branching in a shared action space over frozen pretrained weights.These choices are parameterized by morphology type, span, and intensity.Source: §1 · p. 2
  • 8.06% on five in-distribution benchmarks and 4.45% on four unseen benchmarks are SelfMorph’s mean relative score improvements over the unmodified backbones.It also outperforms a per-benchmark selected fixed action in all 15 settings.Source: §1 · p. 2

2 PARAMETERIZED MORPHOLOGY SPACE

The morphology space organizes computation over a frozen backbone by type, span, and intensity. Its interventions provide distinct forms of added computation, while screening identifies a compact set of useful candidate actions.

2.1 MORPHOLOGY TYPES

  • BASE traverses a span once, while REPEAT makes r+1 serial passes, each building on the preceding output and increasing effective depth without new parameters.REPEAT provides additional opportunities to refine features before passing the state onward.Source: §2.1 · p. 3
  • BRANCH replicates the input into K perturbed paths through a span, then averages their resulting states before the remaining layers.One branch is unperturbed, preserving a direct contribution from the original input.Source: §2.1 · p. 3

2.2 CONSTRUCTING THE ACTION SPACE

  • Favorable REPEAT spans concentrate at intermediate depths, whereas favorable BRANCH spans tend to occur earlier, especially on larger backbones.The paper relates this asymmetry to repeating intermediate feature refinement versus introducing variation early for later processing.Source: §2.2 · p. 4
  • 25 actions per backbone combine screened spans with REPEAT intensities {1, 2, 3} or BRANCH counts {2, 4, 8}.Screening evaluates feasible spans at fixed intensities on two diagnostic datasets separate from the evaluation benchmarks.Source: §2.2 · p. 4

3 SELFMORPH

SELFMORPH selects an action from the frozen backbone’s prompt-time hidden-state trajectory, training the selector with offline utilities that account for task score and analytic cost.

3 SELFMORPH

  • Enumerating every action for each input is prohibitively costly, motivating a lightweight decision from the backbone’s own hidden states.Source: §3 · p. 4

3.1 UTILITY-BASED DISTILLATION

  • Offline action utilities combine observed task score with an analytic cost penalty, while temperature-scaled teacher distributions preserve graded preferences among actions.The cost proxy depends on intervention width and type/intensity; lower teacher temperature concentrates probability on higher-utility actions.Source: §3.1 · pp. 4–5

3.2 TRAJECTORY-CONDITIONED SELECTION

  • During prompt prefill, the selector encodes last-token states from selected depths, aggregates them with a learned query, then predicts morphology type and conditional span–intensity.The type head predicts BASE, REPEAT, or BRANCH; non-BASE types also receive a joint span–intensity prediction.Source: §3.2 · p. 5

3.3 INFERENCE: PROBE–SELECT–MORPH

  • At inference, a BASE prefill constructs the trajectory, the selector chooses an action without candidate response generation, and the chosen morphology runs with prefix reuse and cache sharing to reduce overhead.The probe retains reusable caches and hidden states; cache sharing reduces storage for replicated histories.Source: §3.3 · p. 6

4 EXPERIMENTS

The experiments evaluate SELFMORPH across frozen backbones and benchmarks, finding gains over Base in most settings and over a benchmark-level fixed strategy in every backbone–benchmark pair. Unseen-benchmark tests, ablations, and trajectory analyses examine transfer and how action-space and selector choices relate to performance.

4 EXPERIMENTS

  • Experiments use three frozen backbones and compare SELFMORPH with Base, fixed Repeat and Branch actions, and a benchmark-specific fixed-action prior across five in-distribution and four unseen benchmarks.The benchmark set spans science multiple-choice, mathematics, reasoning, and four unseen tasks.Source: §4 · p. 6

4.1 MAIN RESULTS

  • SELFMORPH improves over Base in most evaluated settings and outperforms Bench Prior in every backbone–benchmark pair.Bench Prior selects one action per benchmark using the training set; Table 1 reports scores on five benchmarks evaluated on completed responses.Source: Table 1 · p. 7, §4.1 · p. 7
  • On Qwen3.5-27B, SELFMORPH achieves higher scores when correct answers must be produced and extracted within a 32-token window.The short-budget evaluation covers three benchmarks and examines early-answer performance.Source: §4.1 · p. 7, Table 2 · p. 7

4.2 GENERALIZATION TO UNSEEN BENCHMARKS

  • On four benchmarks excluded from selector training, SELFMORPH improves over Base in most settings without retraining or refitting.The selector is evaluated on the unseen benchmarks as-is.Source: §4.2 · p. 7, Table 3 · p. 8

4.3 ABLATION STUDIES

  • The 9-action mixed pool outperforms larger single-family pools in most comparisons, while expanding the mixed pool increases oracle headroom faster than selector performance.Expanding the mixed pool improves performance on both Qwen backbones; Gemma matches its 9-action score with 25 actions.Source: §4.3 · p. 8, Table 4 · p. 8
  • The 10-layer hidden trajectory achieves the highest score among tested selector inputs on all three backbones, and almost every hidden-state condition outperforms prompt embeddings.The comparison includes a single intermediate-layer state and uniformly subsampled 3- and 5-layer trajectories.Source: §4.3 · p. 8, Table 5 · p. 8

4.4 ANALYSIS

  • Oracle groups order from base-hopeless through repeat-best and base-enough to branch-best as relative updates decrease and cosine similarity increases.The features are measured across the trajectory window before morphology is applied.Source: §4.4 · p. 9, Figure 2 · p. 9
  • For Qwen3.5-27B, BASE is the oracle choice for most samples, while the lowest intensity receives the largest share of selections within both REPEAT and BRANCH.Some problems require higher intervention intensities when low intensity fails; BASE also wins at zero extra cost for already-correct and unsolved inputs.Source: §4.4 · p. 9, Figure 3 · p. 9

5 CONCLUSION

SELFMORPH treats frozen-model computation as a choice among the original path, serial layer reuse, and local branching, while related approaches generally focus on recurrent or parallel execution. The appendix screens candidate spans with fixed-intensity interventions on diagnostic probes to identify depth regions for later action selection.

5 CONCLUSION

  • SELFMORPH makes frozen-model execution an input-dependent choice among the original path, serial reuse, and local branching, and concludes that computational shape is a useful dimension of inference-time adaptation.Its selector uses offline interventions and the backbone’s prompt trajectory to select morphology type, span, and intensity before generation.Source: §5 · p. 10

A RELATED WORK

  • Prior recurrent-depth models are trained for repeated execution, while work on frozen backbones focuses on skipping or repeating layers; SELFMORPH jointly selects serial reuse and local branching.Related parallel methods use multiple latent streams, whereas SELFMORPH branches within a selected span of a frozen backbone using Gaussian perturbations and arithmetic averaging.Source: A RELATED WORK · p. 12

B SPAN SCREENING

  • Candidate REPEAT and BRANCH spans are screened at fixed intensity on a small, fast diagnostic probe set separate from the formal evaluation benchmarks.The scan enumerates every valid half-open span, with BRANCH excluding spans ending at the final layer because its merge hook requires a layer at position j.Source: B SPAN SCREENING · p. 12, §B.1 · p. 13

B.1 SETUP

  • Screening uses three frozen backbones and two 16-sample probes: large-integer estimation and emotion-intensity prediction, designed to cover distinct capabilities with short outputs.Probe scores use relative numerical error with a digit-length correction, and absolute deviation from reference emotion intensities.Source: §B.1 · p. 12

B.2 HEATMAPS

  • The heatmaps plot mean score gain over BASE by span start and end for REPEAT and BRANCH across three backbones; red indicates improvement and blue degradation.REPEAT uses r=1 and BRANCH uses K=4 in the displayed scans.Source: §B.2 · p. 13, Figure 4 · p. 13

B.3 OBSERVATIONS

Intervention effects vary by span, morphology, backbone, and probe, so candidate actions are screened and ranked rather than assumed to transfer uniformly. The resulting action space uses balanced probe rankings and backbone-specific spans, while evaluation covers in-distribution and held-out benchmarks.

B.3 OBSERVATIONS

  • Many spans have negative gains, including some beginning in the earliest layers, and affected regions vary by morphology, backbone, and diagnostic task.REPEAT’s positive region concentrates at mid-to-late depth across both probes and all three model scales.Source: §B.3 · p. 13
  • 72% of BRANCH spans improve the numerical probe on Qwen3.5-27B, while smaller backbones show broader positive regions extending into mid-depth.The numerical-probe BRANCH peaks are (23, 28) for Qwen3.5-4B and (17, 25) for Gemma4-12B; gains taper toward the output end.Source: §B.3 · p. 14

B.4 FROM SCREENING TO ACTION SPACE

  • The numerical and emotion probes rank spans differently, so candidates are ranked by summed z-scores and validated on larger samples, with widths limited to 10 layers.Final REPEAT spans lie near the middle for Qwen3.5-27B, toward the end for Qwen3.5-4B, and across intermediate and later depths for Gemma4-12B; BRANCH spans are early for Qwen3.5-27B and broader in smaller backbones.Source: §B.3 · p. 14, §B.4 · p. 14, Table 6 · p. 15

C.1 BENCHMARK DESCRIPTIONS

  • Evaluation uses five in-distribution benchmarks and four benchmarks held out entirely from training for zero-shot assessment.The benchmark set includes BBH reasoning tasks and science, mathematics, multi-hop QA, commonsense, and truthfulness tasks.Source: §C.1 · p. 15

C.2 DATASETS AND SPLITS

The in-distribution test samples are held out from selector training and early stopping, while four unseen benchmarks use frozen sample lists without refitting. The appendix also specifies prompt and scoring protocols, oracle-utility collection, and fixed-action baselines.

C.2 DATASETS AND SPLITS

  • Five in-distribution benchmark units span science multiple-choice, mathematics, and reasoning, while unseen benchmarks are evaluated without refitting the production selector.Test samples are held out from training and early stopping, and all backbones use the same frozen sample manifest.Source: §C.2 · pp. 15–16, Table 7 · p. 16

C.3 PROMPT TEMPLATES

  • Prompt protocols vary by benchmark family, with direct-answer instructions and fixed demonstrations for the in-distribution oracle-collection tasks.The main results use generation without an imposed length limit, unlike the short-budget oracle collection and ablations.Source: §C.3 · pp. 16–17

C.4 COLLECTION PROCESS

  • Oracle collection runs each of 9,349 samples against all 25 actions, recording task score, analytic cost, utility, and generated text for each backbone.The BASE pass also records a 10-layer hidden trajectory per sample for the selector’s input.Source: §C.4 · p. 18

C.5 BASELINES

  • Four fixed baselines select one action by highest mean training utility, including a benchmark-prior baseline that assigns one action per benchmark unit.All samples within a benchmark unit share that baseline action.Source: §C.5 · p. 18

D SELECTOR ARCHITECTURE AND TRAINING

The selector consumes a depth-wise hidden trajectory extracted during the BASE forward pass. Its input is normalized and scaled using the backbone’s own layer-normalization gains.

D.1 INPUT REPRESENTATION

  • The selector receives a hidden-state trajectory from 10 layers spanning the mid-to-late portions of each backbone.The sampled layers are backbone-specific and cover different absolute depth indices.Source: §D.1 · p. 18
  • Each layer slice is RMS-normalized and scaled by the consuming layer’s pretrained gain vector, which is frozen rather than trainable.This aligns selector inputs with the numerical semantics used by the backbone at each depth.Source: §D.1 · p. 18

D.2 NETWORK ARCHITECTURE

The selector encodes the prompt’s layerwise hidden states and maps them to a morphology distribution, while inference optimizations reuse compatible computation and cache storage.

D.2 NETWORK ARCHITECTURE

  • The encoder projects per-layer states into depth-aware tokens, and a bidirectional transformer lets a learned query aggregate the trajectory into a decision vector.Gemma4-12B additionally concatenates first-order differences with the projected states.Source: Figure 6 · p. 19, §D.2 · p. 19

2 GELU

  • The head predicts a morphology type and joint span–intensity distributions, which are composed into probabilities over complete actions.Repeat and branch each have a joint head for span and intensity.Source: §D.2 · p. 19, Figure 6 · p. 19, §2 · p. 20

D.3 DIMENSIONS AND PARAMETER COUNTS

  • The selector is small relative to the backbone, and Table 8 lists its trainable dimensions and parameter counts.The table excludes frozen pre-normalization gains.Source: §D.3 · p. 20, Table 8 · p. 20

D.4 TRAINING

  • Training minimizes distillation loss against the utility-based teacher distribution, and early stopping retains the checkpoint with the highest validation mean utility.Validation utility uses the highest-probability action for each input.Source: §D.4 · p. 20

E INFERENCE OPTIMIZATION AND COST MEASUREMENT

  • Prefix reuse and lossless cache deduplication eliminate overlapping prompt computation and identical cache rows without changing the selected morphology or model settings.The optimizations preserve weights, precision, perturbations, and decoding rules.Source: E INFERENCE OPTIMIZATION AND COST… · p. 21

E.1 REUSING COMPUTATION ACROSS PROBE AND COMMIT

  • Probe and commit share compatible prefix computations; checks of 13 prefix-compatible actions on two prompts found unchanged states and generated tokens.For BASE, execution resumes at the probe depth; REPEAT reuses the shared portion before continuing its selected layer sequence.Source: §E.1 · p. 21
  • BRANCH retains its original batch shape because collapsing it outside the span can alter floating-point accumulation; with identical perturbations, the alternative changed 58 of 144 outputs.The implementation therefore uses a separate branch prefill.Source: §E.1 · p. 21

E.2 LOSSLESS DEDUPLICATION OF BRANCH CACHES

  • Identical branch-cache histories outside the span can be stored once and broadcast, while span histories remain separate; Qwen recurrent and convolutional states also remain replicated.In a controlled Gemma K = 4 comparison, deduplication removed 58.5% of retained cache storage at approximately 1.8% additional decode time.Source: §E.2 · pp. 21–22

E.3 COST EVALUATION

The section evaluates analytic morphology costs against measured execution and storage, then illustrates how REPEAT and BRANCH alter hidden-state trajectories.

E.3 COST EVALUATION

  • The analytic cost proxy penalizes intervention width and intensity, with fixed BRANCH multipliers informed by execution-time and cache measurements.The utility cost penalty is applied consistently rather than using workload-dependent wall-clock observations.Source: §E.3 · p. 22
  • Measured REPEAT decode-rate ratios differ from analytic predictions by less than 0.7 percentage points, and cache ratios by less than 0.8% relative error across three backbones.The model predicts increased executed depth from additional layer evaluations; the measurements track that structural depth term.Source: §E.3 · pp. 22–23, Figure 7 · p. 23
  • BRANCH decode throughput stays within 2.4% of BASE across three fixed configurations, while distinct branch histories increase cache storage.The reported throughput panels distinguish decode rate from generation-call throughput, which includes prefill.Source: §E.3 · p. 23, Figure 7 · p. 23
  • For Gemma’s span (0, 10) at K = 4, predicted cache storage is 1.625 times single-replica storage, near the measured ratio of approximately 1.661.The byte-weighted model accounts for unequal layer capacities and replicated recurrent states.Source: §E.3 · p. 23, Figure 8 · p. 24
  • Qwen strategies produce different output lengths, so generation-call rates combine execution speed with how many tokens amortize prefill.Gemma’s matched output lengths make its two timing panels more directly comparable.Source: §E.3 · p. 23
  • Runtime overhead excludes frozen weights; branch-cache sharing benefits Gemma most, while replicated Qwen recurrent states limit deduplication savings.Memory figures report average resident memory rather than temporary peaks.Source: §E.3 · p. 23, Figure 8 · p. 24

F.1 TRAJECTORY GEOMETRY

  • REPEAT extends the depth trajectory through its selected span, whereas BRANCH fans into K paths and merges them into a state that resumes at a different hidden-space position.Figure 9 illustrates these signatures for one ARC-Challenge context on Qwen3.5-27B.Source: §F.1 · p. 24

F.2 TASK-LEVEL DISSOCIATION

REPEAT and BRANCH have distinct task-group effects: each can help where the other hurts, so no single fixed intervention is uniformly preferable.

F.2 TASK-LEVEL DISSOCIATION

  • The per-sample oracle distribution shows that one fixed action is not always optimal, supporting input-level action selection.The analysis records utility for all 25 actions on every training sample.Source: §F.2 · p. 24
  • On Qwen3.5-27B, some BBH subtasks improve under REPEAT but decline under BRANCH, while many science sub-blocks improve under both.Opposite-sign effects also occur on smaller backbones, with different magnitudes and task-group distributions.Source: §F.2 · p. 24
  • Figure 10 plots each task group’s best REPEAT net effect against its best BRANCH net effect; shaded off-diagonal points mark groups where one helps while the other hurts.The axes measure help% − hurt% relative to BASE.Source: §F.2 · p. 24, Figure 10 · p. 25

G ADDITIONAL ABLATION STUDIES

The ablations examine selector architecture, capacity, teacher temperature, and cost weighting, identifying supported design choices and trade-offs on the evaluated setting.

G ADDITIONAL ABLATION STUDIES

  • Removing either selector self-attention or pre-normalization lowers mean score, supporting both components in the evaluated architecture.Self-attention lets the query aggregate cross-layer information using input-dependent relevance.Source: G ADDITIONAL ABLATION STUDIES · pp. 25–26, Table 10 · p. 26
  • Among tested selector sizes, 2.24M parameters achieves the highest score, while the deployed 1.05M setting performs close with fewer parameters.The smallest configuration and configurations above 3M parameters score lower.Source: G ADDITIONAL ABLATION STUDIES · p. 26, Table 11 · p. 26
  • Among tested teacher temperatures, T=0.15 achieves the highest score, favoring an intermediate distribution sharpness for Qwen3.5-27B.Lower temperatures concentrate teacher probability on higher-utility actions; higher temperatures increase entropy.Source: G ADDITIONAL ABLATION STUDIES · p. 26, Table 12 · p. 26
  • A small positive cost coefficient improves mean score over λ = 0 on all three backbones, while larger coefficients generally reduce score.Mean analytic cost generally decreases as λ increases.Source: G ADDITIONAL ABLATION STUDIES · p. 26, Figure 11 · p. 27

Related papers

Loading 2610.38382v1…