Owing to their ability to both effectively integrate information over long time horizons and scale to massive amounts of data, self-attention architectures have recently shown breakthrough success in natural language processing (NLP), achieving state-of-the-art results in domains such as language modeling and machine translation. Harnessing the transformer's ability to process long time horizons of information could provide a similar performance boost in partially observable reinforcement learning (RL) domains, but the large-scale transformers used in NLP have yet to be successfully applied to the RL setting. In this work we demonstrate that the standard transformer architecture is difficult to optimize, which was previously observed in the supervised learning setting but becomes especially pronounced with RL objectives. We propose architectural modifications that substantially improve the stability and learning speed of the original Transformer and XL variant. The proposed architecture, the Gated Transformer-XL (GTrXL), surpasses LSTMs on challenging memory environments and achieves state-of-the-art results on the multi-task DMLab-30 benchmark suite, exceeding the performance of an external memory architecture. We show that the GTrXL, trained using the same losses, has stability and performance that consistently matches or exceeds a competitive LSTM baseline, including on more reactive tasks where memory is less critical. GTrXL offers an easy-to-train, simple-to-implement but substantially more expressive architectural alternative to the standard multi-layer LSTM ubiquitously used for RL agents in partially observable environments.
* Summary
This paper introduces architecture modifications for self-attention to stabilize transformers in reinforcement learning.
The new architecture, Gated Transformer-XL, replaces the order of the layer norm blocks to preserve an identity mapping.
Multiple existing gating layers are proposed to replace the residual connections of transformer.
The new architecture is compared against MERLIN and LSTM on DMlab-30, and further ablation studies are done on Numpad and Memory Maze.
It is noted that they use the recent V-MPO objective to train LSTM and transformer.
Their results show that transformers are able to learn in memory-intensive environments, with some gating combinations surpassing LSTM.
* Decision
This paper presents promising empirical results, however the experiments are limited, making it difficult to place in the broader work.
In addition, the contribution is incremental and not well-motivated.
I would recommend a weak rejection.
Still, I think the paper is well written and could be improved upon.
* Reasons
While the empirical results are impressive, they are not put into context.
DMlab-30 is still a relatively new environment suite and it is difficult to place the result of this paper in the context of broader work.
In addition, the comparison is against LSTM on a new objective.
While Transformer is able to beat LSTM on the same objective, it is unclear whether that is a success of the objective or the architecture.
In Numpad, the transformer architecture shows an improvement over LSTM, but no comparisons are made to any other memory-based agents.
The hyperparameter studies on Memory Maze also show improvements in memory-related tasks, but do not help in understanding of the proposed work.
The choice of architecture modification is also not well motivated.
As the paper mentions, initializing near identity has been shown to be important in the supervised learning literature.
For the topic of this paper however, I do not think this adequetly explains the instability of transformers in reinforcement learning.
In the related work section for example, the paper notes that gating mechanisms have been used to handle the vanishing gradients problem.
The paper also notes that vanishing gradients is not an issue in transformers.
Hence, it is unclear why gating would stabilize transformers for reinforcement learning.
The paper is overall well written and the ideas developed are clear.
Unfortunately, the impressive results on DMlab are not sufficient for both the lack of deeper empirical study and better theoretical motivation for the architecture modifications.