[Paper Review] Changing Model Behavior at Test-Time Using Reinforcement Learning

Augustus Odena, Dieterich Lawson+1 more|arXiv (Cornell University)|Feb 24, 2017

Machine Learning and Data Classification19 citations

TL;DR

This paper proposes a reinforcement learning-based method to dynamically adjust model behavior at test time by using a Composer model that selects from pre-trained modules based on input-dependent constraints. By introducing Policy Preferences, the controller adapts computation usage—such as parameter count or module diversity—on-the-fly, enabling a single model to balance accuracy, efficiency, and resource use without retraining.

ABSTRACT

Machine learning models are often used at test-time subject to constraints and trade-offs not present at training-time. For example, a computer vision model operating on an embedded device may need to perform real-time inference, or a translation model operating on a cell phone may wish to bound its average compute time in order to be power-efficient. In this work we describe a mixture-of-experts model and show how to change its test-time resource-usage on a per-input basis using reinforcement learning. We test our method on a small MNIST-based example.

Motivation & Objective

• Address the challenge of adapting machine learning models to runtime constraints like speed, memory, or power without retraining.
• Enable dynamic, input-specific computation adjustments at test time, rather than hardcoding trade-offs during training.
• Provide interpretability by tracking which modules are used per input and under what preferences.
• Allow runtime modification of model behavior through tunable preferences without retraining.
• Demonstrate that a single model can achieve performance comparable to specialized models across multiple operational regimes.

Proposed method

• The Composer model uses a controller network to select from multiple modules (neural networks) at each metalayer, forming a dynamic computation graph per input.
• The controller is trained via REINFORCE to maximize log-likelihood of correct predictions, using stochastic sampling over module choices.
• Policy Preferences are introduced by augmenting the reward function with a preference-dependent cost term, where the preference γ is input to the controller at test time.
• Two types of preferences are implemented: Glimpse Preferences (to control parameter usage via module size) and Entropy Preferences (to encourage balanced module usage across batches).
• The controller learns to modulate its policy across a distribution of preferences γ during training, enabling runtime adaptation.
• The method supports per-input or per-mini-batch preference changes, allowing real-time trade-offs between performance and resource consumption.

Experimental results

Research questions

• RQ1 Can a single neural network model dynamically adjust its computational complexity at inference time based on input characteristics and user-defined constraints?
• RQ2 To what extent can a reinforcement learning controller learn to make efficient, input-aware module selection without retraining?
• RQ3 How effective is the use of Policy Preferences in enabling runtime adaptation of model behavior, such as reducing parameter usage or balancing module utilization?
• RQ4 Can the model achieve performance comparable to specialized models across different operational regimes (e.g., high-accuracy vs. low-compute) using only one trained instance?
• RQ5 Does the method provide interpretability by revealing which modules are used and under what preferences?

Key findings

• The Composer model trained with both glimpse and entropy preferences successfully adapts its computation usage at test time, achieving a range of average parameter counts across different preference settings.
• The controller learns to use the small module for 'left' digits (which are fully visible in the left half) when preference for low computation is high, preserving accuracy while reducing cost.
• Only when the glimpse preference is very low does the controller begin using the small module for 'right' digits, indicating intelligent, input-aware decisions.
• The gap between the Composer curve and the baseline model (randomly switching between modules) demonstrates that the controller makes smarter, performance-preserving choices than random selection.
• Heatmaps show that the controller’s module selection policy shifts predictably with preference values—favoring the small module for left digits and only later for right digits as cost constraints tighten.
• The method enables dynamic, on-the-fly behavior changes without retraining, and the entropy preference mechanism ensures balanced module usage without requiring a manual annealing schedule.