URM-Energy-Stopping: A New Approach for Reasoning Models Using Energy Convergence to Replace Adaptive Computation Time

This project explores replacing the Adaptive Computation Time (ACT) mechanism in the Universal Reasoning Model (URM) with an energy-based stopping criterion. The core idea is to use an energy function E(input, output) to score prediction quality and stop iteration when energy converges. Compared to ACT's learned stopping probabilities, this method has advantages such as a principled stopping mechanism, built-in MCMC iterative optimization, and energy scores as a confidence metric.

The reasoning ability of large language models is a core topic in AI research. URM achieved a 53.8% pass@1 score on the ARC-AGI benchmark; its cyclic inductive bias and strong nonlinearity are crucial for reasoning tasks, but the ACT mechanism it uses is a learned binary signal. This project asks: Can we replace this learned stopping mechanism with a more principled physical intuition (energy-based model)?

Core Idea

Inspired by Hoover et al.'s 2024 Energy-Based Transformers, we shift the stopping decision from learning when to stop to measuring when to stabilize: introduce an energy function E(input, output) to score prediction quality, use MCMC optimization to find the minimum energy point, and stop when the energy change is below a threshold.

Technical Implementation

• Energy-based URM model: includes MCMC optimization loop, learnable step size, and Langevin dynamics noise;
• Replay buffer: stores diverse MCMC training trajectories to stabilize training;
• Contrastive energy loss: boundary-based loss pushes the energy of correct inputs below that of incorrect ones to prevent energy collapse;
• Configuration management: uses Hydra to manage hyperparameters (e.g., energy convergence threshold, noise standard deviation, etc.).

Trained on the ARC-AGI-1 dataset using a 10×10 downsampled grid and a single RTX3090:

• URM baseline: fast convergence but severe overfitting;
• Energy v0: energy collapse (energy head is constant for all inputs and outputs);
• Energy v1: adding contrastive loss fixes the collapse, and the energy function learns to distinguish correct/incorrect outputs;
• Energy v2: after removing ACT loss, MCMC takes only 1-2 steps, requiring minimum step constraints and threshold tuning.

Key lessons: Contrastive loss is crucial; MCMC needs minimum step constraints; small grids are prone to overfitting and require data augmentation.

Compared to ACT, the advantages of the energy-based method are:

1. Principled stopping mechanism: Energy convergence has clear physical meaning (local energy minimum, similar to physical system stability);
2. Built-in confidence metric: Energy scores directly reflect prediction confidence (lower energy = higher confidence), supporting uncertainty quantification;
3. MCMC iterative optimization: Predictions can be further optimized via gradient descent during inference (similar to iterative denoising in diffusion models);
4. Architecture compatibility: Seamlessly integrates with standard Transformers without modifying the backbone network.

This research is in the early stage, and areas for improvement include:

• Model scale adjustment: The current configuration (hidden dimension 64-128, 2 layers) needs to be adapted to small grids or extended to 30×30 large grids;
• MCMC step tuning: Enforce minimum steps for sufficient iteration;
• Data augmentation: Enhance small grid data to reduce overfitting;
• Hyperparameter search: Systematically optimize contrastive loss weight, boundary values, etc.;
• Fair comparison: Compare energy stopping and ACT performance on matched architectures.

URM-Energy-Stopping is an exploratory project that attempts to replace the ACT mechanism with an energy-based stopping criterion. Although in the early stage, it demonstrates the potential of energy-based methods in reasoning models: principled stopping mechanism, built-in confidence metric, and natural iterative optimization capability. It provides a valuable experimental platform and reference implementation for researchers working on inference-time computation expansion and reasoning model optimization.