This work was done as an experiment for Boaz Barak’s “CS 2881r: AI Safety and Alignment” at Harvard. The lecture where this work was presented can be viewed on YouTube here, and its corresponding blogpost can be found here.

Background

Prompt engineering has become a central idea in working with large language models (LLMs) — the way a question or instruction is phrased can make a big difference in the quality of responses. In this toy experiment, we sought out to see if reinforcement learning (RL) could be a principled way of discovering prompt prefixes that condition models toward more accurate, aligned, or stylistically appropriate outputs.

The Setup

Assume that we have a fixed set of prompt prefixes to choose from. Our idea was to model the problem of choosing the best prefix as a multi-armed bandit:

The policy is a distribution over a set of predefined prompts.
For each incoming task, the policy selects a prompt prefix.
The LLM generates a response.
We compute a reward based on how well the answer matches some criterion (e.g., correctness, clarity, or professionalism).
The reward updates the policy with the UCB algorithm, shifting probability mass toward the more effective personas.

In principle, this setup allows us to converge on prompts that are consistently high-performing across a dataset.

We will highlight below three different small experiments we ran, for different tasks and different sets of prompt options.

The Experiment(s)

Experiment 1: Original Setup

We generated a pool of personas by sampling notable individuals from Wikipedia. These included scientists, artists, historical figures, and contemporary public figures.
The prompt prefix was: "Answer the task below as if you were [Persona X]."
Evaluation Dataset: GSM8K.

Results

The results were somewhat predictable: the model’s accuracy on math problems did not meaningfully depend on the persona. Whether it was Alan Turing, Marie Curie, or Beyoncé, the LLM solved the math problem correctly, and thus the reward model could not be calculated only based on the mathematically correct answer.

That said, we did notice some subtle differences. Alan Turing often “won” under our reward function, since his answers tended to be phrased in a more precise and mathematical style. This hinted that persona prompts could shift the style of answers, even if they didn’t affect raw problem-solving accuracy.

“Alright, le

...

Posts

Wikitag Contributions

Comments