title: How to process TransformerEncoderLayer output in pytorch
body: <p>I am trying to use bio-bert sentence embeddings for text classification of longer pieces of text.</p>
  <p>As it currently stands I standardize the number of sentences in each piece of text (some sentences are only comprised of (&quot;[PAD]&quot;) and run each sentence through biobert to get sentence vectors as they do here:
  <a href="https://pypi.org/project/biobert-embedding/" rel="nofollow noreferrer">https://pypi.org/project/biobert-embedding/</a></p>
  <p>I then run those embeddings through a TrasnformerEncoder with 8 layers and 16 attention heads.</p>
  <p>The TrasnformerEncoder outputs something of shape (batch_size, num_sentences, embedding_size).</p>
  <p>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.</p>
  <p>My loss function is simply nn.CrossEntropyLoss().</p>
  <p>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.</p>
  <p>I then tried only taking the output for the last sentence of the TransformerEncoder output. i.e. TransformerEncoderOutput[:, -1, :].</p>
  <p>This resulted in something similar.</p>
  <p>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.</p>
  <p>This converged even more quickly to just predicting 1 for one of the labels and vanishingly small values for the others.</p>
  <p>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.</p>
  <p>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.</p>
  <p>Is there a glaring flaw in the architecture that I have laid out?</p>