When Can Neural Networks Think Longer to Solve Harder Problems?

Published at AAAI 2026 (oral presentation) · Paper · Code · Talk video · Slides · Poster · maze-dataset package

Joint work with Amandin Chyba Rabeendran, Michael Ivanitskiy, Jordan Pettyjohn, Cecilia Diniz-Behn, Samy Wu Fung, and Daniel McKenzie.

TL;DR: Some neural networks can “think longer” at test time by iterating more. I stress-tested whether that extra compute lets them solve problems harder than anything they saw in training. It does in some ways but not others: the models quietly learned a shortcut heuristic instead of a general algorithm, and diversifying the training data fixed some failures while creating new ones.

The problem

Production models constantly face inputs harder than their training data. Recurrent networks (RNNs) and implicit networks (INNs) offer a tempting fix: their depth is adjustable at test time, so they can spend more compute on harder inputs — like a person thinking longer about a harder puzzle. I investigated when this actually works and how it fails.

Maze-solving is the ideal testbed: classical algorithms provide ground-truth solutions at any difficulty, so I trained on easy mazes and generated harder test mazes along three controlled dimensions. Maze size (e.g., 9×9) makes problems bigger. Percolation (p = 0 to 1) removes walls to create loops, producing multiple solutions where previously there was one. Deadend-start (True or False) controls whether the solution path begins at a dead end. The models were trained only on the easiest corner of this space — 9×9 mazes with p = 0 and deadend-start True — then tested far beyond it. All mazes come from maze-dataset, the open-source package we built for exactly this purpose.

What I did

• Evaluated a pre-trained RNN from Bansal et al. and a custom-trained INN across thousands of out-of-distribution mazes, varying maze size, percolation, and deadend-start.
• Investigated what the models actually learned by studying their failures: correct answers all look alike, but failures have signatures. This is how we found that the RNN had approximately learned dead-end filling. The INN matched no algorithm we tested.
• Used topological data analysis to rigorously measure whether the networks’ internal computations converge to a fixed point, or something else.
• Retrained both architectures across a range of training-data diversity levels, against a standard network 10× the size, to assess how diversification affects generalization.

What I found

1. Test-time compute works — dramatically. With enough iterations, the RNN solves mazes 10× larger than anything in training at near-perfect accuracy, and both iterative models beat a standard network with 10× the parameters.

2. But the model learned a heuristic, not the goal. The RNN’s predictions approximately match dead-end filling, a classic algorithm that provably fails on mazes with loops because it incorrectly retains loops from the input maze in its prediction. Perfect training accuracy concealed that the model never learned a general maze-solving method, only a shortcut.

3. Data diversity is a trade-off, not a free lunch. I diversified the training data by raising its percolation above p = 0, letting the models see mazes with loops during training. Even p = 0.003 unlocked loop-solving — but as training percolation grew, accuracy concentrated around the training distribution and size extrapolation degraded.

4. A standing theoretical assumption was contradicted. Prior work assumed iterates must converge to a fixed point for extrapolation. Topological analysis revealed the RNN often converges to cycles instead while still achieving perfect accuracy — behavior that standard residual plots hide, but PCA projections expose.

Why it matters beyond mazes

Mazes are a toy problem, and that is exactly the point. They are a miniature testbed for some of the biggest concepts in modern AI.

Test-time scaling: RNNs and INNs really can think longer to solve harder problems. The same principle, realized in a very different architecture, drives today’s reasoning models.

Goal misgeneralization: a model can ace training while quietly learning the wrong goal. Ours learned a shortcut heuristic, and no accuracy metric revealed the problem until the distribution shifted. Diagnosis is hard even in a toy setting: we could match the RNN’s behavior to dead-end filling’s, but its mechanism stayed opaque, since two algorithms can agree on every output while computing differently inside. Some researchers fear that someday a superintelligent model subtly pursuing an unintended goal will be catastrophic.

Data diversification: broadening the training data improved generalization in one direction while degrading it in another. More diverse data is not automatically better.

From Colorado to Singapore

I was lucky enough to have this work selected for an oral presentation at AAAI-26 (top 4% of 23,680 submissions). This meant presenting to an international audience of experts on the other side of the world in Singapore. It was an unforgettable way to cap a years-long project.

Skills: Python/PyTorch, large-scale experiment design, out-of-distribution evaluation & stress-testing, model debugging & failure analysis, topological data analysis, scientific communication (peer-reviewed AAAI paper + conference talk).