TQL: Scaling Q-Functions with Transformers by Preventing Attention Collapse

TL;DR

TQL addresses the core problem of scaling transformer-based value functions in reinforcement learning by diagnosing attention collapse as the key instability. It prevents collapse by directly regulating the entropy of attention scores with per-layer and per-token temperature terms toward a target entropy $\bar{H}$, while adding learnable modality embeddings and a compatible policy-extraction method. Empirically, TQL enables large transformers to scale in offline RL tasks, achieving state-of-the-art performance and a substantial scaling gain compared to baselines, with ablations confirming the importance of entropy guidance. The approach offers a simple, general mechanism to unlock the scalability of transformer-based value functions without architectural overhauls, making it practical for pretrained models and real-world offline RL deployments.

Abstract

Despite scale driving substantial recent advancements in machine learning, reinforcement learning (RL) methods still primarily use small value functions. Naively scaling value functions -- including with a transformer architecture, which is known to be highly scalable -- often results in learning instability and worse performance. In this work, we ask what prevents transformers from scaling effectively for value functions? Through empirical analysis, we identify the critical failure mode in this scaling: attention scores collapse as capacity increases. Our key insight is that we can effectively prevent this collapse and stabilize training by controlling the entropy of the attention scores, thereby enabling the use of larger models. To this end, we propose Transformer Q-Learning (TQL), a method that unlocks the scaling potential of transformers in learning value functions in RL. Our approach yields up to a 43% improvement in performance when scaling from the smallest to the largest network sizes, while prior methods suffer from performance degradation.

Figures (14)

• Figure 1: TQL unlocks scaling of value functions in RL by preventing attention collapse. Left: Scaling results of TQL compared with prior approaches with different generative model backbones. The results are reported as relative improvement over the smallest model size, averaged across seeds and all 25 tasks in our evaluation suite. TQL increases in performance as network size increases, while prior approaches are not able to effectively use the extra capacity and drop in performance. Right: Attention entropy with and without TQL, averaged across 25 tasks. TQL scales effectively by preventing attention collapse that occurs with scaling up the transformer architecture for value function training.
• Figure 2: Scaling value networks with TQL.TQL prevents attention collapse in larger networks by controlling the target entropy of attention scores (right). Compared to an unregularized model, TQL exhibits a more uniform attention distribution and a smoother value landscape (left).
• Figure 3: Scaling transformers for value functions results in entropy collapse and worse performance. Left: Visualizations of the success rate and attention entropy of the transformer model without TQL across five cube-double tasks under different model sizes. Right: Q-value landscapes and attention maps for the smallest (0.4M) and largest (26M) models. The larger transformer learns highly non-smooth value surfaces and exhibits high-frequency oscillations and discontinuities that are absent in its smaller counterpart.
• Figure 4: Scaling results. We compare TQL against prior offline RL methods across critic sizes from 0.4M to 26M parameters. The plot reports the average success rate difference compared to the smallest model (0.4M) for each method. While baselines suffer from performance degradation at larger scales, TQL consistently scales well across all environments, outperforming prior methods by a large margin.
• Figure 5: Ablation on the most important components of TQL. We compare TQL against four ablated variants on the cube-double environment: (1) a transformer baseline without attention entropy guidance, (2) using a fixed entropy penalty, and (3) using one temperature instead of layer-wise and token-wise temperatures. The results represent the average performance and standard error across 5 tasks and 3 seeds in the cube-double.