PaperlayerView exact revision on arXiv

Source-linked AI summary

Can You Learn an Algorithm? Generalizing from Easy to Hard Problems with Recurrent Networks

Avi Schwarzschild, Eitan Borgnia, Arjun Gupta, Furong Huang, Uzi Vishkin, Micah Goldblum, Tom Goldstein

arXiv:2106.04537v2cs.LGcs.AI

TL;DR

The paper asks whether recurrent networks can learn reasoning strategies that carry from easy training problems to harder test problems. It trains recurrent models on simpler instances and tests them with additional iterations, finding that more test-time recurrence enables extrapolation across prefix sums, mazes, and chess.

  • Problem

    The paper asks whether recurrent networks can learn scalable reasoning processes that generalize from easy training problems to harder test problems.

    Source: §1 · p. 2

  • Method

    Recurrent networks are trained on easy examples with a fixed iteration budget, then tested on harder examples with additional iterations.

    Source: §1 · p. 2

  • Results

    Across prefix sums, mazes, and chess, recurrent models trained on easy problems can generalize to harder ones by increasing test-time iterations, often outperforming comparable-depth feed-forward models.

    Source: §1 · p. 2, §2 · p. 3

  • Takeaways & Limitations

    The observations indicate that iterative models can generalize beyond their training distribution with more iterations, including when the test distribution is inaccessible.

    Source: §5 · p. 9

    The evidence is representative rather than definitive, and the chess difficulty split assumes higher-rated puzzles require more memory or computation.

    Source: §5 · p. 9

Abstract

from arXiv · show

Deep neural networks are powerful machines for visual pattern recognition, but reasoning tasks that are easy for humans may still be difficult for neural models. Humans possess the ability to extrapolate reasoning strategies learned on simple problems to solve harder examples, often by thinking for longer. For example, a person who has learned to solve small mazes can easily extend the very same search techniques to solve much larger mazes by spending more time. In computers, this behavior is often achieved through the use of algorithms, which scale to arbitrarily hard problem instances at the cost of more computation. In contrast, the sequential computing budget of feed-forward neural networks is limited by their depth, and networks trained on simple problems have no way of extending their reasoning to accommodate harder problems. In this work, we show that recurrent networks trained to solve simple problems with few recurrent steps can indeed solve much more complex problems simply by performing additional recurrences during inference. We demonstrate this algorithmic behavior of recurrent networks on prefix sum computation, mazes, and chess. In all three domains, networks trained on simple problem instances are able to extend their reasoning abilities at test time simply by "thinking for longer."

1 Introduction

The paper asks whether recurrent networks can learn scalable reasoning processes that extrapolate from easy training problems to harder test problems. It situates this question among work on adaptive computation and studies it through prefix sums, mazes, and chess.

1 Introduction

  • Recurrent networks can improve on harder problems by increasing iterations without adding parameters or retraining, unlike feed-forward networks whose layer-specific behaviors cannot be repeated.Source: §1 · p. 2
  • Because training and test problems differ in size or difficulty, this is extrapolation across disjoint distributions rather than conventional same-distribution generalization.Classical algorithms offer a functional analogue: they scale to arbitrary problem sizes at the cost of longer runtime.Source: §1 · p. 2
  • The goal is to train recurrent architectures to encode a scalable problem-solving method rather than memorize an input-output mapping.Source: §1 · p. 2
  • Recurrent models trained on easy problems can solve harder problems by taking more test-time iterations, and with enough iterations often outperform comparable-depth feed-forward models.The study also visualizes recurrent outputs to inspect the discovered problem-solving process.Source: §1 · p. 2

1.1 Related works

  • Prior adaptive-computation work largely studies inference on stationary distributions, whereas this work extends both compute budgets and problem difficulty beyond training.Related approaches include adaptive compute limits for sequence inputs and variable-depth networks.Source: §1.1 · p. 2
  • The case studies use prefix sums, maze solving, and chess to test whether learned processes generalize to harder problems, rather than to build the best solver for each task.The cited discussion distinguishes this aim from chess systems relying on hand-crafted search algorithms.Source: §1.1 · p. 3

2 Dataset descriptions

The experiments use three reasoning tasks with measurable difficulty, enabling training on easier instances and testing on harder ones. Their datasets represent prefix sums, maze paths, and chess-puzzle moves in task-specific input and target formats.

Prefix sums

  • Prefix-sum datasets contain 10,000 unique uniform random binary strings at each of 32, 44, and 48 bits, with cumulative-sum-modulo-two output strings.Figure 1 provides an input-and-target example.Source: §2 · p. 3, Figure 1 · p. 3

Mazes

  • Maze training uses 50,000 9 × 9 mazes and testing uses 10,000 13 × 13 mazes; each target is a binary mask of the shortest path, scored correct only on an exact match.Figure 2 shows a maze input and pixel labels indicating whether each pixel lies on the optimal path.Source: §2 · p. 3, Figure 2 · p. 3

Chess puzzles

  • Chess inputs encode piece locations in 8 × 8 × 12 arrays, while targets mark the optimal move’s origin and destination; training uses 600,000 puzzles rated below 1,385 and testing 100,000 rated above it.Figure 3 shows a puzzle input and a target in which white represents one and black zero.Source: Figure 3 · p. 3, §2 · p. 4

3 Model architectures & training

The models use ResNet-based architectures, with recurrent models sharing weights across residual blocks. Training minimizes mean cross-entropy over binary output labels, and model depth is measured from recurrent iterations plus encoder and head layers.

3 Model architectures & training

  • Recurrent and feed-forward models use ResNet-based architectures, but recurrent models share weights between residual blocks; the recurrent portion is called the recurrent module.The architectures are fully convolutional, with one-dimensional convolutions for prefix sums.Source: §3 · p. 4
  • Each recurrent iteration repeats a four-layer residual block, so effective depth equals four times the iteration count plus non-recurrent encoder and head layers.A 10-iteration prefix-sum model has effective depth 1 + 10 × 4 + 3 = 44 layers.Source: §3 · p. 4
  • Training uses mean cross-entropy across binary classification outputs, with hyperparameters selected to obtain convergent models.Source: §3 · p. 4

4 Recurrent networks can generalize from easy to hard problems

Across prefix sums, mazes, and chess, recurrent models generalize from easier training problems to harder test problems, and additional test-time iterations can improve performance. Their intermediate outputs show progress toward solutions across iterations.

4.1 Prefix sums

  • With more test-time iterations, recurrent prefix-sum models reach upwards of 90% accuracy on harder strings, while comparable and even much deeper feed-forward models fail to generalize similarly.Figure 4 compares accuracy on longer test strings for recurrent models and feed-forward models of comparable effective depth.Source: Figure 4 · p. 5, §4.1 · p. 5
  • Across prefix sums, added recurrence improves generalization beyond what receptive-field growth alone explains, because the compared feed-forward models have the same receptive field as recurrent models at 16 iterations.The authors attribute the improvement to a learned process that extrapolates beyond the training distribution.Source: §4.1 · pp. 5–6

4.1.1 Iterative outputs

  • Prefix-sum outputs progress toward the solution one iteration at a time, resolving earlier bits before later ones, like a sequential algorithm that scans the string.The illustrated model was trained with 10 iterations on 32-bit strings and shown on a 40-bit input.Source: §4.1.1 · p. 6, Figure 5 · p. 6

4.2 Mazes

  • Recurrent models generalize to large mazes better than feed-forward counterparts at fixed effective depths, and models trained with 20 iterations exceed 70% accuracy using five additional test-time iterations.The reported result is for models trained on small mazes and tested on large ones.Source: §4.2 · pp. 6–7, Table 2 · p. 7, Figure 6 · p. 7
  • Dilated filters slightly improve maze performance, but the benefits of recurrence remain clear and the performance difference is even larger.Dilation changes the receptive field without adding parameters or depth.Source: §4.2 · p. 7, Table 3 · p. 7

4.2.1 Iterative outputs

  • Maze outputs refine across iterations: the model finds two routes, sharpens confidence on the path, and matches the target at iteration seven, though some pixels flip at iteration nine.The figure illustrates a model trained to solve small mazes in six iterations.Source: §4.2.1 · p. 7, Figure 7 · p. 8

4.3 Chess puzzles

  • Recurrent models solve more chess puzzles than feed-forward counterparts, and their performance improves when they think for longer at test time.The gains are described as modest compared with those in the prefix-sum and maze settings.Source: Figure 8 · p. 8, §4.3 · p. 9

4.3.1 Iterative outputs

  • On chess puzzles, successive iterations increase confidence first in the piece to move and then in its destination, although the learned process is harder to interpret than those for sums and mazes.In the example, the model considers several rook moves and a queen check before narrowing its choice.Source: §4.3.1 · p. 9, Figure 9 · p. 9

5 Discussion

The paper asks whether learned models can generalize like classical algorithms, while emphasizing that its representative cases do not settle the question. It discusses possible uses of extra test-time computation, a working-memory analogy, and the challenge of deciding when to stop thinking.

5 Discussion

  • The broader question is whether learned models can generalize to harder or larger problems in the way classical algorithms do.The paper frames its results as raising this question rather than resolving it.Source: §5 · p. 9
  • The study offers representative evidence that recurrence can help models extrapolate from easy training data to harder test examples, but does not provide a definitive answer about algorithm-like generalization.For chess, the authors assume puzzle ratings track algorithmic complexity while noting they may be weak surrogates.Source: §5 · p. 9
  • When the harder test distribution is inaccessible, iterative models that generalize with more iterations may be useful, including in robot training from simulation.The authors also mention practicing easy math before attempting difficult unsolved questions.Source: §5 · p. 9
  • The recurrent process is conceptually analogous to manipulating representations in working memory, but the experiments do not claim that the models use human-like mechanisms.In the analogy, recurrent blocks perform transformations and their activations serve as memory.Source: §5 · p. 10
  • The authors ask whether networks can think longer and whether added computation can keep improving performance, potentially with an awareness of when to stop.They suggest stop-awareness as a possible future direction.Source: §5 · p. 10

6 Conclusion

The conclusion reports that recurrent models extrapolate from sequential reasoning tasks to more complex problems than those seen in training. It connects this behavior to algorithmic reasoning and leaves exploration of other network designs for future work.

6 Conclusion

  • Recurrent models solve sequential reasoning tasks and extrapolate to problems more complex than their training examples.The models are recurrent versions of popular architectures, inspired in part by classical theories of mind.Source: §6 · p. 10
  • An in-depth investigation of alternative neural network designs is left for future work.Source: §6 · p. 10
  • The results suggest that algorithmic behaviors might be learned end-to-end for tasks traditionally solved with hand-crafted algorithms.The paper names reduction trees for prefix sums, depth- or breadth-first search for mazes, and Monte Carlo tree search for chess; it presents gameplay without hand-crafted tree search as a possibility.Source: §6 · p. 10

A.1 Datasets

This section introduces the datasets and directs readers to their downloadable collection.

A.1 Datasets

A.1.1 Prefix sum data

Prefix-sum inputs are generated as fixed-size samples of binary strings, with dataset sparsity increasing as string length grows.

A.1.1 Prefix sum data

  • Each binary string of length n is generated from a random integer in [0, 2n), expressed with n digits.Source: §A.1.1 · p. 13
  • The dataset contains 10,000 strings sampled without replacement, making it sparser in the ambient hypercube as n increases.The number of possible points grows exponentially with n while dataset size remains fixed.Source: §A.1.1 · p. 13
  • The authors restrict string lengths to n > 13 to avoid duplicate data points.Source: §A.1.1 · p. 13

A.1.2 Maze data

The appendix describes maze and chess data construction and the corresponding model architectures. Recurrent variants share weights across residual blocks, unlike their feed-forward counterparts.

A.1.2 Maze data

  • Depth-first search generates maze grids by carving paths that visit every cell; path lengths vary, and fewer than 0.5% of samples are duplicates.Maze targets label each pixel as on or off the optimal path.Source: §A.1.2 · p. 13, Figure 10 · p. 13

A.1.3 Chess puzzle data

  • Chess inputs encode board positions as 8 × 8 × 12 arrays, while targets mark the origin and destination of the first correct move.The puzzles derive from analysis of 200 million games using Stockfish 12/13 NNUE at 40 meganodes; the data is released under CC0.Source: §A.1.3 · p. 13

A.2 Architectural details

  • Prefix-sum models use 1D convolutions and maze models use 2D convolutions; recurrent versions share weights across residual blocks.The feed-forward and recurrent maze models otherwise share the described architecture; chess models match maze models except for a 512-channel first layer taking 8 × 8 × 12 inputs.Source: §A.2 · p. 14
  • Dilated-filter experiments change convolution dilation and padding to preserve each layer’s output dimensions.The models do not use batch normalization or bias terms.Source: §A.2 · p. 14

A.3 Training hyperparameters

The appendix reports training settings and compute use, alongside a comparison showing that substantially deeper feed-forward models still generalize poorly to longer prefix-sum inputs.

A.3 Training hyperparameters

  • Prefix-sum models use Adam for 500 epochs, while maze models use SGD for 200 epochs and chess models use SGD for 140 epochs.All use weight decay of 2e-4; prefix-sum training also uses gradient clipping at magnitude 1.0 because of instability.Source: §A.3 · p. 14

A.4 Compute resources

  • The reported total compute is approximately 30 GPU hours for prefix sums, one GPU week for mazes, and three GPU weeks for chess.Experiments used Nvidia GeForce RTX 2080Ti GPUs.Source: §A.4 · pp. 14–15

B.1 Deeper feed-forward models

  • 53% on 40-bit and 27% on 44-bit data were the best accuracies among feed-forward models of depths 132, 264, and 528.All fit the 32-bit training data well, and the deepest model did not attain either best result.Source: §B.1 · p. 15

B.2 In-distribution tests

The section reports in-distribution performance, transfer to smaller inputs, and findings on dilated filters; figures also visualize longer-input prefix sums and harder chess puzzles.

B.2 In-distribution tests

  • Both model classes achieve >99% on prefix sums and >97% on mazes in easy in-distribution tests; on chess, recurrent models reach 92% versus 84% for feed-forward models.On harder datasets, both classes achieve >99% on prefix sums and >97% on mazes; additional iterations slightly reduce easy-case accuracy without improving it.Source: §B.2 · p. 15
  • 100% of 16-bit test examples are correct after 4 iterations when models trained on 44-bit prefix sums with 10 iterations are tested on the smaller inputs.The section reports that models trained on larger instances can solve smaller cases in fewer iterations than used during training.Source: §B.2 · p. 15

B.4 Dilated filters

  • Dilated filters let prefix-sum models fit N-bit training data with fewer than N layers, unlike non-dilated convolutions, suggesting the learned algorithm depends on receptive field.Dilation can increase receptive field without adding parameters or depth.Source: §B.4 · p. 15

C Visualizations

  • Figure 11 compares recurrent models with feed-forward models of comparable effective depth on 36- and 48-bit inputs; Figure 12 shows chess-puzzle performance across test-index ranges.The chess models were trained on the first 600,000 puzzles; the figure includes later ranges and indices 100,000 to 150,000.Source: Figure 11 · p. 16, Figure 12 · p. 16

C.1 Prefix sums

The visualizations show intermediate model outputs for prefix sums, mazes, and chess puzzles. The checklist records the authors’ stated scope and reporting choices.

C.1 Prefix sums

  • The prefix-sum visualization shows outputs across 11 iterations on 40-bit strings from a model trained for 10 iterations on 32-bit strings.It plots confidence at each output index, with input and target vectors at the left and right.Source: §C.1 · p. 16, Figure 13 · p. 17

C.2 Mazes

  • Maze visualizations show inputs, targets, and outputs across iterations for a model trained to solve small mazes in six iterations.The plotted confidence indicates whether each pixel belongs to the optimal path.Source: §C.2 · p. 17, Figure 14 · p. 18

C.3 Chess puzzles

  • Chess visualizations plot the model’s confidence that each pixel marks one of the two positions defining a move across different iterations.The figures show intermediate outputs for the chess model.Source: §C.3 · p. 17, Figure 15 · p. 19

D NeurIPS submission checklist

  • The checklist says limitations are discussed and societal impacts are not, because the investigation concerns toy AI domains.It also marks code, data, instructions, training details, compute resources, and asset licensing as reported.Source: D NeurIPS submission checklist · p. 20
Loading 2106.04537v2…