Universal Transformers Need Memory: Depth-State Trade-offs in Adaptive Recursive Reasoning

• URL: http://arxiv.org/abs/2604.21999v2
• Date: Mon, 27 Apr 2026 14:17:07 GMT
• Title: Universal Transformers Need Memory: Depth-State Trade-offs in Adaptive Recursive Reasoning
• Abstract summary: We study learned memory tokens as computational scratchpad for a single-block Universal Transformer (UT) on SudokuExtreme.<n>We find that memory tokens are empirically necessary: across all configurations tested, no configuration without memory tokens non-trivial performance.
• Score: 0.0
• License: http://creativecommons.org/licenses/by-sa/4.0/
• Abstract: We study learned memory tokens as computational scratchpad for a single-block Universal Transformer (UT) with Adaptive Computation Time (ACT) on Sudoku-Extreme, a combinatorial reasoning benchmark. We find that memory tokens are empirically necessary: across all configurations tested -- 3 seeds, multiple token counts, two initialization schemes, ACT and fixed-depth processing -- no configuration without memory tokens achieves non-trivial performance. The optimal count exhibits a sharp lower threshold (T=0 always fails, T=4 is borderline, T=8 reliably succeeds for 81-cell puzzles) followed by a stable plateau (T=8-32, 57.4% +/- 0.7% exact-match) and collapse from attention dilution at T=64. During experimentation, we identify a router initialization trap that causes >70% of training runs to fail: both default zero-bias initialization (p ~ 0.5) and Graves' recommended positive bias (p ~ 0.73) cause tokens to halt after ~2 steps at initialization, settling into a shallow equilibrium (halt ~ 5-7) that the model cannot escape. Inverting the bias to -3 ("deep start," p ~ 0.05) eliminates this failure mode. We confirm through ablation that the trap is inherent to ACT initialization, not an artifact of our architecture choices. With reliable training established, we show that (1) ACT provides more consistent results than fixed-depth processing (56.9% +/- 0.7% vs 53.4% +/- 9.3% across 3 seeds); (2) ACT with lambda warmup achieves matching accuracy (57.0% +/- 1.1%) using 34% fewer ponder steps; and (3) attention heads specialize into memory readers, constraint propagators, and integrators across recursive depth. Code is available at https://github.com/che-shr-cat/utm-jax.

Related papers

• On the Design of Qwen3.8-Next Architecture: Evaluation, Efficiency, and Training Stability [57.03393882933515]
  We describe the architecture and ablations of Qwen3.8-Flash-Next, a sparse mixture-of-experts model with 125B parameters, 6B activated per token, and 51B parameters of n-gram embedding tables held off the accelerator.<n>On fourteen pre-training benchmarks the model leads the 397B-A17B predecessor on eight and trails it on the rest by at most 2.6 points, at 1/3 the activated parameters, 1/3 the training tokens, and roughly 1/9 the training FLOPs.
  arXiv  Detail & Related papers  (2026-08-31T06:35:07Z)
• Compress and Forget: bitsandbytes Quantization Amplifies Proactive Interference in LLMs [0.23438564092609357]
  Proactive interference (PI) is a documented failure mode in large language models in which retrieval of a repeatedly overwritten value degrades as prior overwrites accumulate.<n>PTQ is now the default deployment path for open-weight models, yet its effect on this failure mode has not been tested.<n>We evaluate three precision levels (FP16, INT8, INT4/NF4, via bitsandbytes) across three architecturally distinct instruction-tuned models.
  arXiv  Detail & Related papers  (2026-08-19T06:17:13Z)
• Per-Token Fixed-Point Convergence in Depth-Recurrent Transformers [0.0]
  We show a training-free rule that halts each token once its output stabilizes attains uniform depth-8 quality at 4.94 average loops.<n>A linear router trained on convergence labels harvested from the same model requires nearly full depth and yields no reduction.<n>We report average depth as a FLOP proxy with a three-point wall-clock bracket rather than a realized speedup.
  arXiv  Detail & Related papers  (2026-07-15T23:38:44Z)
• CARVE: Content-Aware Recurrent with Value Efficiency for Chunk-Parallel Linear Attention [0.0]
  Recurrent models must forget in order to remember, yet the state of the art decides what to erase without consulting what is stored.<n>This memory-blind gating is one of three coupled defects in the leading delta-rule architecture (GDN-2)<n>We introduce CARVE, which resolves all three problems through one principle: erase only on the key axis.
  arXiv  Detail & Related papers  (2026-06-25T16:16:51Z)
• Automated Proving of Shannon-Type Entropy Inequalities via Fine-Tuned Language Models and Guided Tree Search [50.16356451328644]
  Proving Shannon-type entropy inequalities is a fundamental task in information theory.<n>We investigate whether small-scale large language models, fine-tuned on atomic proof steps, can automate this process.<n>We show that GPT-5.5 solves 1.7% samples under zero-shot prompting while Psitip solves 33.3% samples.
  arXiv  Detail & Related papers  (2026-06-04T05:43:12Z)
• Right Makes Might: Aligning Verified Hidden States Empowers RL Reasoning [55.264863369127774]
  Current methods reduce each correct rollout to a single reward bit, ignoring the geometric structure shared among their hidden states.<n>We propose Hidden-Align, an auxiliary loss function that aligns the last-layer hidden states of correct rollouts at the anchor token during RL training, with zero overhead in both training and inference.<n>On eight mathematical reasoning benchmarks, Hidden-Align improves average pass@1 over the DAPO baseline by 3.8, 6.2, and 5.4 percentage points on Qwen3-1.7B, 4B, and 14B respectively, with consistent pass@k gains across all three scales.
  arXiv  Detail & Related papers  (2026-06-02T06:51:15Z)
• ChunkFT: Byte-Streamed Optimization for Memory-Efficient Full Fine-Tuning [58.54940026861599]
  textscChunkFT is a memory-efficient fine-tuning framework.<n>textscChunkFT enables gradient computation for arbitrary sub-tensors without modifying the network architecture.<n>textscChunkFT consistently outperforms existing memory-efficient baselines.
  arXiv  Detail & Related papers  (2026-05-20T13:44:44Z)
• Less Is More: Cognitive Load and the Single-Prompt Ceiling in LLM Mathematical Reasoning [0.0]
  We study prompt engineering for formal mathematical reasoning in the context of the SAIR Equational Theories Stage 1 competition.<n>The task requires deciding whether one equational law implies another over all magmas.<n>Over five weeks, we designed, tested, and analyzed more than 40 prompt variants.
  arXiv  Detail & Related papers  (2026-04-20T22:55:23Z)
• THEIA: Learning Complete Kleene Three-Valued Logic in a Pure-Neural Modular Architecture [0.0]
  THEIA is a 2.75M modular neural architecture that learns the complete Kleene three-valued logic (K3) truth table from task data without external symbolic inference or hand-encoded K3 gate primitives.<n> Transformer baselines also reach >99% on all 39 rules, and flats match THEIA on Phase-1 accuracy within 0.04pp.
  arXiv  Detail & Related papers  (2026-04-13T10:44:15Z)
• ImplicitMemBench: Measuring Unconscious Behavioral Adaptation in Large Language Models [60.14219417402433]
  Existing memory benchmarks for LLM agents evaluate explicit recall of facts, yet overlook implicit memory where experience becomes automated behavior without conscious retrieval.<n>We introduce ImplicitMemBench, the first systematic benchmark evaluating implicit memory through three constructs drawn from standard cognitive-science accounts of non-declarative memory.<n>Our 300-item suite employs a unified Learning/Priming-Interfere-Test protocol with first-attempt scoring.
  arXiv  Detail & Related papers  (2026-04-09T10:26:32Z)
• DeepPrune: Parallel Scaling without Inter-trace Redundancy [53.62015294143274]
  Over 80% of parallel reasoning traces yield identical final answers, representing substantial wasted computation.<n>We propose DeepPrune, a novel framework that enables efficient parallel scaling through dynamic pruning.<n>Our work establishes a new standard for efficient parallel reasoning, making high-performance reasoning more efficient.
  arXiv  Detail & Related papers  (2025-10-09T17:24:54Z)
• VeriThinker: Learning to Verify Makes Reasoning Model Efficient [52.74493506816969]
  Large Reasoning Models excel at complex tasks using Chain-of-Thought (CoT) reasoning.<n>Their tendency to overthinking leads to unnecessarily lengthy reasoning chains.<n>We introduce VeriThinker, a novel approach for CoT compression.
  arXiv  Detail & Related papers  (2025-05-23T14:17:56Z)
• Retrosynthetic Planning with Dual Value Networks [107.97218669277913]
  We propose a novel online training algorithm, called Planning with Dual Value Networks (PDVN) PDVN alternates between the planning phase and updating phase to predict the synthesizability and cost of molecules. On the widely-used USPTO dataset, our PDVN algorithm improves the search success rate of existing multi-step planners.
  arXiv  Detail & Related papers  (2023-01-31T16:43:53Z)
• Peeling the Onion: Hierarchical Reduction of Data Redundancy for Efficient Vision Transformer Training [110.79400526706081]
  Vision transformers (ViTs) have recently obtained success in many applications, but their intensive computation and heavy memory usage limit their generalization. Previous compression algorithms usually start from the pre-trained dense models and only focus on efficient inference. This paper proposes an end-to-end efficient training framework from three sparse perspectives, dubbed Tri-Level E-ViT.
  arXiv  Detail & Related papers  (2022-11-19T21:15:47Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.