What to do with Transformer Encoder output?

Question

I'm in the middle of learning about Transformer layers, and I feel like I've got enough of the general idea behind them to be dangerous. I'm designing a neural network and my team would like to include them, but we're unsure how to proceed with the encoded sequences and what the right way to plug the them into the next layer in the model would be. We would like to process it such that we can plug the encoded sequence into a FC layer immediately after the Transformer Encoder.

If we just use a batch size of 1, for the sake of the argument, our encoded sequence output after being processed by the Transformer Encoder has shape tuple of (L,E), where L is the input sequence length and E is the embedded dimension size. I've seen some vague description of using some max/avg/conv1d pooling on the Encoded sequence, but nothing super clear about what that means. If I'm following this correctly, would I apply the max/avg/1conv1d pooling such that the pooling result gives me an resulting vector with shape tuple (E,), or would I pool along the other dimension?

Answer 1

The typical approach for this is follow BERT's approach: add an extra special token at the beginning of the input sequence (in BERT it is [CLS]) and only use the output of the network at that position as input to your fully connected layer. The output at the rest of the positions is ignored.

You can see a nice illustration of this approach in the notorious blog post The Illustrated BERT, which explains very visually all the details about BERT:

In the illustration, you can see the model input at the bottom and how it has been added a special [CLS] token at the beginning and then the output of the model at that position is then used for a classification task.

During training, the model will learn to condense the needed information from the whole sentence into the output of the first position.

Another alternative, as you pointed out, is to have global average pooling over all the outputs. This was the norm in the LSTM times before Transformers came. I am not aware of any articles comparing the performance of both approaches but, nowadays, with Transformers, everybody uses the BERT approach.

Both BERT's approach and the global average/max pooling approach achieve your goal: collapsing the variable length sequence of vectors into a single vector that you can then use for classification.