How to process TransformerEncoderLayer output in pytorch

user_13238615 0
Viewed 1127

I am trying to use bio-bert sentence embeddings for text classification of longer pieces of text.

As it currently stands I standardize the number of sentences in each piece of text (some sentences are only comprised of ("[PAD]") and run each sentence through biobert to get sentence vectors as they do here: https://pypi.org/project/biobert-embedding/

I then run those embeddings through a TrasnformerEncoder with 8 layers and 16 attention heads.

The TrasnformerEncoder outputs something of shape (batch_size, num_sentences, embedding_size).

I then try to decode this with a Linear layer and map it to my classes (of which there are 7) and softmax the output to get probabilities.

My loss function is simply nn.CrossEntropyLoss().

At first, I summed over dimensions 1 of the TransformerEncoder output to get something of size (batch_size, embedding_size). This invariable led to my network converging on always predicting one of the labels with absolute certainty. Usually the most common label in the dataset.

I then tried only taking the output for the last sentence of the TransformerEncoder output. i.e. TransformerEncoderOutput[:, -1, :].

This resulted in something similar.

I then tried running my Linear layer on each of the outputs of TransformerEncoder output to produce a tensor of size (batch_size, num_sentences, 7). I then sum over dim 1 (makes a tensor of size (batch_size, 7) and softmax as usual. The idea here is that every sentence gets to vote for the label after being informed about its place in the sequence.

This converged even more quickly to just predicting 1 for one of the labels and vanishingly small values for the others.

I feel like I am misunderstanding out to use the output of a pytorch Transformer somehow. My learning rate is very low, 0.00001, and that helped delay the convergence but it converged eventually anyway.

What this suggests to me is that my network is incapable of figuring anything out about the text and is just learning to find the most common labels. I would guess that this is either a problem with my loss function or a problem with how I am using the Transformer.

Is there a glaring flaw in the architecture that I have laid out?

transformer-model bert-language-model pytorch
1 Answers

So the input and output shape of the transformer-encoder is batch-size, sequence-length, embedding-size). There are three possibilities to process the output of the transformer encoder (when not using the decoder).

  1. you take the mean of the sequence-length dimension:
x = self.transformer_encoder(x)
x = x.reshape(batch_size, seq_size, embedding_size)      
x = x.mean(1)
  1. sum it up as you said:
x = self.transformer_encoder(x)
x = x.reshape(batch_size, seq_size, embedding_size)      
x = x.sum(1)
  1. using a recurrent neural network to combine the information along the sequence-length dimension:
x = self.transformer_encoder(x)
x = x.reshape(batch_size, seq_len, embedding_size)      

# init hidden state
hidden = torch.zeros(layers, batch_size, embedding_size).to(device=device)  
x, hidden = self.rnn(x, hidden)
x = x.reshape(batch_size, seq_size, embedding_size)

# take last output
x = x[:, -1]

Taking the last element of the Transformer output isnt a good idea I think. Because then you only take 1 / seq-len of the information. But using rnn, the last output has still information of every other output.

Id say that taking the mean is the best idea.

And to the learning rate: For me it always worked very much better when I used warmup training. If you dont know what that is: You start at a low learning rate, for example 0.00001 and you increase it until you have reached some target lr, for example 0.002. And from then you just decay the lr as usual.

user_89093530
Related