# Learning to Stop While Learning to Predict

arXiv: https://arxiv.org/abs/2006.05082 · HTML: https://mlsanity.com/paper/2006.05082/

## Key facts

- Cited by 31 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 9 Jun 2020.

## Built on by (most cited first, 30 of 31)

- [Learning to Optimize: A Primer and A Benchmark](https://mlsanity.com/paper/2103.12828/) (2021)
- [Adaptive Inference through Early-Exit Networks: Design, Challenges and Directions](https://mlsanity.com/paper/2106.05022/) (2021)
- [Task-Oriented Communication for Multi-Device Cooperative Edge Inference](https://mlsanity.com/paper/2109.00172/) (2021)
- [Learning Iterative Reasoning through Energy Diffusion](https://mlsanity.com/paper/2406.11179/) (2024)
- [Learning Iterative Reasoning through Energy Minimization](https://mlsanity.com/paper/2206.15448/) (2022)
- [Energy-Based Transformers are Scalable Learners and Thinkers](https://mlsanity.com/paper/2507.02092/) (2025)
- [Fast yet Safe: Early-Exiting with Risk Control](https://mlsanity.com/paper/2405.20915/) (2024)
- [Realistic Evaluation of Transductive Few-Shot Learning](https://mlsanity.com/paper/2204.11181/) (2022)
- [Few-Shot Learning with Class Imbalance](https://mlsanity.com/paper/2101.02523/) (2021)
- [Window-Based Early-Exit Cascades for Uncertainty Estimation: When Deep Ensembles are More Efficient than Single Models](https://mlsanity.com/paper/2303.08010/) (2023)
- [Jointly-Learned Exit and Inference for a Dynamic Neural Network : JEI-DNN](https://mlsanity.com/paper/2310.09163/) (2023)
- [BAGEN: Are LLM Agents Budget-Aware?](https://mlsanity.com/paper/2606.00198/) (2026)
- [Continuous-Time Meta-Learning with Forward Mode Differentiation](https://mlsanity.com/paper/2203.01443/) (2022)
- [OccamNets: Mitigating Dataset Bias by Favoring Simpler Hypotheses](https://mlsanity.com/paper/2204.02426/) (2022)
- [Artificial Intelligence for Operations Research: Revolutionizing the Operations Research Process](https://mlsanity.com/paper/2401.03244/) (2024)
- [Revisiting Single-gated Mixtures of Experts](https://mlsanity.com/paper/2304.05497/) (2023)
- [Adaptive Deep Neural Network Inference Optimization with EENet](https://mlsanity.com/paper/2301.07099/) (2023)
- [Improving Anytime Prediction with Parallel Cascaded Networks and a Temporal-Difference Loss](https://mlsanity.com/paper/2102.09808/) (2021)
- [Tiny Models are the Computational Saver for Large Models](https://mlsanity.com/paper/2403.17726/) (2024)
- [Learning to Accelerate Approximate Methods for Solving Integer Programming via Early Fixing](https://mlsanity.com/paper/2207.02087/) (2022)
- [ATHEENA: A Toolflow for Hardware Early-Exit Network Automation](https://mlsanity.com/paper/2304.08400/) (2023)
- [FIANCEE: Faster Inference of Adversarial Networks via Conditional Early Exits](https://mlsanity.com/paper/2304.10306/) (2023)
- [DyCE: Dynamically Configurable Exiting for Deep Learning Compression and Real-time Scaling](https://mlsanity.com/paper/2403.01695/) (2024)
- [ModelPred: A Framework for Predicting Trained Model from Training Data](https://mlsanity.com/paper/2111.12545/) (2021)
- [EERO: Early Exit with Reject Option for Efficient Classification with limited budget](https://mlsanity.com/paper/2402.03779/) (2024)
- [How Sensitive are Meta-Learners to Dataset Imbalance?](https://mlsanity.com/paper/2104.05344/) (2021)
- [Non-Myopic Active Feature Acquisition via Pathwise Policy Gradients](https://mlsanity.com/paper/2605.05511/) (2026)
- [ERDE: Entropy-Regularized Distillation for Early-exit](https://mlsanity.com/paper/2510.04856/) (2025)
- [Scaling Competence, Shrinking Reasoning: Cognitive Signatures in Language Model Learning](https://mlsanity.com/paper/2511.21743/) (2025)
- [Supraglacial Lake Fate Is Knowable Long Before the Season Ends](https://mlsanity.com/paper/2608.30113/) (2026)

## Abstract (arXiv)

There is a recent surge of interest in designing deep architectures based on the update steps in traditional algorithms, or learning neural networks to improve and replace traditional algorithms. While traditional algorithms have certain stopping criteria for outputting results at different iterations, many algorithm-inspired deep models are restricted to a ``fixed-depth'' for all inputs. Similar to algorithms, the optimal depth of a deep architecture may be different for different input instances, either to avoid ``over-thinking'', or because we want to compute less for operations converged already. In this paper, we tackle this varying depth problem using a steerable architecture, where a feed-forward deep model and a variational stopping policy are learned together to sequentially determine the optimal number of layers for each input instance. Training such architecture is very challenging. We provide a variational Bayes perspective and design a novel and effective training procedure which decomposes the task into an oracle model learning stage and an imitation stage. Experimentally, we show that the learned deep model along with the stopping policy improves the performances on a diverse set of tasks, including learning sparse recovery, few-shot meta learning, and computer vision tasks.