As the attention blocks, we use the Transformer [1]} attention layers with the modifications inspired by [2]}. Unlike in Transformer, where the outputs of the multi-head attention blocks are passed through a feed-forward block directly, we first concatenate the attention outputs with the input features before passing them to the feed-forward block followed by Layernorm [3]} (see Fig. REF b). The feed-forward multilayer perceptron (MLP) is then followed by a residual connection.
All the MLPs in our network have one ReLU-activated hidden layer.
