Stochastic Bandits with ReLU Neural Networks
Kan Xu, Hamsa Bastani, Surbhi Goel, Osbert Bastani
TL;DR
This work addresses stochastic bandits where the reward is modeled by a one-layer ReLU neural network, and shows that a two-phase approach—exploration to enter a linear regime followed by a transformed-feature linear bandit—achieves a minimax regret of $\tilde{O}(\sqrt{T})$. By leveraging the piecewise linear structure, the authors convert the problem to a linear bandit in a higher-dimensional feature space and design OFU-ReLU and its batching variant OFU-ReLU+ to remove dependence on unknown parameters. They establish a parameter-estimation bound for the neurons (up to sign) and connect small generalization error to accurate neuron recovery, enabling online learning with provable guarantees. Empirical results on synthetic ReLU-bandits demonstrate notable improvements over linear OFUL and NeuralUCB baselines, suggesting practical potential for ReLU-based bandit algorithms in limited-time regimes.
Abstract
We study the stochastic bandit problem with ReLU neural network structure. We show that a $\tilde{O}(\sqrt{T})$ regret guarantee is achievable by considering bandits with one-layer ReLU neural networks; to the best of our knowledge, our work is the first to achieve such a guarantee. In this specific setting, we propose an OFU-ReLU algorithm that can achieve this upper bound. The algorithm first explores randomly until it reaches a linear regime, and then implements a UCB-type linear bandit algorithm to balance exploration and exploitation. Our key insight is that we can exploit the piecewise linear structure of ReLU activations and convert the problem into a linear bandit in a transformed feature space, once we learn the parameters of ReLU relatively accurately during the exploration stage. To remove dependence on model parameters, we design an OFU-ReLU+ algorithm based on a batching strategy, which can provide the same theoretical guarantee.