Tiny Recursive Message Passing
Tommaso Salvatori ⋅ Giovanni de Felice ⋅ Alexis Toumi
Abstract
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 trends have, however, shown that better results can be achieved using recursive transformers equipped with inference procedures such as deep supervision, truncated recursion, and deduction lattices, which make extremely deep loops stable to train. Such models reach state of the art performance while remaining extremely parameter efficient. In this work, we show how message passing models can be used with the training and inference machinery of modern recursive models: deep supervision, a discrete deduction lattice, and, optionally, truncated recursion. The result is a recursive message passing architecture that reaches above $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 obtain similar efficiency gains on Shakashaka, the PPBench puzzle where LLMs struggle the most, by exploiting the planar structure of the problem.
Chat is not available.
Successful Page Load