PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning to Stop While Learning to Predict

Xinshi Chen, Hanjun Dai, Yu Li, Xin Gao, Le Song

arXiv:2006.05082v1cs.LGstat.ML

TL;DR

Algorithm-inspired deep models commonly use a fixed depth even though inputs may need different amounts of computation. The paper jointly learns a feed-forward predictor and variational stopping policy through oracle learning and imitation, and reports effectiveness across diverse tasks. Its label-conditioned oracle or posterior is unavailable at test time.

  • Problem

    Algorithm-inspired deep models generally use fixed depth, although different inputs may need different numbers of steps to avoid over-thinking or unnecessary computation.

    Source: §1 · p. 1

  • Method

    The method jointly learns a feed-forward predictive model and variational stopping policy, training an oracle model before imitating its stopping behavior with a sequential policy.

    Source: §1 · p. 2, §4.2 · p. 5

  • Results

    Experiments demonstrate effectiveness in prediction accuracy and inference efficiency across sparse recovery, task-imbalanced few-shot meta-learning, and computer vision.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The framework models and trains deep networks with input-specific depth across a wide range of applications.

    Source: §6 · p. 9

    The oracle stopping distribution and label-conditioned posterior require the true label, so they cannot be used at test time.

    Source: §3.4 · p. 4, §4.2 · p. 5

Abstract

from arXiv · show

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.

1. Introduction

Algorithm-inspired deep models often use fixed depth, although inputs may need different numbers of steps to save computation or avoid over-thinking. The paper proposes learning an input-specific stopping policy alongside a shared predictive model, using a two-stage training procedure.

1. Introduction

  • Fixed-depth algorithm-inspired models ignore that different inputs may need different numbers of steps, potentially wasting computation or over-thinking.Source: §1 · p. 1
  • The proposed steerable architecture jointly learns a shared feed-forward predictor and a stopping policy that selects an input-specific number of layers.Source: §1 · p. 2
  • The stopping time is treated as a latent variable in a variational Bayes objective, with prediction learning tied to marginal likelihood and stopping-policy learning to posterior approximation.Source: §1 · p. 2
  • Training is decomposed into oracle model learning, which uses label information, and imitation learning, which trains a sequential policy usable at test time without labels.Source: §1 · p. 2
  • The paper reports experiments across sparse recovery, task-imbalanced few-shot meta-learning, and computer vision, assessing prediction accuracy and inference efficiency.Source: §1 · p. 2

2. Related Works

Prior algorithm-unrolling methods generally fix the iteration count across inputs, while related work explores task adaptation, early exits, and optimal stopping. The paper positions adaptive stopping as complementary to these approaches.

2. Related Works

  • Unrolled algorithm-based models assign a fixed iteration count across inputs, whereas this method adapts stopping time to input variation.Source: §2 · pp. 2–3
  • For task-imbalanced meta-learning, the paper proposes task-specific adaptation steps alongside prior task-adaptive parameters and metric scaling.Source: §2 · p. 3
  • Image-recognition early exits often use heuristic confidence policies, and their predictor training does not account for the stopping policy's effect.Source: §2 · p. 3
  • Optimal-stopping methods based on HJB equations face dimensionality challenges, while prior deep-learning work learns stopping without jointly considering predictor parameters.Source: §2 · p. 3

3. Problem Formulation

The section formulates adaptive depth as a shared predictor with a sequential stopping policy, whose decisions induce a distribution over stopping times. It casts stopping time as a latent variable in a variational Bayes model and relates the training objective to β-VAE under stated assumptions.

3.1. Steerable Model

  • The T-layer predictor Fθ generates states x1 through xT, and its computation can stop at different input-dependent times.Source: §3.1 · p. 3
  • At each step, πφ uses input x and current state xt to choose a stop probability, allowing remaining computation to be saved after stopping.Source: §3.1 · p. 3

3.2. From Sequential Policy To Stop Time Distribution

  • Sequential stop probabilities induce qφ over stopping times; for t<T, qφ(t) combines survival to t with stopping at t, while qφ(T) is residual survival.Source: §3.2 · p. 3

3.3. Optimization Objective

  • The optimization objective averages predictive loss over the discrete stop-time distribution qφ(t) and includes entropy regularization weighted by β.Source: §3.3 · p. 4

3.4. Variational Bayes Perspective

  • The stopping time is latent with a uniform prior across layers, θ controls the likelihood, and the true posterior requires label y, unavailable at test time.Source: §3.4 · p. 4
  • With negative log-likelihood loss and a uniform prior over t, minimizing the stated loss is equivalent to maximizing the β-VAE objective.Source: §3.4 · p. 4
  • The variational Bayes view supports a two-step procedure for effectively learning the predictive model and stopping policy.Source: §3.4 · p. 4

4. Effective Training Algorithm

The training procedure replaces slow alternating optimization with an oracle-learning stage followed by imitation, with optional joint fine-tuning. The oracle accounts for input-specific stopping, while the learned sequential policy avoids reverse-KL mode collapse and can be trained more efficiently using approximations.

4. Effective Training Algorithm

  • The procedure addresses slow, schedule-sensitive alternating optimization by learning an oracle stopping distribution first, then training a policy to imitate it; joint fine-tuning is optional.The optional joint β-VAE fine-tuning stage did not improve much over the first two stages.Source: §4 · pp. 4–5

4.1. Oracle Stop Time Distribution

  • The oracle q*θ is a closed-form stopping distribution that selects a layer using loss or likelihood, with exploration controlled by β.It is equivalent to the posterior over stopping time given the label and input, and introduces no parameters beyond θ.Source: §4.1 · p. 5

4.2. Stage I. Predictive Model Learning

  • Stage I learns θ using the oracle’s optimal, data-specific stopping distribution, and gradients can propagate through that differentiable distribution.Because the oracle requires the true label and is not a sequential policy, it cannot be used at test time.Source: §4.2 · p. 5

4.3. Stage II. Imitation With Sequential Policy

  • Stage II fixes θ and trains the sequential policy to mimic the oracle; reverse KL may collapse to a common stopping time, while forward KL performed better experimentally.The reverse-KL objective can be framed as maximum-entropy reinforcement learning with stop/continue actions.Source: §4.3 · pp. 5–6
  • Forward KL performed better experimentally than reverse KL, which may collapse to a common stopping time across inputs.The reverse-KL failure mode is a distribution that captures one commonly best average stopping time rather than spread-out, input-specific times.Source: §4.3 · p. 6

4.5. Implementation Details For Efficient Training

  • Training costs more computation and memory because both stage objectives sum over layers; Monte Carlo sampling and MAP targets are proposed efficiency options.The MAP estimate can also improve accuracy by focusing policy learning on each sample’s best layer.Source: §4.5 · p. 6

5. Experiments

Experiments assess adaptive stopping in sparse recovery, task-imbalanced few-shot learning, and image denoising, with an additional image-recognition exploration. The reported results include improved recovery, meta-learning accuracy, and denoising generalization, while image recognition exposes difficulty imitating the oracle stopping policy.

5. Experiments

  • The experiments compare adaptive stopping in sparse recovery, few-shot meta learning, and image denoising, and also explore image recognition.The comparisons use ablations to examine stopping-policy effects under the same predictive architecture and the training algorithm's effectiveness against alternating EM.Source: §5 · p. 6

5.1. Learning To Optimize: Sparse Recovery

  • LISTA-stop improves recovery over LISTA and, with ≤20 iterations, outperforms 100-iteration ISTA and FISTA.The comparison is reported for recovery performance, especially on more difficult tasks at SNR 20.Source: §5.1 · p. 7, Table 2 · p. 7

5.2. Task-imbalanced Meta Learning

  • MAML-stop consistently outperforms vanilla MAML on task-imbalanced tasks, with a 3.5% accuracy improvement on the more difficult MiniImagenet setting.On the vanilla meta-learning setting, MAML-stop achieves comparable or better performance.Source: Table 4 · p. 7, §5.2 · p. 8, Table 5 · p. 8

5.3. Image Denoising

  • DnCNN-stop outperforms DnCNN, with significantly better generalization at noise levels 65 and 75 that were unseen during training.BM3D and WNNM are evaluated with known noise levels, unlike the deep models' blind-denoising setting, so that comparison is not completely fair.Source: §5.3 · p. 8, Table 6 · p. 8

5.4. Image Recognition

  • With 14 internal classifiers, oracle stopping raises VGG16 accuracy to 83.26%, but the learned policy is hard to imitate and reaches accuracy similar to SDN's heuristic policy.The authors identify imitation of the oracle policy as a bottleneck in this exploration.Source: Table 7 · p. 8, §5.4 · p. 8

6. Conclusion

The paper presents a generic framework for input-specific depth controlled by a stopping policy and reports extensive experiments across a wide range of applications.

6. Conclusion

  • The framework models and trains deep learning models with input-specific depth determined by a stopping policy.The conclusion identifies broader incorporation of algorithmic aspects into deep models as a future direction.Source: §6 · p. 9
  • Experiments across a wide range of applications are reported to demonstrate the effectiveness of both the model and training algorithm.Source: §6 · p. 9
  • The conclusion proposes exploring whether other algorithmic aspects can be incorporated into deep learning models for performance or theoretical understanding.Source: §6 · p. 9

A.1. Proof of Lemma 1

The appendix connects the variational objective to a loss minimization and maximum-entropy reinforcement learning formulation, then records experimental setup details for sparse recovery and meta-learning.

A.1. Proof of Lemma 1

  • Under the stated likelihood and uniform stopping-time prior, maximizing the β-VAE objective is equivalent to minimizing L.The likelihood is defined as ℓ(y, x_t; θ) := −log p_θ(y|t, x), and the prior over t is uniform.Source: §A.1 · p. 12

A.2. Equivalence of reverse KL and maximum-entropy RL

  • The appendix formulates stopping as an action and specifies reward and transition components for relating the optimization to maximum-entropy reinforcement learning.The action is whether to stop or continue; continuing transitions to F_θ(x_t).Source: §A.2 · p. 13

B.1. Learning To Learn: Sparse Recovery

  • Sparse recovery uses a fixed 3,000-sample test set split evenly across three SNR levels and evaluates performance by NMSE in dB.The synthetic setup uses m = 250, n = 500 and Gaussian sensing-matrix entries, with normalized columns.Source: §B.1 · p. 13

B.2.1. DETAILS OF SETUP

  • For imbalanced few-shot experiments, MAML training uses 60,000 outer-loop updates, while task datasets vary the number of training observations across the specified shot range.The setup uses 5 inner SGD steps for Omniglot and 10 for MiniImagenet; MAML-stop uses 10 for both.Source: §B.2.1 · p. 13

B.2.2. MEMORY EFFICIENT IMPLEMENTATION

The straightforward MAML-stop implementation incurs memory costs that grow with the maximum number of inner updates; stochastic EM and gradient-storage control make training nearly as memory efficient as MAML.

B.2.2. MEMORY EFFICIENT IMPLEMENTATION

  • With a 10-step maximum, straightforward MAML-stop training uses 10× the memory for CNN prediction on the meta-test set compared with vanilla MAML.The larger maximum is preferred to use the approach’s automated stopping capacity, but requires meta-test loss at every inner update.Source: §B.2.2 · p. 13
  • Stochastic EM samples a stopping time from the oracle distribution and optimizes the predictive likelihood for that sampled step.Source: §B.2.2 · p. 14
  • Computing the oracle distribution without gradient storage, then recomputing the sampled-step likelihood, makes MAML-stop almost as memory efficient as MAML.Source: §B.2.2 · p. 14

B.2.3. STANDARD META-LEARNING TASKS

The standard few-shot section includes MAML-stop as an ablation against vanilla MAML and describes a shared hyperparameter-selection procedure.

B.2.3. STANDARD META-LEARNING TASKS

  • The standard few-shot experiments include MAML-stop and compare it mainly with vanilla MAML for an ablation study.Source: §B.2.3 · p. 14
  • MAML and MAML-stop use the same batch size, training epochs, and learning rate, with selected configurations chosen by grid search on 600 held-out tasks.The search varies unrolling steps, β, and inner-update learning rate for Omniglot 20-way and MiniImagenet 5-way experiments.Source: §B.2.3 · p. 14

B.3.1. IMPLEMENTATION DETAILS

Denoising training used a two-stage loss schedule, followed by policy training and joint training with reused hyperparameters; the appendix also provides denoising visualizations and computing setup details.

B.3.1. IMPLEMENTATION DETAILS

  • Denoising models were trained for 50 epochs with unweighted loss, then 50 more with weighted loss; the exploration coefficient β was set to 0.1.Training used 403K augmented 50 × 50 patches, batch size 256, and Adam at 1e−4; policy training used Adam at 1e−4, and joint training reused the hyperparameters.Source: §B.3.1 · p. 14

B.3.2. VISUALIZATION

B.4. Computing infrastructure

  • Experiments typically used one or two V100 GPUs, each assigned six CPU threads and 64 GB of CPU memory.The typical CPU processor was an Intel Xeon Platinum 8260L.Source: §B.4 · p. 15
Loading 2006.05082v1…