ReWork: Efficient Iterative Reasoning with Working Memory
Abstract
Hard reasoning problems yield to iteration: a model drafts a candidate solution, then repeatedly revises it. Two group of revision methods coexist in the literature: latent recursion deepens the computation in the latent space, and diffusion revises the predictions over denoising steps. Iteration of either kind places two competing demands on model state, which must maintain the partial solution built so far yet and keep the plasticity to change wherever that solution is wrong. A recurrent network that carries a single hidden state forward asks one representation to serve both demands: persistence protects progress but establishes stale computation, while plasticity enables correction but erases what is already built. To mitigate this, we introduce ReWork with refreshing working memory for latent recurrence. To split the two demands into two states with different lifetimes a token-aligned sensory memory persists across latent recursions and holds the evolving answer, while compact working memory reads that answer, integrates global evidence, writes an update back, and is refreshed for the next recurrence. By nesting the latent recurrence and discrete diffusion, ReWork re-estimates every token across denoising steps, and latent recursion refines the sensory memory within each latent recurrence. Experiments on Maze-Hard show that a 36.8M-parameter ReWork solves 82.0\% of mazes exactly compared to size-matched token-aligned recurrent model (72.8\%), while costing under 48\% of the attention a token-aligned update would use. On Sudoku-Extreme the same separation nearly triples exact solve (60.4\% against 21.1\%). Exact accuracy grows along both loops and equal-compute allocations are not interchangeable, so recursion and denoising act as complementary axes of test-time compute, and a learned halting rule removes about half of the recurrent computation at unchanged accuracy. An iterative reasoner, these results suggest, should couple a persistent sensory memory with disposable working memory.