Input stacking. An important hyperparameter for transformer models is what dimension to use self-attention over. Previous work applies it across states, actions, and rewards as separate tokens (or even individual state and action dimensions) [1]}, [2]}; this can increase the effective sequence length that we would need to input a trajectory snippet of length \(k\) : for example if treating states, actions, and rewards separately (have self-attention act on each independently), the sequence length would be \(3k\) . While this is not an issue for the short context windows we use in our experiments, this seems wasteful: the main bottleneck for transformer models is usually the computational cost of self-attention, which scales quadratically in the sequence length. To obviate this problem, we stack states, actions, and rewards for each timestep, treating them as single inputs. This way, we are making self-attention happen only across timesteps, reducing the self-attention sequence length required to \(k\) . This also seems like a potentially advantageous inductive bias for improving performance. Though we did not test this systematically, preliminary experiments did show that input stacking reduced validation loss.
