Use Your INSTINCT: INSTruction optimization for LLMs usIng Neural bandits Coupled with Transformers
Xiaoqiang Lin, Zhaoxuan Wu, Zhongxiang Dai, Wenyang Hu, Yao Shu, See-Kiong Ng, Patrick Jaillet, Bryan Kian Hsiang Low
TL;DR
The paper addresses the problem of automatically optimizing prompts for black-box LLMs, where costly API calls constrain exploration. It introduces INSTINCT, a neural-bandit framework that replaces Gaussian-process surrogates with a neural surrogate and couples it to the hidden representations of a pre-trained transformer to predict instruction quality. The approach uses NeuralUCB to balance exploration and exploitation, and employs pre-computation over a discretized soft-prompt domain to reduce costs while adapting the intrinsic prompt dimensionality via random projections. Empirically, INSTINCT yields consistent improvements over state-of-the-art baselines across instruction induction tasks and zero-shot chain-of-thought prompts, demonstrates robustness across white-box/black-box LLM pairings, and accommodates in-context learning enhancements, offering a scalable, data-efficient path for automatic instruction tuning.
Abstract
Large language models (LLMs) have shown remarkable instruction-following capabilities and achieved impressive performances in various applications. However, the performances of LLMs depend heavily on the instructions given to them, which are typically manually tuned with substantial human efforts. Recent work has used the query-efficient Bayesian optimization (BO) algorithm to automatically optimize the instructions given to black-box LLMs. However, BO usually falls short when optimizing highly sophisticated (e.g., high-dimensional) objective functions, such as the functions mapping an instruction to the performance of an LLM. This is mainly due to the limited expressive power of the Gaussian process (GP) which is used by BO as a surrogate to model the objective function. Meanwhile, it has been repeatedly shown that neural networks (NNs), especially pre-trained transformers, possess strong expressive power and can model highly complex functions. So, we adopt a neural bandit algorithm which replaces the GP in BO by an NN surrogate to optimize instructions for black-box LLMs. More importantly, the neural bandit algorithm allows us to naturally couple the NN surrogate with the hidden representation learned by a pre-trained transformer (i.e., an open-source LLM), which significantly boosts its performance. These motivate us to propose our INSTruction optimization usIng Neural bandits Coupled with Transformers (INSTINCT) algorithm. We perform instruction optimization for ChatGPT and use extensive experiments to show that INSTINCT consistently outperforms baselines in different tasks, e.g., various instruction induction tasks and the task of improving zero-shot chain-of-thought instructions. Our code is available at https://github.com/xqlin98/INSTINCT.
