Reproduce “Changing Model Behavior at Test-Time Using Reinforcement Learning” on OpenResearch

About this paper

Augustus Odena, Dieterich Lawson, Christopher Olah

This research introduces a deep learning framework called the "Composer Model" that uses a Reinforcement Learning-trained controller and dynamic Policy Preferences to enable a single model to adapt its computational behavior at test-time. It demonstrated that resource allocation could be intelligently adjusted per input based on external preferences, achieving various accuracy-computation trade-offs on a modified MNIST dataset.

Key insights

• A single trained model can dynamically adjust its computational budget and behavior at test-time by varying an external "preference" parameter, eliminating the need for retraining or multiple specialized models.
• Reinforcement Learning effectively trains a controller to make intelligent, input-dependent decisions about module selection, adapting resource usage to maximize performance under a given preference.
• By training the controller with a distribution of preference values, it learns to generalize and respond appropriately to new, unseen preference settings at inference time.

Summary by alphaXiv.