PaperlayerView exact revision on arXiv

Source-linked AI summary

SpiralFormer: Looped Transformers Can Learn Hierarchical Dependencies via Multi-Resolution Recursion

Chengting Yu, Xiaobo Shu, Yadao Wang, Yizhen Zhang, Haoyi Wu, You Wu, Rujiao Long, Ziheng Chen, Yuchi Xu, Wenbo Su, Bo Zheng

arXiv:2602.11698v2cs.LG

TL;DR

Recursive Transformers face a gap against non-recursive baselines while typically operating at fixed full-token resolution. SpiralFormer introduces multi-resolution recursion, and the paper reports hierarchical dependencies plus superior parameter and compute efficiency across 160M–1.4B models.

  • Problem

    Existing recursive Transformers often underperform equally expensive non-recursive baselines and leave compressed, lower-resolution computation largely unexplored.

  • Method

    SpiralFormer applies one shared Transformer core across a multi-resolution recursion schedule that downsamples token states into latent chunks and causally upsamples updates.

  • Results

    SpiralFormer consistently outperforms looped and non-looped baselines in parameter and compute efficiency across model scales from 160M to 1.4B.

  • Takeaways & Limitations

    Attention-based probing indicates that multi-resolution recursion enables hierarchical dependencies and iteration-wise functional specialization, making sequence resolution a potential scaling axis.

  • Takeaways & Limitations

    A no-overlap variant creates deployment-time pipelining opportunities but reduces quality, and closing that quality gap remains open.

Abstract

from arXiv · show

Recursive (looped) Transformers decouple computational depth from parameter depth by repeatedly applying shared layers, providing an explicit architectural primitive for iterative refinement and latent reasoning. However, early looped Transformers often underperform non-recursive baselines of equal compute. While recent literature has introduced more effective recursion mechanisms to mitigate this gap, existing architectures still operate at a fixed, full-token resolution, neglecting the potential efficiency of computing over compressed latent representations. In this paper, we propose SpiralFormer, a looped Transformer that executes recurrence under a multi-resolution recursion schedule. We provide probing evidence that multi-resolution recursion enables the model to learn hierarchical dependencies by inducing iteration-wise functional specialization across different scales. Empirically, SpiralFormer achieves better parameter and compute efficiency than both looped and non-looped baselines across model scales from 160M to 1.4B, establishing sequence resolution as a potential axis for scaling recursive architectures.

1. Introduction

SpiralFormer addresses the fixed-resolution limitation of recursive Transformers by making compressed latent computation explicit within recursion. It uses multi-resolution recurrence to support hierarchical dependencies and achieves improved parameter and compute efficiency across model scales.

  • Motivation: Recursive Transformers reuse shared layers to increase computational depth and support iterative refinement with fewer parameters.Looping repeatedly updates hidden representations before token emission.
  • Motivation: Early looped Transformers often underperform equally expensive non-recursive baselines, while existing improvements retain full-token resolution.This leaves sequence resolution largely unexplored as a recursion axis.
  • Motivation: Latent reasoning suggests that compressed chunk-level representations can summarize multi-token information and improve token efficiency.This motivates introducing compressed representations directly into recursive computation.
  • Approach: SpiralFormer applies one shared Transformer core across a multi-resolution schedule, downsampling hidden states, processing latent slots, and causally upsampling updates.The architecture makes multi-token latent compression an explicit recursive operation.
  • Findings: Probing evidence indicates that multi-resolution recursion produces hierarchical, scale-dependent dependencies through iteration-wise functional specialization.Attention statistics shift coherently across iterations and heads specialize under resolution changes.
  • Findings: SpiralFormer consistently outperforms looped and non-looped baselines in parameter and compute efficiency from 160M to 1.4B.The result supports sequence resolution as an additional scaling axis for recursive architectures.

2. Method

SpiralFormer combines a loop-shared Transformer backbone with coarse-to-fine recursion over compressed chunk-level representations. Each iteration downsamples, processes, causally upsamples, and integrates updates while preserving autoregressive causality.

  • Backbone: Recursive Transformers repeatedly apply a loop-shared block, increasing computational depth without proportionally increasing parameters.The loop transition maps one hidden state to the next through shared computation and a topology update.
  • Backbone: The Middle-cycle backbone uses pre-loop, shared loop, and post-loop Transformer blocks, with topology operator U integrating each iteration's update.Anchor and MeSH instantiate alternative state-fusion mechanisms.
  • Multi-resolution recursion: SpiralFormer performs each iteration as downscale, shared-core processing, upscale, causal shift, and loop-state update.The same core operates at the resolution specified for that iteration.
  • Causality: A right-shift corrects the upsampled update because chunk aggregation can include future tokens, preserving strict autoregressive causality.The default shift s_t = g_t − 1 yields a single-token overlap between producing and receiving chunks.
  • Compression and reconstruction: Chunking partitions token states into resolution-dependent chunks, optionally offsets boundaries, aggregates each chunk into a latent vector, and allocates latent outputs back to token positions.The default design uses a half-chunk offset, self-aggregation, output-dependent allocation, and gain scaling λ_t = √g_t.
  • Resolution schedule: A coarse-to-fine schedule starts with short latent sequences and progressively increases resolution to refine token-level representations.The effective sequence length is controlled by resolution schedule {r_t}, with examples beginning at r_0 = 1/8 or 1/16.
  • Evaluation: The evaluation compares SpiralFormer with full-resolution LOOPEDFORMER and standard Pythia across parameter, FLOP, perplexity, and downstream-accuracy measures.Experiments use the Pythia family from 160M to 1.4B parameters.

3. Experiments

SpiralFormer is evaluated against standard and full-resolution recursive Transformers across model scales, compute budgets, recurrence ratios, and component ablations. Multi-resolution recursion improves efficiency and quality, with coarse-to-fine scheduling and balanced parameter sharing emerging as favorable choices.

  • Experimental setup: Models span 160M–1.4B Pythia scales and compare standard Pythia, full-resolution LOOPEDFORMER, and multi-resolution SPIRALFORMER variants.SPIRALFORMER-B matches LOOPEDFORMER’s recursive layer allocation, whereas SPIRALFORMER-L matches the non-recursive baseline’s parameter count.
  • Main results: SPIRALFORMER-B consistently improves perplexity and few-shot accuracy over LOOPEDFORMER while reducing FLOPs by ∼7–11% across sizes.Examples include 4.59→4.11 at 410M and 9.67→8.95 at 1B, without increasing parameters.
  • Main results: SPIRALFORMER-L reduces FLOPs by ∼3–10% at matched parameters while improving perplexity and few-shot performance over standard Pythia.At 1.4B, FLOPs fall from 14.08→13.13 while 5-shot accuracy rises from 51.93→54.37.
  • Scaling effects: Both SPIRALFORMER variants establish better loss–compute frontiers than LOOPEDFORMER and Pythia, with matched-parameter advantages generally widening at larger scales.The scaling comparison covers validation loss versus FLOPs and downstream 0-shot accuracy versus total parameters.
  • Recurrence ratio: Validation loss follows a U-shaped recurrence-ratio curve, reaching its minimum around 30–40% and degrading when parameter sharing becomes excessive.The pattern is consistent at 8e19 and 16e19 FLOPs on the 410M model.
  • Ablation study: Ablations favor MeSH topology, coarse-to-fine scheduling, learned self-aggregation with output-dependent allocation, and the half-chunk offset as the default configuration.Fine-to-coarse scheduling and the Anchor topology degrade performance; alternative chunk offsets can perform comparably but may create non-uniform per-token compute.

4. Analyzing Hierarchical Dependencies in Multi-Resolution Recursion

SpiralFormer’s probes test whether coarse-to-fine recursion changes attention behavior across resolutions. The results show a systematic shift from diffuse, global interactions at coarse resolutions toward selective, local refinement at fine resolutions, supporting hierarchical dependencies across loops.

  • Measurements: The probes measure loop-dependent attention behavior using key-marginal entropy and Local Attention Mass (LAM) on attention heads.Dynamic heads are the top 40% ranked by cross-loop range for each metric; statistics are averaged over 500 Pile validation samples.
  • Attention shifts: Lower key-marginal entropy at higher resolutions indicates that later loops concentrate attention on fewer salient keys.This is consistent with fine-resolution loops selectively refining dependencies rather than maintaining broad interactions.
  • Attention shifts: LAM shifts systematically across resolutions, providing evidence that later iterations increasingly emphasize local token-level dependencies.LAM measures the probability mass assigned to resolution-aligned causal local neighborhoods.
  • Hierarchical dependencies: The observed cross-loop changes support hierarchical, scale-dependent dependencies induced by multi-resolution recursion.The analysis tests the intended coarse-to-fine computation pattern through complementary attention-based probes.
  • Hierarchical dependencies: Coarse loops are more diffuse and less local, whereas fine loops are more selective and more local despite shared core parameters.Together, these signatures support global aggregation at coarse resolution followed by local refinement at fine resolution.

5. Conclusion

SpiralFormer combines recursive shared layers with multi-resolution computation, connecting looped recursion to latent reasoning and sequence compression. The paper concludes that this design improves parameter and compute efficiency while exposing sequence resolution as a scaling axis for recursive architectures.

  • Conclusion: SpiralFormer runs loop-shared layers under a multi-resolution recursion schedule.The architecture replaces fixed full-resolution recursion with computation over changing sequence resolutions.
  • Conclusion: SpiralFormer consistently outperforms looped and non-looped baselines in parameter and compute efficiency from 160M to 1.4B model scales.This is the paper’s principal empirical conclusion.
  • Conclusion: Sequence resolution emerges as a potential axis for scaling recursive architectures.The conclusion links the efficiency result with probing evidence that multi-resolution recursion supports hierarchical dependencies.
  • Connection to latent reasoning: Latent reasoning methods motivate compressing complex computation into a limited number of high-capacity latent slots.SpiralFormer operationalizes this compression morphology inside a recursive loop rather than treating it only as a distillation strategy.
  • Related architectural context: Hierarchical architectures address Transformer scaling by introducing compressed or multi-scale representations.This places SpiralFormer among approaches that use sequence compression to address the quadratic sequence-length bottleneck.

C. Right-shift for strict causality and overlap regimes

The right-shift converts chunk-level updates into causally valid autoregressive updates. Its size determines whether updates are non-causal, causally overlapping by one token, or chunk-wise non-overlapping.

  • Causality mechanism: A right shift is applied to chunk-expanded updates because chunk aggregation can include future tokens relative to earlier positions.The shift enforces that the update written at position i does not depend on tokens greater than i.
  • Overlap regimes: Shifts st ≤ gt−2 can violate strict causality because an update may depend on later tokens in the same chunk.The failure occurs when the shifted source remains within the receiving chunk.
  • Overlap regimes: Setting st = gt−1 preserves strict causality while creating a single-token overlap between the generating and receiving chunks.The overlapping position is the unique boundary position whose update may depend on its current chunk.
  • Overlap regimes: Shifts st ≥ gt preserve strict causality and eliminate chunk overlap, yielding a chunk-wise parallel regime.Every receiving position depends only on chunks strictly before its own chunk.

D. Inference-Time Motivation for Chunk Offset

Chunk offsets matter at inference because low-resolution computations are triggered periodically rather than at every token. A half-chunk offset shifts and interleaves trigger positions, helping make per-token compute more uniform without changing parameters or causality enforcement.

  • Inference-time behavior: Under autoregressive decoding, cached low-resolution computations are triggered only at periodic positions determined by chunk structure.This creates non-uniform per-token compute across resolution levels.
  • Inference-time behavior: With st = gt−1, exactly one position per chunk triggers a new low-resolution computation.Other positions reuse summaries from the previous chunk.
  • Non-uniform compute: A coarse-to-fine schedule with chunk sizes 8, 4, 2, and 1 produces nested trigger sets, so some tokens trigger multiple resolutions while others trigger only full resolution.For example, token i = 7 triggers multiple resolutions, whereas i = 0 triggers only the full-resolution iteration.
  • Chunk offset: The half-chunk offset shifts triggering positions and tends to interleave triggers across resolutions, making compute allocation more uniform in practice.Offsets change grouping and triggering patterns but do not alter model parameters or the right-shift causality mechanism.

E. Inference-Time Parallelism in the Parallel (no-overlap) Regime

In the no-overlap regime, choosing s_t ≥ g_t creates a chunk-level lag that separates current-token full-resolution work from background multi-resolution computation. This enables asynchronous preprocessing of earlier chunks, potentially moving extra recursion compute off the per-token critical path.

  • Chunk-level lag: s_t ≥ g_t ensures updates written to a chunk depend only on strictly earlier chunks, creating a chunk-level time lag.The current chunk cannot influence its own shifted update.
  • Parallel schedule: The full-resolution pathway computes current-token logits while low-resolution recursion processes earlier-chunk summaries in the background.The low-resolution branch performs downscaling, the shared core pass, and upscaling for future-token use.
  • Caching: Once a chunk is finalized, its chunk-level latent can be computed asynchronously, cached, and reused by later tokens.This shifts multi-resolution work away from later tokens’ critical paths.
  • Systems implication: The no-overlap regime is analogous in spirit to cross-loop parallelism because removing within-chunk dependencies permits earlier-chunk recursion to advance concurrently.The comparison is presented as a related systems optimization, complementary to the architectural gains of multi-resolution recursion.

F. Implementation Details

SpiralFormer is evaluated as a decoder-only GPT-NeoX model with shared loop layers and a coarse-to-fine recursion schedule. The implementation uses causal chunk processing, learned down/up-scaling, standard language-model training, and broad validation and few-shot evaluation.

  • Training setup: All models are pretrained from scratch for one epoch on 250B tokens from a deduplicated Pile subset using the GPT-NeoX tokenizer.The vocabulary contains 50,257 tokens.
  • Architecture: SpiralFormer uses a decoder-only GPT-NeoX backbone with pre-loop, shared loop, and post-loop blocks.The shared loop core is executed across T iterations under a resolution schedule.
  • Recursion configuration: The default schedule is coarse-to-fine, with causal right shifts, half-chunk offsets, and learnable self-aggregation and output-dependent upsampling.The default right shift is s_t = g_t − 1, and both Anchor and MeSH topologies are evaluated.
  • Optimization: Training uses AdamW, cosine learning-rate decay with warmup, batch size 512, sequence length 4096, BF16, FlashAttention-2, and DeepSpeed ZeRO Stage 0.The optimizer uses β1 = 0.9, β2 = 0.95, and weight decay 0.01.
  • Evaluation: Evaluation reports validation perplexity on three corpora and 0-shot and 5-shot accuracy across nine downstream tasks.The downstream suite includes both Lambada variants, HellaSwag, PIQA, WinoGrande, ARC-Easy, ARC-Challenge, SciQ, and continuation-MMLU.
  • Decoding: The implementation includes chunk-triggered recursive updates with cached key-value states for autoregressive decoding.Algorithm 2 separates prompt prefill from sequential token generation and initializes chunk buffers and latent caches.

H. Detailed Downstream Results

The detailed downstream table reports task-level accuracy for every model size and shot setting, while its average column reproduces the main results. Results are organized to compare topology variants and identify the best configurations within each setting.

  • Table coverage: Table 3 reports 0-shot and 5-shot accuracy on nine tasks for each model size, with the average accuracy in the final column.The average column corresponds to the task accuracy reported in Table 1.
  • Comparison scheme: Bold and underlined entries identify the best and second-best results within each model size and shot setting, respectively.Model notation distinguishes Anchor (*) and MeSH (†) topologies.

I.1. Additional Analysis: Where Cross-Loop Changes Concentrate

The additional analysis localizes cross-loop variability to particular attention heads rather than distributing it uniformly. Dynamic heads differ between entropy and attention-mass metrics, indicating metric-dependent responses to the coarse-to-fine schedule.

  • Range definition: For each attention head, the cross-loop range measures the difference between its maximum and minimum metric values across recursion loops.Metric values are averaged over 500 Pile validation sequences before the range is computed.
  • Dynamic heads: Dynamic heads are the top 40% of heads ranked by cross-loop range, separately for key-marginal entropy and Local Attention Mass.The ranking is performed independently for each metric.
  • Localization result: Large changes in both ∆H and ∆LAM concentrate in specific layer–head regions, but the highlighted dynamic heads differ across metrics.The heatmaps use layer–head cells and min–max normalization to compare patterns across the model.

I.2. Additional Analysis: LoopedFormer Controls (Distribution & Head-wise)

The controls compare SpiralFormer’s multi-resolution recursion against a full-resolution LoopedFormer baseline using matched attention-statistics probes. The full-resolution control shows weaker, less structured specialization, while all-head analysis confirms that SpiralFormer’s qualitative cross-loop shifts are not limited to selected dynamic heads.

  • Distribution-level trends: The control repeats SpiralFormer’s attention probes on a full-resolution LoopedFormer baseline with all iterations operating at token resolution.The evaluation uses 500 Pile validation sequences and distributions over dynamic heads selected by cross-loop range.
  • Distribution-level trends: Full-resolution LoopedFormer exhibits qualitatively different cross-loop behavior from SpiralFormer, indicating that resolution changes shape iteration-wise specialization.
  • Head-wise localization control: The full-resolution control has weaker and less structured head-wise specialization than SpiralFormer across key-marginal entropy and Local Attention Mass variability.Dynamic heads are the top 40% by cross-loop range for each metric, with heatmaps min–max normalized to [0, 1].
  • Attention statistics over all heads: All-head analysis preserves the qualitative shifts: increasing resolution lowers key-marginal entropy and raises Local Attention Mass.Including weakly varying heads dilutes effect sizes, but the trends remain visible under the identical 500-sequence evaluation protocol.
  • Attention statistics over all heads: The all-head persistence supports the claim that the hierarchical dependency pattern is not an artifact of filtering to dynamic heads.
Loading 2602.11698v2…