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

Perry Dong, Kuo-Han Hung, Alexander Swerdlow, Dorsa Sadigh, Chelsea Finn
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:25856-25876, 2026.

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-dong26d, title = {{TQL}: Scaling Q-Functions with Transformers by Preventing Attention Collapse}, author = {Dong, Perry and Hung, Kuo-Han and Swerdlow, Alexander and Sadigh, Dorsa and Finn, Chelsea}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {25856--25876}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/dong26d/dong26d.pdf}, url = {https://proceedings.mlr.press/v306/dong26d.html}, 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.} }
Endnote
%0 Conference Paper %T TQL: Scaling Q-Functions with Transformers by Preventing Attention Collapse %A Perry Dong %A Kuo-Han Hung %A Alexander Swerdlow %A Dorsa Sadigh %A Chelsea Finn %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-dong26d %I PMLR %P 25856--25876 %U https://proceedings.mlr.press/v306/dong26d.html %V 306 %X 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.
APA
Dong, P., Hung, K., Swerdlow, A., Sadigh, D. & Finn, C.. (2026). TQL: Scaling Q-Functions with Transformers by Preventing Attention Collapse. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:25856-25876 Available from https://proceedings.mlr.press/v306/dong26d.html.

Related Material