PaperlayerView exact revision on arXiv

Source-linked AI summary

PonderNet: Learning to Ponder

Andrea Banino, Jan Balaguer, Charles Blundell

arXiv:2107.05407v2cs.LGcs.AIcs.CC

TL;DR

Standard neural networks generally scale computation with input size rather than problem complexity, motivating adaptive computation. PonderNet learns probabilistic halting and a geometric-prior training objective, achieving strong results across parity, question answering, and multi-step reasoning while using less compute in a bAbI comparison.

  • Problem

    Most neural networks do not adapt their computational budget to the inherent complexity of the task, limiting computation-aware learning.

  • Method

    PonderNet learns a probabilistic halting policy with a differentiable objective combining prediction loss and a geometric-prior regularizer.

  • Results

    PonderNet achieved higher accuracy than ACT with less computation on parity, nearly perfect extrapolation accuracy where ACT remained at chance, and state-of-the-art-level results on question answering and multi-step reasoning.

  • Takeaways & Limitations

    PonderNet can increase computation for harder inputs, extrapolate beyond training lengths, and match or exceed benchmark results using less compute in the reported comparisons.

Abstract

from arXiv · show

In standard neural networks the amount of computation used grows with the size of the inputs, but not with the complexity of the problem being learnt. To overcome this limitation we introduce PonderNet, a new algorithm that learns to adapt the amount of computation based on the complexity of the problem at hand. PonderNet learns end-to-end the number of computational steps to achieve an effective compromise between training prediction accuracy, computational cost and generalization. On a complex synthetic problem, PonderNet dramatically improves performance over previous adaptive computation methods and additionally succeeds at extrapolation tests where traditional neural networks fail. Also, our method matched the current state of the art results on a real world question and answering dataset, but using less compute. Finally, PonderNet reached state of the art results on a complex task designed to test the reasoning capabilities of neural networks.1

1. Introduction

PonderNet adapts computation to task complexity by learning a probabilistic halting policy. It addresses limitations of prior adaptive-computation methods through differentiable, unbiased training and distinct architectural, loss, and inference choices.

  • Motivation: Problem complexity can vary independently of input size, but most machine-learning algorithms do not adapt their computational budget accordingly.This adaptation is called pondering.
  • Limitations of prior work: Adaptive Computation Time learns computation through a scalar halting probability, but its training is unstable and sensitive to the accuracy–computation trade-off hyper-parameter.Its computation-cost gradient also back-propagates only through the last computational step, producing a biased gradient estimate.
  • PonderNet: PonderNet reformulates halting as a probabilistic model, making the method fully differentiable with low-variance and unbiased gradient estimates.The reformulation affects the architecture, loss, and inference procedure.
  • Architecture: The halting node predicts the probability of stopping conditional on not having halted earlier, yielding an overall geometric halting distribution.This explicitly models the probability of halting at each computational step.
  • Loss and inference: PonderNet encourages exploration rather than directly regularizing the number of computation steps, allowing efficient computation to emerge through an Occam’s-razor effect.Inference is probabilistic in both the number of computational steps and the prediction produced.

2. Methods

PonderNet augments recurrent computation with a learned halting process and trains it using a prediction loss plus a geometric-prior regularizer. Its probabilistic halting distribution determines both dynamic computation depth and the final prediction.

  • Architecture: PonderNet modifies a neural network’s forward pass and adds a loss function to learn a function from supervised input–output data.The architecture can use recurrent step functions such as MLPs, LSTMs, or transformers.
  • Step recurrence: At each step, the step function produces a prediction, the next hidden state, and a scalar halting probability, and it is recurrently applied up to N times.An encoder can alternatively initialize the hidden state from the input before recurrence.
  • Halting process: PonderNet represents halting as a two-state Markov process whose transition probability is stopping at step n given that no earlier step halted.The process starts in the continue state, and halting is terminal.
  • Halting distribution: The unconditioned halting probabilities form a generalized geometric distribution, and the prediction is sampled from the output produced at the step where halting occurs.This differs from ACT’s weighted averaging across computational steps.
  • Finite computation: Finite unrolling requires normalizing the halting probabilities or assigning remaining probability to the final step.During training, N can be set as the smallest step whose cumulative halting probability exceeds 1 − ε, with ε set near 0.
  • Training loss: The total loss combines expected prediction loss with a KL-divergence regularizer against a truncated geometric prior over halting.The prior biases the expected number of steps toward 1/λp and gives non-zero probability to possible step counts, promoting exploration.
  • Inference: At evaluation, the network samples whether to continue or halt at each step and returns the prediction from the halting step, or from step N when the maximum is reached.This makes both computation length and prediction probabilistic.

3. Results

Across parity, extrapolation, bAbI, and paired associative inference, PonderNet achieved strong accuracy while adapting computation to task demands and using less compute in reported comparisons.

  • Parity: PonderNet achieved better parity accuracy than ACT with more efficient thinking time and less total training computation.Figure 1 compares accuracy, ponder steps, and total forward passes across PonderNet, ACT, and a non-adaptive RNN.
  • Parity: PonderNet remained effective across prior settings, solving parity with λp = 0.1 and settling to roughly 3 average thinking steps.The method failed only when λp = 0.9, corresponding to a prior average of roughly 1 thinking step.
  • Parity: Almost perfect extrapolation accuracy was achieved by PonderNet, while ACT remained at chance level.PonderNet increased its thinking time to 5 steps on inputs beyond the training range.
  • bAbI: On bAbI, PonderNet matched state-of-the-art accuracy while using 1658 steps versus 10161 for Universal Transformer.The reported comparison also states that PonderNet achieved the results faster and with lower average error.
  • Paired associative inference: On paired associative inference, PonderNet matched MEMO and achieved higher accuracy than Universal Transformer despite sharing its architecture.The results are reported as averaged inference-trial accuracy in Table 2.

4. Discussion

PonderNet reformulates adaptive halting probabilistically, changing the architecture, loss, inference procedure, and scope of computation. The paper reports extrapolation on parity and high accuracy in question answering and multi-step reasoning.

  • 4. Discussion: PonderNet increased computation to extrapolate beyond training data on parity and achieved the highest accuracy in complex question-answering and multi-step reasoning domains.The method can be added to recurrent architectures by augmenting the step function with a halting unit and adding a loss term.
  • 4. Discussion: PonderNet reformulates how neural networks learn to halt by using a probabilistic halting policy.This probabilistic formulation affects the architecture, forward computation, training loss, and deployment.
  • 4. Discussion: PonderNet and ACT share the same step function, but use the halting node differently during forward computation.ACT treats λ_n as the overall halting probability, whereas PonderNet distinguishes training unroll length from the learned halting policy.
  • 4. Discussion: During evaluation, PonderNet samples a halt state and returns the current prediction, unlike ACT’s weighted-average prediction procedure.PonderNet computes weighted losses across steps during training, but samples halting during evaluation.
  • 4. Discussion: PonderNet proposes holistic pondering for inputs such as graph message-passing steps rather than independent pondering for each sequence item.The proposal is motivated by attention mechanisms that process arrays with dynamic shapes.

A.2 Training loss

PonderNet replaces ACT’s proxy-based computation penalty with an exactly differentiable expected-step objective and a prior-matching regularizer that encourages exploration and efficient computation.

  • A.2 Training loss: PonderNet computes and differentiates the expected number of computational steps exactly, unlike ACT’s proxy based on the remainder.ACT uses the remainder because its maximum-step criterion is not differentiable with respect to the halting probabilities.
  • A.2 Training loss: Matching a prior halting distribution encourages alternative halting strategies, robustness, and avoiding computation beyond the prior.The KL regularizer is also expressed in the same units as information-theoretic losses such as cross-entropy.
  • A.2 Training loss: PonderNet computes loss for every possible computational step and minimizes the weighted expectation, whereas ACT averages predictions before computing loss.This difference makes combining multiple networks easier in ACT than in PonderNet.
  • A.2 Training loss: PonderNet uses λ_p and β as loss hyper-parameters, while ACT uses τ to trade off accuracy and computational complexity.Although τ and β both weight regularization terms, their effects differ because the regularizers have different interpretations.
  • A.2 Training loss: PonderNet trains and evaluates differently: training aggregates losses across steps, while evaluation samples halting and may perform another computational step.ACT uses the same weighted-average prediction procedure during training and evaluation.

B.1 Training and evaluation details

The parity experiment compares common recurrent architectures and optimization settings while varying ACT’s computation penalty and fixing key PonderNet settings.

  • B.1 Training and evaluation details: The experiment used the parity task described by Graves (2016).The task evaluates whether models can predict whether the input contains an odd number of ones.
  • B.1 Training and evaluation details: All models used a one-hidden-layer RNN with 128 tanh units and a logistic sigmoid output, optimized with Adam at learning rate 0.0003.Training used binary cross-entropy loss.
  • B.1 Training and evaluation details: ACT was evaluated across 20 τ settings, spanning [2e-4, 2e-2] plus 0, while both ACT and PonderNet used N = 20.PonderNet used β = 0.01.

C.1 Training and evaluation details

The question-answering experiment used preprocessed English Question Answer data with padded query-story batches and evaluated mean accuracy on sampled examples, including comparisons with published baselines.

  • C.1 Training and evaluation details: The experiment used the English Question Answer dataset with standardized preprocessing of text, punctuation, spacing, answers, and query separation.Queries were stripped from the text and supplied separately to the system.
  • C.1 Training and evaluation details: Training sampled mini-batches of 128 queries with corresponding stories, represented as padded arrays of up to 320 stories and 11 tokens per sentence.Queries had shape 128 × 11 tokens, while stories had shape 128 × 320 × 11.
  • C.1 Training and evaluation details: Models were trained for 2e4 epochs with 100 batch updates per epoch and evaluated on 10,000 sampled examples using mean accuracy.The evaluation forward pass followed the training procedure.
  • C.1 Training and evaluation details: Published MEMO results came from Banino et al. (2020), while Universal Transformer results came from Dehghani et al. (2018).

C.2 Transformer architecture and hyperparameters

The experiments use a Universal Transformer small implementation for bAbI, with hyperparameters and search ranges documented in Tables 3 and 4.

  • C.2 Transformer architecture and hyperparameters: The bAbI experiments use the Universal Transformer small implementation and its described hyperparameters.The implementation is attributed to Dehghani et al. (2018).
  • C.2 Transformer architecture and hyperparameters: Table 3 reports the hyperparameters used for bAbI experiments.
  • C.2 Transformer architecture and hyperparameters: Table 4 reports the hyperparameter ranges searched with PonderNet on bAbI.

D.1 PAI - Task details

The PAI task uses ImageNet-derived image embeddings arranged into paired associations, then tests direct and indirect cue–match reasoning against lures.

  • D.1 PAI - Task details: The PAI dataset uses ImageNet images embedded by a pretrained ResNet, with length-three A-B-C sequences and 1e6 training images, 1e5 evaluation images, and 2e5 testing images.Each batch entry selects N = 16 sequences.
  • D.1 PAI - Task details: Each memory stores pairwise associations such as A_iB_i and B_iC_i, producing M = 32 rows with two 1000-dimensional embeddings per row.
  • D.1 PAI - Task details: Queries contain a cue, a match, and a lure; direct queries use the same memory slot, whereas indirect queries require linking facts across different memory locations.An indirect example is A1 as cue, C1 as match, and C6 as lure.
  • D.1 PAI - Task details: The network receives concatenated cue, match, and lure embeddings as a 3 × 1000-dimensional vector, with randomized match and lure positions.Batches are balanced between direct and indirect queries, and targets are ImageNet class IDs of the matches.
  • D.1 PAI - Task details: The memory, query, and target sizes are 32 ∗2 ∗1000, 1 ∗3 ∗1000, and 1, respectively.

D.2 PAI - Architecture details

The PAI architecture combines a Universal Transformer-like encoder–decoder with memory, while learning a shared bounded pondering time through PonderNet.

  • D.2 PAI - Architecture details: The encoder and decoder layer count is learned jointly, constrained to be equal, and bounded by a maximum number of layers N.This learned layer count is the architecture’s “pondering time.”
  • D.2 PAI - Architecture details: At each step, PonderNet updates its state with one encoder application and applies the decoder an equal number of times to produce the prediction.
  • D.2 PAI - Architecture details: The PAI experiments optimize weights with Adam, using a maximum learning rate of 0.0003, linear warm-up during the first epoch, and mini-batches of 128.
  • D.2 PAI - Architecture details: Table 5 lists the hyperparameters used for PAI experiments.
  • D.2 PAI - Architecture details: Table 6 lists the hyperparameter ranges searched with PonderNet on PAI.

D.3 PAI - Results based on query type

The PAI results section reports evaluation details and compares PonderNet with MEMO and Universal Transformer on the length-three A-B-C task.

  • D.3 PAI - Results based on query type: The reported PAI results come from evaluation sets containing 600 items per set.
  • D.3 PAI - Results based on query type: Table 7 reports results for the paired associative task with sequence length three, A-B-C.
  • D.3 PAI - Results based on query type: The MEMO and Universal Transformer comparison results were taken from Banino et al. (2020).
  • D.3 PAI - Results based on query type: PonderNet is introduced as a method that adapts neural-network computational complexity to the task being solved.
Loading 2107.05407v2…