PonderNet: Learning to Ponder
Paper · 2021 · Andrea Banino, Jan Balaguer and Charles Blundell
Summary
Standard neural networks spend computation in proportion to the size of the input, but not to the complexity of the task itself. The authors from DeepMind propose PonderNet, an algorithm that decides on its own how many "pondering" steps a particular example requires. Any network (MLP, LSTM, transformer) is applied recurrently as a step function. At each step it produces a prediction and a probability of halting. From these probabilities, the distribution of the halting step is derived exactly; it generalizes the geometric distribution. The loss function combines two terms: the expected prediction error over all steps and the KL divergence from a geometric prior distribution. The second term does not directly penalize the number of steps, but rather encourages exploration of different options. The method is fully differentiable. Unlike ACT, it yields unbiased gradients, and unlike REINFORCE, it yields low-variance gradients. On the synthetic parity task, PonderNet outperforms ACT and passes extrapolation tests. On bAbI it reaches state-of-the-art performance with less computation, and on PAI it matches the specialized MEMO model. The value of the work lies in providing a simple and stable way to adapt the amount of computation: it suffices to add a halting node and a single term to the loss function.
Key points
- Parity, extrapolation. The model was trained on vectors with 1–48 non-zero elements and tested on 49–96. PonderNet achieved near-perfect accuracy, while ACT remained at chance level. The number of pondering steps increased on its own to about 5, i.e., almost doubled.
- Robustness to the hyperparameter λp. The task went unsolved only with λp = 0.9 (a prior of about 1 step). With λp = 0.1 (a prior of 10 steps), the network converged on its own to about 3 steps. ACT, by contrast, is highly sensitive to τ and shows large variance across runs.
- bAbI (joint 10k, 20 tasks). Transformer+PonderNet has a mean error of 0.15 and solves all 20 tasks. Universal Transformer has an error of 0.29, and in total it uses 10,161 steps versus 1658 for PonderNet.
- PAI (training on A-B-C, testing the A-C association). PonderNet's accuracy is 97.86, MEMO's is 98.26, and Universal Transformer's is 85.60.
- Technical details. During training, the maximum number of steps N is determined by a threshold: the cumulative halting probability must exceed 1 − ε, where ε = 0.05. During inference, the halting decision at each step is sampled from a Bernoulli distribution.
Fields, topics and concepts
📚 Shelves 3The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksComputer Vision and Pattern Recognition · 3,621 works
Cite
Cite as Andrea Banino, Jan Balaguer and Charles Blundell. 2021. PonderNet: Learning to Ponder. arXiv:2107.05407.
@misc{banino2021pondernet,
title = {{PonderNet: Learning to Ponder}},
author = {Andrea Banino, Jan Balaguer and Charles Blundell},
year = {2021},
doi = {10.48550/arxiv.2107.05407},
eprint = {2107.05407},
archivePrefix = {arXiv},
url = {https://doi.org/10.48550/arxiv.2107.05407}
}References
- Andrea Banino, Adrià Puigdomènech Badia, Raphael Köster, Martin J. Chadwick, Vinicius Zambaldi, Demis Hassabis, Caswell Barry, Matthew Botvinick, Dharshan Kumaran, and Charles Blundell. MEMO: A deep network for flexible combination of episodic memories. In International Conference on Learning Rep… find
- Tolga Bolukbasi, Joseph Wang, Ofer Dekel, and Venkatesh Saligrama. Adaptive neural networks for efficient inference. In Proceedings of the 34th International Conference on Machine Learning-Volume 70, pages 527–536. JMLR. org, 2017. find
- Victor Campos Camunez, Brendan Jou, Xavier Giró Nieto, Jordi Torres Viñals, and Shih-Fu Chang. Skip RNN: learning to skip state updates in recurrent neural networks. In Sixth International Conference on Learning Representations: Monday April 30-Thursday May 03, 2018, Vancouver Convention Center, V… find
- Junyoung Chung, Sungjin Ahn, and Yoshua Bengio. Hierarchical multiscale recurrent neural networks. arXiv preprint arXiv:1609.01704, 2016.
- Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint arXiv:1901.02860, 2019. find
- Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprint arXiv:1807.03819, 2018. find
- UT MEMO PonderNet 85.60 98.26(0.67) 97.86(3.78) J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei. ImageNet: A Large-Scale Hierarchical Image Database. In CVPR09, 2009. find
- Alex Graves. Adaptive computation time for recurrent neural networks. arXiv preprint arXiv:1603.08983, 2016. find
- Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016.
- Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
- Sainbayar Sukhbaatar, Jason Weston, Rob Fergus, et al. End-to-end memory networks. In Advances in Neural Information Processing Systems, pages 2440–2448, 2015. find
- Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998–6008, 2017.
- Petar Veličković, Rex Ying, Matilde Padovano, Raia Hadsell, and Charles Blundell. Neural execution of graph algorithms. arXiv preprint arXiv:1910.10593, 2019. find
- Jason Weston, Antoine Bordes, Sumit Chopra, Alexander M Rush, Bart van Merrinboer, Armand Joulin, and Tomas Mikolov. Towards AI-complete question answering: A set of prerequisite toy tasks. arXiv preprint arXiv:1502.05698, 2015.
- Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):229–256, 1992. find
- Adams Wei Yu, Hongrae Lee, and Quoc V Le. Learning to skim text. arXiv preprint arXiv:1704.06877, 2017. find
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/banino-2021-pondernet-learning-ponder/