Reinforcement Learning with Generative Models for Compact Support Sets
Nico Schiavone, Xingyu Li
TL;DR
The paper tackles classification in data-scarce regimes by generating a small, targeted synthetic support set with reinforcement learning. An agent controls a pretrained generative model through a domain-class dictionary to craft prompts, optimized by PPO with a reward r(s,s') = $\Delta$Acc(s→s') − $\Delta$Entropy(s→s'), producing at most $N_{\textrm{syn}}$ images added to the training data. Key contributions include (i) a novel RL–generative framework for compact synthetic support sets, (ii) an entropy-aware reward design that promotes balanced improvement, and (iii) empirical gains on CIFAR-10 and Tiny ImageNet across multiple backbones with minimal synthetic data. The results demonstrate data-efficient augmentation that yields up to roughly 1% accuracy gains without extra labeling. This approach highlights a scalable way to integrate reinforcement learning with large pretrained generative models for targeted data augmentation in low-data scenarios.
Abstract
Foundation models contain a wealth of information from their vast number of training samples. However, most prior arts fail to extract this information in a precise and efficient way for small sample sizes. In this work, we propose a framework utilizing reinforcement learning as a control for foundation models, allowing for the granular generation of small, focused synthetic support sets to augment the performance of neural network models on real data classification tasks. We first allow a reinforcement learning agent access to a novel context based dictionary; the agent then uses this dictionary with a novel prompt structure to form and optimize prompts as inputs to generative models, receiving feedback based on a reward function combining the change in validation accuracy and entropy. A support set is formed this way over several exploration steps. Our framework produced excellent results, increasing classification accuracy by significant margins for no additional labelling or data cost.