Adaptive Computation Time (ACT)

Summary: ACT (Graves 2016, adapted in HRM/TRM) lets models dynamically choose computation depth per input. A Q-head predicts "halt" vs "continue" values; Q-learning trains it. During training, segments run up to M_{max}; at inference, M_{max} can be increased for test-time scaling without retraining. Sudoku shows strong scaling; ARC minimal.

HRM/TRM Integration (§2.4, Fig 5)

Q-Head & Halting Policy

  • Q-head: hat{Q}^m = Sigma(theta_Q^T z_H^{mNT}) → (hat{Q}_{halt}, hat{Q}_{continue})
  • State: z^m (final segment state)
  • Actions: {halt, continue}
  • Rewards: halt → mathbf{1}\{hat{y}^m = y\}; continue → 0

Q-Targets (Eq. 185-198)

hat{G}_halt^m = mathbf{1}\{hat{y}^m = y\}

hat{G}_continue^m = begin{cases} hat{Q}_halt^{m+1} & m ≥ M_{max} \\ max(hat{Q}_halt^{m+1}, hat{Q}_continue^{m+1}) & else end{cases}

Q-Loss (Eq. 203-204)

L^m_{ACT} = Loss(hat{y}^m, y) + BCE(hat{Q}^m, hat{G}^m)

Stochastic M_{min} (Eq. 181)

  • With prob varepsilon: M_{min} ~ Uniform(\{2, …, M_{max}\}) — encourages longer thinking
  • With prob 1-varepsilon: M_{min} = 1

Halting Conditions

Halt if: m > M_{max} OR (hat{Q}_halt > hat{Q}_continue AND m ≥ M_{min})

Training vs Inference

Phase M_{max} Behavior
Train Fixed (e.g., 8) Learns halting policy; stochastic M_{min} encourages exploration
Inference Can increase (e.g., 8 → 16) Model handles more steps; accuracy improves (Fig 5c)

Results (Fig 5, Sudoku-Extreme-Full)

Metric Fixed Steps ACT
Avg compute steps Grows with M_{max} Low & stable even as M_{max} increases
Accuracy Improves with M_{max} Matches fixed at lower avg compute
Inference scaling Requires retraining Zero-shot: train M_{max}=8, test M_{max}=16 → gains

Task-dependent:

  • Sudoku: Strong scaling (needs backtracking depth)
  • ARC: Minimal (few transformations needed)

Stability (Section 6)

Deep Q-learning typically needs replay buffers + target networks. HRM achieves stability via:

  1. Post-Norm architecture (RMSNorm) — bounded activations
  2. AdamW — L_∞ constrained optimization (||theta||_∞ ≤ 1/lambda)
  3. Bounded params + weight decay + post-norm → Q-learning convergence (Gallici et al. 2025)

Pseudocode (from HRM Fig 4 + TRM Fig 2)

# Per segment inside deep supervision loop
z, y_pred, q = hrm(z, x)

# Task loss
loss = ce_loss(y_pred, y_true)

# ACT halt loss
target_halt = (y_pred == y_true).float()
loss += 0.5 * bce_loss(q[:,0], target_halt)

# ACT continue loss (extra forward)
_, _, q_next = hrm(z, x)
if step == N_sup - 1:
    target_cont = torch.sigmoid(q_next[:,1])
else:
    target_cont = torch.sigmoid(torch.max(q_next[:,0], q_next[:,1]))
loss += 0.5 * bce_loss(q[:,1], target_cont)

# Early stopping
if q[:,0] > q[:,1]:
    break

Key Claims with Sources

Claim Source Locator Confidence
Q-head on H-state predicts halt/continue arxiv-2506.21734 §2.4, Eq. 175-179 0.95
Q-learning with stochastic M_{min} arxiv-2506.21734 §2.4, Eq. 181 0.95
Halting: Q_halt > Q_continue AND m ≥ M_{min} arxiv-2506.21734 §2.4 0.95
Inference-time scaling: increase M_{max} without retraining arxiv-2506.21734 §6, Fig 5(c) 0.95
Sudoku: strong scaling; ARC: minimal arxiv-2506.21734 §6, Fig 5 0.95
Stability via Post-Norm + AdamW (bounded params) arxiv-2506.21734 §6 0.9
TRM uses identical ACT mechanism arxiv-2510.04871 Fig 2, §2.1 0.9

Related Concepts

Sources

  • arxiv-2506.21734: HRM paper (§2.4, §6, Figure 5)
  • arxiv-2510.04871: TRM paper (Figure 2, §2.1)