🤖 AI Summary
This work addresses the suboptimal decision-making of existing Transformer-based agents in non-stationary, partially observable environments, where reliance solely on observation similarity for attention retrieval fails to distinguish between distinct action–reward histories under identical observations. To overcome this limitation, the paper introduces the Utility-Augmented Transformer (UAT), which formally characterizes the “feedback-blind retrieval” problem and incorporates a utility-state-modulated attention mechanism. This mechanism explicitly integrates action–reward history into query, key, and value projections to guide contextual retrieval. The proposed architecture features a zero-gate degeneracy property, strictly generalizing the representational capacity of observation-only Transformers, and offers theoretical guarantees under Lipschitz continuity and finite-horizon assumptions. Empirical results demonstrate that UAT significantly outperforms current baselines across four non-stationary benchmark tasks, particularly excelling in high-noise regimes and rapid adaptation scenarios.
📝 Abstract
Sequential decision making in non-stationary and partially observable environments requires rapid adaptation to latent regime changes. However, existing Transformer decision models face a structural bottleneck in the retrieval mechanism: even when reward is used for training or exposed as an input token, attention retrieval remains primarily driven by observation-derived similarity. We formalize this limitation as feedback-blind retrieval, and formally show that, on feedback-informative tasks, observation-equivalent histories with different action-reward outcomes cannot be distinguished by any observation-only attention, resulting in suboptimal choice. To address this mismatch, we propose the Utility-Augmented Transformer (UAT), a new feedback-conditioned retrieval attention architecture in which a compact utility state modulates the query, key, and value projections, allowing action-reward history to directly alter context retrieval during the forward pass. UAT also enjoys an exact zero-gate degradation property that recovers the Vanilla Transformer when feedback is uninformative. Under finite-horizon compactness and Lipschitz assumptions, we prove that UAT strictly enlarges the observation-only Transformer class and can uniformly approximate feedback-dependent decision maps. Across four non-stationary benchmarks: synthetic navigation with hidden goal shifts, non-stationary sepsis treatment, cross-market portfolio allocation, and delayed-feedback recommendation, UAT consistently improves performance over observation-only, test-time adaptation, and input-level feedback baselines, with particularly large gains in noisier regimes that require stronger adaptation.