TLDR: Decision-Pretrained Transformers (DPTs) often struggle to generalize to new environments, especially with biased training data. Researchers propose the Prediction-Powered Transformer (PPT) framework, which introduces “in-context curiosity” – an exploration-inspired regularizer. PPT augments DPT with an auxiliary reward predictor, using prediction error as a curiosity signal to encourage broader exploration during offline pretraining. This leads to improved robustness and better generalization in varied test environments, particularly in multi-armed bandit tasks.
Large language models (LLMs) are increasingly being integrated into decision-making tasks, often through a framework known as Decision-Pretrained Transformers (DPTs). While DPTs show great promise, a significant challenge they face is their limited ability to generalize beyond the specific data they were trained on. This means they often struggle when encountering new situations or environments that differ from their pretraining data distribution.
This limitation is particularly evident when DPTs are trained on biased or undiverse datasets. In such scenarios, they tend to learn specific patterns that work well within the training distribution but fail dramatically in out-of-distribution (OOD) settings. This issue is especially critical in simpler tasks like multi-armed bandits (MABs), where robust generalization is key.
To address this, researchers Huitao Yang from the University of California, Los Angeles, and Guanting Chen from the University of North Carolina at Chapel Hill, have introduced a novel approach called in-context curiosity. This concept is integrated into a new framework they propose: the Prediction-Powered Transformer (PPT). Their work, detailed in the paper “In-Context Curiosity: Distilling Exploration for Decision-Pretrained Transformers on Bandit Tasks”, aims to enhance the robustness and generalization capabilities of DPTs.
What is In-Context Curiosity?
Inspired by the idea of curiosity in online reinforcement learning (RL), in-context curiosity is a lightweight, exploration-driven regularizer applied during the offline pretraining phase. Unlike traditional curiosity methods that rely on real-time interaction with an environment, this approach embeds the curiosity signal directly into the offline training objective.
The PPT framework augments a standard DPT with an auxiliary reward predictor. This predictor’s job is to estimate the expected reward for each possible action. The core idea is to use the prediction error – the difference between the predictor’s estimate and the actual reward – as an intrinsic curiosity signal. When the predictor is uncertain about an action’s outcome (i.e., its prediction error is high), the policy is encouraged to explore that action more.
How PPT Works
In PPT, there are two main components: an autoregressive transformer that acts as the policy model, and a sequential predictor. During training, the predictor learns to estimate the mean reward for each action. The policy model is then trained using a standard negative log-likelihood loss, but with an added weighted curiosity term. This curiosity term is derived from the squared error between the predictor’s estimated rewards and the true mean rewards. Essentially, the policy is biased towards actions whose outcomes are less predictable, thereby encouraging broader exploration during training.
At test time, the predictor first generates reward estimates, which are then fed into the policy alongside the observed history. The policy then selects actions based on both past observations and these new predictions, enabling curiosity-driven exploration even in new environments.
Empirical Findings and Benefits
The researchers conducted experiments on Gaussian multi-armed bandit environments, comparing PPT against standard DPT. The results showed that PPT significantly improves generalization and robustness. It effectively moderates the performance degradation observed in DPTs, especially when test environments exhibit higher variance in rewards or when pretraining data has limited diversity (referred to as “tricky datasets”).
PPT demonstrated a slower rate of performance decline as environment variance increased, indicating better stability. The strength of this exploration signal is controlled by a parameter called lambda (λ). Moderate values of λ were found to improve robustness without sacrificing too much in-distribution performance. While the training cost of PPT is roughly double that of DPT (due to training both the policy and predictor), this can be reduced if a pretrained predictor is available.
Also Read:
- The Double-Edged Sword: How LLM Training Boosts Performance But Fosters Greed in Decision-Making
- Enhancing LLM Agent Training with Principle-Based Process Rewards and Normalization
Limitations and Future Directions
Despite its promising results, PPT is not a complete solution to the exploration problem. Its effectiveness is still fundamentally linked to the quality and coverage of the pretraining data. In highly variable test environments, the advantage of PPT gradually diminishes. Additionally, the current method ideally requires access to exact reward information, though preliminary tests suggest that proxy estimators can also yield competitive performance.
Future work includes extending the framework beyond bandit tasks to more complex in-context reinforcement learning scenarios with state transitions. This would involve designing curiosity mechanisms that account for uncertainty in environment dynamics, not just reward functions. Further optimization of the curiosity algorithm itself, such as developing adaptive strategies for selecting the curiosity coefficient (λ), is also a key area for future research.
