Source-linked AI summary
Tiny Recursive Message Passing
Tommaso Salvatori, Matteo Forasassi, Alexis Toumi, Giovanni de Felice
TL;DR
Structured reasoning models must exploit task structure and generalize beyond training instances, while full-trajectory BPTT constrains message-passing depth. This paper adapts recursive training and inference procedures to GNNs, combining them with graph structure; the resulting models match or exceed larger neural baselines, with truncation costs that depend on wiring.
Problem
Neural models for structured reasoning need to exploit problem structure and generalize to larger instances, while full-trajectory BPTT is memory-expensive and limits practical recursion depth.
Source: §1 · p. 1Method
The paper adapts deep supervision, truncated recursion, adaptive halting, and deduction lattices to typed message-passing GNNs.
Source: §1 · p. 2Results
16k–65k-parameter models match or exceed neural baselines up to an order of magnitude larger, and often match or outperform current state of the art on larger zero-shot instances.
Source: §1 · p. 2, §6 · p. 10Takeaways & Limitations
Truncation becomes costlier as graph diameter grows, making wiring a trade-off between robustness to truncation and transfer across problem sizes.
Source: §6 · p. 10
Abstract
from arXiv · showhide
Message passing neural networks have shown remarkable performance in algorithmic and constraint reasoning, due to their ability to encode the structure of the task directly in the architecture. Recent work has shown that better results can be achieved using recursive transformers trained with deep supervision and truncated recursions, which make extremely deep loops stable to train, and equipped with test-time procedures such as discrete deduction lattices. In this work, we apply the same procedures to message passing networks, and show that the two ingredients are complementary, and lead to parameter efficient models: we obtain about $99\%$ test accuracy on Sudoku-Extreme with a model of $16$k parameters that trains in about $10$ minutes on an H100, or in $9$ hours on a CPU. We further show that the cost of truncating gradients grows with the number of rounds information needs to cross the graph, making the choice of wiring a trade-off between robustness to truncation and transfer across problem sizes. We conclude by showing that our models transfer zero-shot to larger instances than those seen during training, outperforming popular neural baselines on N-queens and Boolean satisfiability with a fraction of their parameters.
1 Introduction
The paper asks whether recursive training and inference machinery can make structured message-passing models more parameter-efficient while preserving generalization to larger tasks. It adapts that machinery to recursive GNNs and reports complementary gains from combining it with graph structure, alongside a wiring-dependent cost of gradient truncation.
1 Introduction
- The paper provides a training and inference procedure for recursive GNNs, whose node memories can be paused and resumed across computation segments.Its architectural design is close to recurrent relational networks; the contribution also analyzes how the procedure interacts with graph wiring.Source: §1 · p. 2
- Deep supervision, detached recursion, and adaptive halting carry over to recursive GNNs, while adaptive halting and deduction lattices require reformulation.The reformulations use per-vertex halting logits pooled with a smooth minimum and a lattice over residual SAT formulas.Source: §1 · p. 2
- Combining graph structure with recursive machinery outperforms a recursive lattice transformer and a BPTT-trained recursive GNN with the same graph on Sudoku-Extreme and maze-unique.The transformer has 32k parameters; the comparison supports the complementarity of the two ingredients.Source: §1 · p. 2
- Gradient truncation becomes costlier as graph propagation distance grows: dense wirings tolerate it, whereas factor graphs need longer gradient windows.Porting a lattice-based method to message passing largely increases robustness to truncated gradients.Source: §1 · p. 2
- The models often match or outperform current state of the art on zero-shot larger N-Queens and SAT instances, and outperform all neural methods reported for the 4-SAT benchmark.N-Queens training uses 10 × 10 boards and evaluates up to 16 × 16; SAT training uses medium-size SR and 3-SAT formulas and evaluates on hard ones.Source: §1 · p. 2
2 Related Works
Earlier message-passing models exploit graph structure and weight sharing, while recursive reasoning models have developed ways to train deep computations without backpropagating through the full trajectory. This work draws on both lines of work and adapts their procedures to message passing.
Recurrent message passing
- Recurrent message-passing precedents include recurrent relational networks for Sudoku and NeuroSAT for SAT, but both use full-trajectory BPTT, practically limiting depth.Weight-tied networks can run longer at test time, which prior work has used to solve larger or harder instances.Source: §2 · p. 2
- Figure 1 depicts typed modules with shared weights within each type and contrasts an invariant pooling module with latent memory against an equivariant Deep Sets module without memory.The pooling module computes output when iteration converges; the Deep Sets module does not carry latent memory.Source: Figure 1 · p. 3
Deep recursive computation
- Recursive reasoning work uses detached recursion, deep supervision, and adaptive halting; recent inference procedures add noise or project outputs into a lattice to guide search.This paper implements, adapts, and tests these procedures for message-passing models.Source: §2 · pp. 2–3
3 Background
A typed message-passing model repeatedly updates vertex memories from incoming messages and fixed inputs, then sends new messages along graph edges. This section describes two modules for interchangeable neighbors and contrasts the usual full-BPTT training pipeline with recursive-model machinery.
3 Background
- At each round, a vertex reads incoming neighbor messages and, when present, its memory, uses a fixed input, and sends new messages to its neighbors.Vertices of the same type share a module; repeating rounds lets information travel across the graph.Source: §3 · p. 3
- Pooling modules handle interchangeable neighbors when a vertex needs one shared answer, summing message encodings with memory and updating the memory through a GRU.The new memory is projected to produce the same outgoing message for every neighbor.Source: §3 · p. 4
- Deep Sets modules handle interchangeable neighbors when a vertex must answer each one individually, using a pooled summary to decode each neighbor's outgoing message.The module is permutation equivariant and has no memory or input.Source: §3 · p. 4
- The standard forward pass embeds discrete inputs, runs a fixed number of rounds, and computes outputs from vertex memories; the paper replaces full BPTT with deep supervision, truncated recursion, and adaptive halting.The usual training gradients are computed through all rounds, which the recursive machinery avoids doing at full cost.Source: §3 · p. 4
4 Methods
The method adapts recursive-model machinery to message passing through nested inner, outer, and deep-supervision loops, with truncated gradients and optional adaptive halting or discrete-lattice inference. A per-vertex answer state carries information across loops.
Answer state
- A per-vertex answer state is concatenated into the pooling GRU, held fixed during the inner loop, and updated at its end.The paper reports that a smaller dedicated update layer worked better than an independent layer or the shared-module TRM version.Source: §4 · p. 4
Looped iterations
- Each deep-supervision step runs K outer-loop cycles of R message-passing rounds, with the answer state held fixed within each inner loop.Increasing loop depth adds computation without changing model parameters or state dimensions.Source: §4 · p. 5
Truncated gradients
- Truncated gradients detach earlier outer-loop states so gradients flow only through the retained final segment, reducing training memory without changing forward dynamics.The model still executes earlier iterations normally; only temporal credit assignment is limited.Source: §4 · p. 5
Adaptive computation time
- The adaptive computation time head learns when to stop a datapoint's supervision loop, using vertex logits pooled by a smooth minimum for a graph-level decision.Its training loss is binary cross-entropy on each vertex logit, against whether that vertex's decoded value is already correct.Source: §4 · p. 5
Confidence selection
- For test-time selection, the model can generate P predictions from identical inputs with Gaussian noise and use the halting head to choose the highest-confidence solution.Noise is added to the answer state after each deep-supervision step.Source: §4 · pp. 4–5
Discrete lattice
- The discrete lattice carries surviving candidates across supervision steps; projection removes low-probability candidates, detects conflicts, and may commit an undecided cell.Conflicted partial solutions are reset and retried; a trial can end in timeout or a wrong prediction.Source: §4 · pp. 5–6
5 Experiments
The experiments compare graph wirings and recursive training procedures on Sudoku, mazes, N-Queens, and SAT. They test accuracy, robustness to truncated gradients, and zero-shot transfer to larger instances.
5 Experiments
- Peer graphs exchange information between constrained cells in one round, while factor graphs use two rounds and longer paths but support joint constraint reasoning and transfer across problem sizes.Factor graphs use k edges for a k-variable constraint rather than k(k −1)/2, and fixed ports enable size transfer at the cost of greater sensitivity to truncation.Source: §5 · p. 6
5.1.1 Results
- About 99% Sudoku-Extreme accuracy is reached by both wirings with full gradients; under truncation, the peer graph stays above 98%, while Maze-unique exceeds 99% only with the longest gradient window.Maze-unique needs more cycles and rounds, consistent with its larger graph diameter.Source: Figure 2 · p. 7, §5.1.1 · p. 8
5.1 Sudoku-extreme and Maze-unique
- A 16k-parameter graph model trains on Sudoku-extreme in about ten minutes on one H100 or nine hours on 32 CPU threads; tuned baselines stay below 20% board accuracy, while the peer graph model outperforms all baselines at that scale.Source: §5.1 · p. 7
5.2 N-Queens
- On 16 × 16 N-Queens boards, the best configuration solves 63% on average; with a verifier, the 12 × 12 configuration reaches near-perfect accuracy up to 20 × 20.The RRN baseline falls to 2% at 16 × 16 for T = 64 and retains 26% for T = 128.Source: Figure 3 · p. 8, §5.2 · p. 9
5.1.1 Results
- With only g rounds carrying gradients, Sudoku-Extreme accuracy more than doubles from g = 12 to g = 24; the needed gradient horizon grows with graph propagation distance and diameter.The lattice provides intermediate supervision that can support improvement with fewer gradient-carrying steps.Source: §5.1.1 · p. 8
5.2 N-Queens
- The N-Queens factor graph gives every cell four ports independent of board size, and the model uses a recurrent hidden state, answer state, and ACT head without a lattice.The experiments train on 10 × 10 boards and evaluate larger boards; the peer graph is not used because it does not generalize to larger boards.Source: §5.2 · p. 9
5.3 Boolean satisfiability (SAT)
- The SAT model uses a variable-clause factor graph with polarity-typed edges and size-independent parameters; its lattice tracks undecided, committed, or conflicting variables and conditions loss on commitments.At inference, the lattice reads assignments from commitments and readout, then repairs them; it returns no wrong assignments, only possible timeouts.Source: §5.3 · pp. 9–10
- On SAT, the lattice is strongest in aggregate, solving over 99% of easy G4SATBench formulas and leading neural methods on both RandCSPBench families, while no-lattice rollouts transfer better to hard formulas.Best-of-64 reaches 4–5× NeuroSAT on hard G4SATBench splits; on RandCSPBench it draws level with QuerySAT at a fifth of its parameters.Source: §5.3 · p. 10, Table 2 · p. 10
6 Conclusion
The paper finds that recursive message-passing models combine graph structure with recursive-transformer training and inference machinery to match or exceed much larger neural baselines. It also identifies a wiring trade-off: gradient truncation becomes costlier as graph diameter grows, while wiring affects transfer across problem sizes.
6 Conclusion
- Models with 16k–65k parameters match or exceed neural baselines up to an order of magnitude larger.The conclusion attributes this to combining recursive-transformer machinery with typed message passing and graph structure.Source: §6 · p. 10
- Truncation becomes more costly as graph diameter grows, making wiring a trade-off between robustness to truncation and transfer across problem sizes.The conclusion identifies learned heuristics inside exact solvers as a natural next step for constraint-structured problems.Source: §6 · p. 10
A Pseudocode
- Inference iterates message-passing cycles over candidate lattices, projects candidates, and restarts from the initial lattice after a conflict.The pseudocode also permits Best-of-P inference, with P = 1 and σ = 0 giving a deterministic single pass.Source: A Pseudocode · p. 15
B Recursive message passing with truncated gradients
- Truncated-gradient training detaches the first D cycles and backpropagates through the final T = C − D cycles; their credited message-passing rounds define the credit horizon.Setting D = 0 recovers full back-propagation through time, with h = CR.Source: B Recursive message passing with… · p. 16
B.1 Architecture
The architecture shares recurrent message-passing modules across vertices and computation steps, so parameter count is independent of the number of cycles and rounds. It supports distinct lattice and adaptive-halting systems, and Sudoku factor and peer graphs route information differently despite encoding the same constraints.
B.1 Architecture
- Shared modules are reused across vertices and every round, cycle, and supervision step, so parameter count does not depend on computation and weights can be iterated longer at test time.All models are recurrent message-passing networks on a fixed graph whose cell vertices are puzzle cells.Source: §B.1 · p. 16
- Sudoku factor and peer graphs encode identical constraints but differ in routing: cell-to-cell exchange takes two rounds in the factor graph and one in the peer graph.Their cell-to-cell diameters are four rounds and two rounds, respectively.Source: §B.1 · p. 16
- The architecture is equivariant to Sudoku’s stated geometric symmetries and the eight square symmetries of the Maze, but not to digit relabellings, which are handled by augmentation.The geometric properties are numerically verified to agreement of 10−12 in double precision.Source: §B.1 · pp. 16–17
- Cell updates pool incoming messages invariantly to neighbor order, update memory, and broadcast a message to adjacent vertices.The memory updates at every round, while the answer state is refreshed once per cycle.Source: §B.1 · p. 17
B.2 Two readouts: the discrete lattice and adaptive halting
- The same network supports lattice and ACT systems, which differ in what carries between supervision steps, output readout, and training.The lattice and ACT systems are introduced as two distinct ways to use the shared network.Source: §B.2 · p. 18
B.2.1 Lattice model
The lattice system represents each cell by its surviving candidate symbols and combines learned propagation with branching and chain search. The ACT system instead decodes symbols directly and can halt when the pooled cell-level confidence indicates the board is ready.
B.2.1 Lattice model
- The lattice model stores the symbols still possible for each cell, with givens as singletons and blank cells initialized to all task-specific symbols.Its recurrent state is rebuilt at every supervision step, so only the lattice carries between steps.Source: §B.2.1 · p. 18
- Candidates are only removed, making the lattice a monotone DPLL-style solver with learned propagation and branching heuristics.The training mix includes raw puzzles, correct partial fills, and corrupted partial fills so the model learns completion and contradiction recognition.Source: §B.2.1 · pp. 18–19
- At test time, independent stochastic chains project candidates; conflict chains restart, and the first chain to solve supplies the answer, with timeouts counted as failures.The stated compute budget caps the number of supervision steps per puzzle.Source: §B.2.1 · p. 19
B.2.2 ACT model
- The ACT model decodes cells from the symbol head without a lattice, projection, or search, and its smooth-minimum halting head requires every cell to be individually confident.Its halting head reads the answer state without shaping the representation.Source: §B.2.2 · p. 19
- ACT training carries and detaches recurrent state between supervision steps, while inference supports fixed-compute, halting, and best-of-k rollout selection.Best-of-k selects using model confidence; pass@k is reported separately as an oracle upper bound.Source: §B.2.2 · p. 19
B.3 The truncated gradient
Truncated gradients reduce the differentiated portion of a recurrent message-passing trajectory: early cycles run without a computational graph, and gradients cover only the final window. Forward computation remains unchanged, while memory and backward cost scale with the credit horizon.
B.3 The truncated gradient
- The first D cycles are detached, and training minimizes losses from cycles D+1 through C while treating the state at cycle D as constant.The clue embedding is recomputed with gradients at the start of the differentiated window, and each differentiated cycle is supervised.Source: §B.3 · p. 20
- Truncation leaves forward computation unchanged and makes activation memory and backward cost scale with h rather than CR.At test time the network is identical to one trained with full gradients; D affects how its parameters were obtained.Source: §B.3 · p. 20
- In both systems, one optimization step has credit horizon h; for the peer-graph Sudoku example, five of sixty message-passing rounds receive gradients.The example uses C × R = 12 × 5 and D = 11.Source: §B.3 · p. 20
B.4 Sudoku-Extreme
The Sudoku-Extreme experiments compare lattice and ACT recipes under different gradient horizons, with results showing strong accuracy from truncated-gradient lattice training and improved ACT performance from longer supervision steps. Longer credit horizons, however, are harder to train.
Data and augmentation
- The benchmark uses 1,000 training puzzles, 500 validation puzzles, and 1,000 test puzzles; each puzzle has a unique solution.Training puzzles receive fresh uniform digit relabelings; this is the only augmentation.Source: §B.4 · p. 20
Lattice results
- 99.1% test accuracy at 24k parameters lets the truncated peer graph exceed the full-gradient reference, while its tuned 32k recipe falls to 98.1% versus 99.4%.The 16k factor graph exceeds 99% with full gradients; at 12k it scores 82.9% versus 98.5% for the peer graph, even with twice the credit horizon.Source: §B.4 · p. 21, Table 5 · p. 21, Table 6 · p. 22
ACT results
- 98.7% mean exact test accuracy with deterministic halting rises to 99.9% with best-of-64 selection, matching the pass@64 oracle.When a correct answer appears among rollouts, the halting head finds it.Source: §B.4 · p. 22
- Longer credit horizons are harder to train: the 8 × 12, h = 24 arm diverged at peak learning rate 1.5 × 10−3, and Table 8 was obtained at 7.5 × 10−4, the main recipe learning rate.Warm-started h = 36 and h = 54 runs diverged at the default rate (gradient norms up to 4.5 × 10^6) and needed a four-fold smaller one.Source: §B.4 · p. 22
B.5 Maze-unique
Maze-unique experiments test lattice and ACT models on mazes with a single shortest path. Lattice models retain high test accuracy across a wide parameter range, while ACT achieves near-perfect accuracy with deterministic rollouts and perfect accuracy under best-of-64 selection.
Data
- The 30 × 30 benchmark contains mazes with a unique shortest path of 100–140 cells, and every maze is checked for uniqueness.Training and test sets each contain 1,000 released mazes; validation uses 100 separately generated mazes.Source: §B.5 · p. 22
Augmentation
- The maze models are exactly equivariant to square symmetries, so geometric augmentation is omitted and only entrance–exit swaps are used.Swapping the endpoints preserves the solution because paths are undirected.Source: §B.5 · p. 22
Lattice results
- 98.9%–99.9% of test mazes are solved across lattice models from 128k down to 9k parameters, with shrinking models taking two to three times as many steps to reach perfect validation.The size study uses the 24 × 8, D = 20 recipe.Source: §B.5 · p. 23, Table 10 · p. 24
- Halving the credit horizon lowers accuracy at every size, while removing a third of the forward computation preserves accuracy when the horizon stays at four cycles.The authors report that differentiated-window length, rather than unroll depth or parameter budget, governs training quality.Source: §B.5 · p. 23, Table 11 · p. 24
ACT results
- 99.90% of test mazes are solved with deterministic rollouts at 256 deep-supervision steps, and 100.0% with best-of-64 selection at 64 steps.Pass@k reaches 100% by k = 32; the halting head averages 1.41 deep-supervision steps, or roughly 200 message-passing rounds.Source: §B.5 · p. 24
B.6 Practical notes
The practical notes report precision differences associated with fused Maze kernels and give representative training times for Sudoku, Maze, and Sudoku ACT runs.
Precision
- The fused reference arithmetic agrees to about 10^-15 in double and 10^-5 in single precision, while the Maze bf16 kernels change gradients by 3–10% in relative terms.Maze appendix results use bf16 kernels; Sudoku ACT uses single precision and Sudoku lattice uses bf16 autocast.Source: §B.6 · p. 24
Compute
- A 16k-parameter Sudoku lattice run takes about 15 minutes, a 128k Maze lattice run about half an hour, and a Sudoku ACT run eight to nine hours on an H100.The Sudoku ACT timing is for 180,000 steps with 8 × 12 deep-supervision steps.Source: §B.6 · p. 24
C N-Queens: experimental setup
The N-Queens setup addresses multiple valid solutions by constructing puzzles from complete boards, splitting data by symmetry orbit, and training against a selected consistent solution. It also evaluates transfer from 10 × 10 training to larger boards without fine-tuning.
Puzzles and correctness
- A prediction is correct when it forms a valid N-queen board and preserves the clues; each generated puzzle is solvable by construction.A 10 × 10 puzzle can have up to 93 consistent solutions.Source: C N-Queens: experimental setup · p. 25
Graphs
- The peer graph connects cells sharing a line, whereas the factor graph represents rows, columns, and diagonals as unit vertices and has four cell neighbors at every board size.The factor graph's parameter count does not depend on N; peer-graph cell degree grows with N.Source: C N-Queens: experimental setup · p. 25
Splitting by symmetry orbits
- The 10 × 10 split groups puzzles by dihedral orbit to prevent symmetry leakage, yielding 13,442 training and 2,372 test puzzles without a separate validation set.The split is shared across runs.Source: C N-Queens: experimental setup · p. 25
Size-transfer probes
- Size-transfer probes use test-only boards from N = 11 to 16, and models trained on 10 × 10 are evaluated there without fine-tuning.The factor graph permits this transfer because its parameters do not depend on board size.Source: C N-Queens: experimental setup · p. 25
Supervising many answers
- Training uses the consistent solution closest to the model's current prediction as its target, rather than sampling a solution randomly at each step.The paper reports that random targets drove deterministic models toward invalid per-cell averages, while all reported models use the minimum-over-solutions target.Source: C N-Queens: experimental setup · p. 25
Training protocol
- All models are ACT models capped at 16k parameters; peer-graph training uses full gradients, while factor-graph training uses truncated gradients.The RRN baseline uses the same factor graph, widths, data, target, and parameter budget, but trains with full backpropagation through time and no halting.Source: C N-Queens: experimental setup · p. 26
D Boolean satisfiability: experimental setup
The SAT experiments use two benchmark families that probe formula difficulty across sizes and near the satisfiability threshold, with correctness evaluated by exact clause checking. The shared factor-graph design supports variable-size formulas, while lattice and ACT recipes use distinct supervision and readout procedures.
Evaluation
- SAT formulas are conjunctions of clauses, and correctness is decided exactly by checking whether the returned assignment satisfies every clause.Satisfiable formulas have many valid assignments, so none is used as a target.Source: D Boolean satisfiability: experimental… · p. 26
Datasets
- RandCSPBench sweeps clause-to-variable ratios across five sizes, while G4SATBench targets threshold difficulty using estimated critical clause counts or solver-determined unsatisfiability.RandCSPBench includes 3-SAT and 4-SAT; G4SATBench uses easy, medium, and hard ranges, with hard test-only, and evaluation decodes test splits in full once.Source: D Boolean satisfiability: experimental… · p. 26
Graph
- Each formula is represented as a factor graph with variable memories and answer states, clause vertices using mean-pooled Deep Sets, and polarity-specific edge ports.The architecture is independent of variable count, clause count, and clause length; both recipes use widths (48, 64, 64, 16), with 64,689 lattice and 64,722 ACT parameters.Source: D Boolean satisfiability: experimental… · pp. 26–27
Lattice recipe
- The lattice recipe trains with 12 × 6 deep-supervision steps and a two-cycle credit horizon, using truth-value possibilities and a projection that commits assignments.Its loss combines a lattice-conditioned clause objective with cross-entropy tying readouts to commitments; steps begin blank with probability 0.20 or from random satisfiable commitments.Source: D Boolean satisfiability: experimental… · p. 27
ACT recipe
- The ACT recipe uses three supervision segments per step, each with 12 × 6 steps and a random differentiated window spanning 36 of 72 rounds.Its halting head pools over clauses with a correction for clause count and is trained against exact clause and formula satisfaction.Source: D Boolean satisfiability: experimental… · p. 27
Training protocol
- Both recipes use AdamW with a 400-step warm-up, cosine learning-rate decay, gradient clipping, EMA, batches of 128, and single precision.Lattice models train for 20,000 steps; ACT models train for 10,000 steps on easy splits and 20,000 otherwise.Source: D Boolean satisfiability: experimental… · p. 27