Generative Recursive Reasoning

In plain words

A recursive reasoning model solves a problem by running the same small network over and over on a hidden "scratchpad" of its own state, refining it step by step, instead of writing out a long answer. The trouble is that these models are deterministic: the same input always follows the same path and lands on one answer. When a puzzle has several valid answers, the model tends to find only one of them.

This paper, GRAM, adds controlled randomness to every refinement step. Each run now wanders along a different path, so one puzzle can yield several different answers. During training the model sees the correct answer and learns how much randomness to add and where, then at test time it drops the answer and draws its own paths. Running several paths side by side and taking the most common answer is a second way to get better results, besides running more steps. The same machinery can also start from an empty input and produce new examples, such as Sudoku boards.

TL;DR

GRAM turns a deterministic recursive reasoner into a latent-variable generative model. Each latent transition adds a learned Gaussian perturbation, trained with amortized variational inference and truncated gradients. At inference, N trajectories are drawn from the prior in parallel and majority-voted §2.3. Author-reported results beat the recursive baselines (Looped Transformer, HRM, TRM) on Sudoku-Extreme, ARC-AGI-1 and ARC-AGI-2, and on N-Queens and Graph Coloring coverage. Use it when your task has several valid outputs or benefits from inference-time sampling; expect costly sequential training.

Problem

Deterministic recursive models such as Looped Transformers, HRM and TRM follow a single latent trajectory and converge to a single prediction, so they cannot represent several valid solutions for one input §1. On multi-solution constraint tasks, the training targets for one input conflict, and a deterministic model must collapse them into one output §4.2.

Extra depth is also a poor scaling axis for latency: recursion steps run one after another. A matched-compute comparison on Sudoku-Extreme reports 184 ms per puzzle for GRAM with 16 iterations and 20 samples, against 2,755 ms for TRM with 320 iterations §4.1 tab. 10.

Assumptions: the method targets small, structured problems with hard constraints. All experiments are controlled benchmarks, and the paper compares against small recursive baselines, not frontier LLMs, because their training data and budgets are not comparable §1.

Method

Encoder. The input x is embedded once, e_x = f_enc(x), and the embedding is reused by every transition §2.1. The official code scales the token embedding by sqrt(D) and prepends 16 puzzle-embedding tokens gram.py.

Latent state. The state is a pair z = (h, l). The high-level part h is updated once per transition; the low-level part l is refined K times inside each transition §2.1.

One transition. Given the previous state z_{t-1} = (h_{t-1}, l_{t-1}):

1. Refine the low-level state K times with the high-level state held fixed: l_{t,k} = f_L(h_{t-1}, l_{t,k-1}, e_x). This part is deterministic §2.1.
2. Compute a deterministic high-level proposal u_t = f_H(h_{t-1}, l_t) §2.1.
3. Sample guidance ε_t from a Gaussian whose mean and variance are functions of u_t: ε_t ~ N(μ_θ(u_t), σ_θ²(u_t) I). The mean steers the trajectory and the variance sets how much it explores §2.1.
4. Set h_t = u_t + ε_t and z_t = (h_t, l_t) §2.1.

Noise enters only at the high level. A placement ablation that injects noise at every low-level step fails outright on N-Queens 8x8 app. D.8.

Supervision step and depth. A supervision step is T transitions. A full run applies N_sup supervision steps, each one starting from the final state of the previous one §2.1. The decoder reads only h. Depth can be raised at inference by running more supervision steps §2.3 fig. 4.

Variational training. Training maximizes an ELBO over latent trajectories with a posterior q_φ(ε | u, y) that sees the target and a prior p_θ(ε | u) that does not. The posterior and prior are separate SwiGLU heads of the same shape, and the posterior head takes [u_t ; repr(y)] as input §2.2 app. B.1. Because the decoder reads only the terminal state, the likelihood depends only on the terminal state §2.2.

Truncated gradients. Exact gradients would need backpropagation through all T × N_sup transitions. Instead, gradients flow only through the last transition of each supervision step, and the state is detached between supervision steps. The resulting surrogate per step is the reconstruction term minus the KL between posterior and prior guidance at the last transition. The paper calls this biased but memory-efficient §2.2 app. A.2.

Halting (optional). An auxiliary head maps the state to a halting decision. At inference the trajectory stops when the halt probability exceeds one half app. A.3. The paper uses a learned Q-style target for it app. A.3.

Width (parallel sampling). At inference, N trajectories are drawn from the prior, each decoded to an answer, and the most frequent answer is returned §2.3 alg. 2. No extra parameters or training are involved. A learned best-of-N selector based on a value head is discussed as an alternative app. A.4.

Unconditional generation. The same process with an empty input embedding defines p_θ(x). The Sudoku generator starts from an all-blank board and the image generator from an empty conditioning signal §2.1 §4.3.

# GRAM training step  [alg. 1]  [§2.2]
# Input:  input x, target y; N_sup, T, K, beta
# Output: updated parameters

def transition(state, input_emb, noise_head):
    # Refine the low-level state K times, high-level held fixed  [§2.1]
    high, low = state
    for k in range(K):
        low = f_low(low + high + input_emb)
    # Deterministic high-level proposal  [§2.1]
    proposal = f_high(high + low)
    # Stochastic guidance from the given head: prior or posterior  [§2.1]
    noise = sample(noise_head(proposal))
    return (proposal + noise, low), proposal

def train_step(input_x, target_y):
    # Encode the input once; every transition reuses it  [§2.1]
    input_emb = encode(input_x)
    # Posterior guidance is conditioned on the target in training  [§2.2]
    posterior = posterior_head(target_y)
    prior = prior_head
    state = initial_state
    for step in range(N_sup):
        # The first T-1 transitions get no gradient  [§2.2] [alg. 1]
        for t in range(T - 1):
            state, _ = stop_gradient(transition(state, input_emb, posterior))
        # Only the last transition of the step is differentiated  [§2.2]
        state, proposal = transition(state, input_emb, posterior)
        # Reconstruction term on the decoded high-level state  [alg. 1]
        loss = cross_entropy(decode(state.high), target_y)
        # KL between posterior and prior guidance at the last step  [§2.2]
        loss += beta * kl(posterior(proposal), prior(proposal))
        # Halting head learns whether the current answer is correct  [app. A.3]
        loss += bce(halt(state.high), is_correct)
        update(parameters, loss)
        # No gradient flows across supervision steps  [alg. 1]
        state = stop_gradient(state)
        if halt_probability(state) > 0.5:
            break
    return loss

# GRAM inference with parallel trajectory sampling  [alg. 2]  [§2.3]
# Input:  input x; N trajectories, N_sup, T, K
# Output: majority-voted prediction

def predict(input_x):
    # Encode the input once, as in training  [alg. 2]
    input_emb = encode(input_x)
    predictions = []
    # Trajectories are independent, so they can run as one batch  [app. D.2]
    for i in range(N):
        state = initial_state
        for step in range(N_sup):
            for t in range(T):
                # Target is unknown at test time: use the prior  [alg. 2]
                state, _ = transition(state, input_emb, prior)
        predictions.append(argmax(decode(state.high)))
    # Most frequent answer wins; no parameters added  [§2.3]
    return most_common(predictions)

Hyperparameters

Values are taken from the paper unless the row says code. Where the paper is silent, the table says so.

Results

Author-reported. Setups: Sudoku-Extreme and N-Queens and Graph Coloring use the paper's own reproduction of each baseline, averaged over 3 seeds. ARC-AGI rows for HRM and TRM are taken from their papers tab. 8.

Ablation on Sudoku-Extreme and N-Queens 8x8 at N=20 tab. 3. Removing the stochastic guidance drops GRAM to the TRM result on Sudoku (83.6) and to 66.8 on N-Queens. Stochasticity without guidance gives 96.0 on Sudoku but 62.1 on N-Queens. Guidance without stochasticity fails completely (0.0 on both). A direct-prediction variant gives 70.7 on Sudoku and 98.8 on N-Queens. Random initialization of TRM gives 91.2 and 71.4 tab. 3 app. D.7.

Figures are described here, not reproduced. fig. 1 contrasts three runs of a deterministic model, which collapse onto one trajectory, with GRAM runs that reach two different valid N-Queens solutions. fig. 3 shows puzzle-benchmark bars with the N=1 and N=20 splits for Sudoku-Extreme. fig. 4 plots accuracy against recursion depth for several sample counts and against FLOPs.

Compute tab. 9: GRAM costs about 3.5% more FLOPs per iteration than TRM (0.0265 vs 0.0256 TFLOPs).

Paper vs code

Code: ahn-ml/gram-release at f2e32b9b11737890bfd4f39ead4cbea90835f91b. The project page linked from the abstract points to it. Its LICENSE is MIT with a 2025 Samsung Electronics copyright line, which should be checked before reuse. Configuration defaults are in config/, and the per-task training scripts override them. The discrepancies below come from reading the scripts and the model file. The paper and the code differ in the following ways:

• Loss weights and cross-entropy variant. Paper: cross-entropy, KL with weight β, and a halting BCE with weight 1 alg. 1. Code: the language-model loss uses stablemax cross-entropy by default, and the total loss is lm_loss + 0.5 * (q_halt_loss + q_continue_loss) + kl_coef * kl_loss + value_coef * value_loss gram.yaml losses.py losses.py.
• Halting target. Paper: halting label is whether the prediction equals the target, 1[ŷ = y]. Code: for Sudoku, N-Queens and Graph Coloring, the label is a task verifier that accepts any valid solution, not only the one in the data losses.py. Otherwise it is label equality.
• Where the noise goes. Paper: ε is added directly to the hidden state, with mean and variance computed from u_t. Code: the variational layer samples a 256-dimensional latent per position, projects it with a linear layer to the model width, and adds the result to h gram.py gram.py. The code also predicts log-variance, not variance gram.py.
• Initial state. Paper: sample z_0 once from N(0, I) and keep it fixed. Code: a truncated normal at ±2 standard deviations, scaled to std 1, saved as a buffer gram.py common.py.
• FFN width. Paper: FFN hidden size 512 app. B.1. Code: SwiGLU expansion 4 gram.yaml. The width formula gives 1,365 before rounding to a multiple of 256, so each gate and up projection is 1,536 wide, by our computation from layers.py.
• Positional encoding for Sudoku. Paper: RoPE for the encoder app. B.1. Code: Sudoku script sets pos_encodings=none and mlp_t=True (SwiGLU core instead of attention) sudoku.sh.
• Weight decay and puzzle-embedding learning rate for ARC. Paper: weight decay 1.0 for all tasks app. B.2. Code: the ARC scripts do not override weight decay or puzzle-embedding learning rate, so they use the defaults of 0.1 and 1e-2 arc1.sh cfg_pretrain.yaml cfg_pretrain.yaml.
• ARC hardware. Paper: 8 GPUs for ARC-AGI tab. 7. Code: ARC scripts launch with 4 processes per node arc1.sh.
• N-Queens 8x8 epochs. Paper: 3K epochs tab. 7. Code: 5000 epochs nqueens_8.sh.
• KL balance. Paper: coefficient 0.8 applies generally app. B.2. Code: the default is off gram.yaml. The scripts we read (Sudoku, N-Queens 8x8, MNIST, ARC-AGI-1) set 0.8 sudoku.sh nqueens_8.sh mnist.sh arc1.sh. The Sudoku-generation and ARC-AGI-2 scripts were not checked for this setting.
• Value head. The Sudoku script enables the value head for the learned selector sudoku.sh. The main text of the paper does not use it for the reported numbers.
• ARC voting. Paper: majority vote §2.3. The ARC pass@k evaluator sorts candidates by count, then by average Q value, and treats the top entry as the majority arc_pass_k.py. The paper's reported ARC numbers use pass@2 over augmented inputs app. C.1, not pure majority vote.
• Evaluation weights. The scripts set EMA on, but this card did not check whether evaluation loads the EMA weights.

Pitfalls

• Baseline numbers differ from published ones. The text quotes published TRM at 87.4% and HRM at 55.0% on Sudoku-Extreme, while the comparison table uses the paper's own 3-seed reproductions, 83.6% and 56.8% §4.1 tab. 8. Quote the reproduced numbers when comparing.
• N=20 is not a fair comparison with N=1 baselines. Majority voting uses 20 forward passes. The paper reports a matched-FLOPs comparison that favors GRAM tab. 10, but any N=1 comparison must say so.
• ARC-AGI numbers do not use parallel sampling. They use 1,000 augmented inputs with pass@2 app. C.1. Parallel sampling is reported separately app. D.3.
• N-Queens 10x10 accuracy is not a GRAM win. At N=1, the AR baseline scores 90.0 against GRAM's 89.7 tab. 1. GRAM leads on coverage and at N=20.
• Halting label may accept different solutions. For tasks with several valid answers, the code's halting target accepts any valid solution losses.py. Changing the verifier changes what the model learns to halt on.
• Training cost. Deep supervision makes training sequential. The paper says this limits training efficiency against Transformers and is a barrier to scaling §5. Looped TF on ARC-AGI was not trained because one run was estimated at about 97 days app. C.1.
• Width spends compute, not time. GRAM's N=20 run uses about the same FLOPs as 320 TRM iterations, but batching the samples gives roughly 15 times lower latency tab. 10. Latency is measured with one puzzle in flight on the paper's hardware.
• Stability of the KL term. The standard KL divides by the exponential of the prior log-variance. The code guards only the case where variance is forced to zero losses.py. The paper uses KL balance to prevent posterior collapse app. B.2. Keep log-variance in a sane range if you port this.
• Precision. The recursive core runs in bfloat16 gram.yaml. The variational layer and the stablemax loss run in float32 or float64 gram.py losses.py.
• Stochastic training trajectories. During training, the posterior guidance is also used in the no-gradient cycles before the last one gram.py. Implementations that use the prior there will not match the paper's training procedure alg. 1.
• Learning-rate warmup is a code default. The paper does not state a warmup, but the code uses 2000 steps with a constant rate afterwards cfg_pretrain.yaml. Apply the same schedule when reproducing.
• Generation is measured on derived data. Sudoku generation trains on 50K complete solutions derived from the Sudoku-Extreme training split, with no augmentation. Its 99.05% validity is over 100K samples, and all valid boards were unique app. B.2 app. D.6.
• Figures carry no numbers here. The figure descriptions above come from captions and text. Use the tables for numeric claims.
• Code license. The LICENSE copyright line names Samsung Electronics. Check the repository's history before copying more than a few lines from it.

Lineage

• Builds on: Looped Transformers, arXiv:2311.12424 §1; HRM, arXiv:2506.21734; TRM, arXiv:2510.04871. GRAM keeps the deep-supervision and halting setup of HRM and TRM and makes the latent transition stochastic §2.1 app. A.3.
• Compared against: Looped TF, HRM, TRM (deterministic recursion); AR and MDLM (generative models); D3PM and a VAE (unconditional generation) §4. None of these baselines has a card in this wiki yet.
• Superseded by: not known from content already in this wiki.