title: The last layers of longformer for document embeddings
body: <p>What's the right way to return a limited number of layers using the longformer API?</p>
  <p>Unlike this case in basic <a href="https://stackoverflow.com/questions/63461262/bert-sentence-embeddings-from-transformers">BERT</a>, it's not clear to me from the return type how to get only the last N layers.</p>
  <p>So, I run this:</p>
  <pre><code>from transformers import LongformerTokenizer, LongformerModel
  
  text = &quot;word &quot; * 4096 # long document!
  
  tokenizer = LongformerTokenizer.from_pretrained('allenai/longformer-base-4096')
  model = LongformerModel.from_pretrained(&quot;allenai/longformer-base-4096&quot;)
  
  encoded_input = tokenizer(text, return_tensors=&quot;pt&quot;, max_length=4096, truncation=True)
  output = model(**encoded_input)
  </code></pre>
  <p>And I get dimensions like so from my return:</p>
  <pre><code>&gt;&gt;&gt; output[0].shape
  torch.Size([1, 4096, 768])
  
  &gt;&gt;&gt; output[1].shape
  torch.Size([1, 768])
  </code></pre>
  <p>You can see the shape of [0] is curiously similar to my number of tokens.  I believe that slicing this would just give me fewer tokens, not just the last N layers.</p>
  <h2>Update from answer below</h2>
  <p>Even asking for <code>output_hidden_states</code>, the dimensions still look off, and it's not clear to me
  how to reduce these to vector sized, 1-d embedding.  Here's what I mean:</p>
  <pre><code>encoded_input = tokenizer(text, return_tensors=&quot;pt&quot;, max_length=4096, truncation=True)
  output = model(**encoded_input, output_hidden_states=True)
  </code></pre>
  <p>Ok, now let's look into output[2], the third item of the tuple:</p>
  <pre><code>&gt;&gt;&gt; len(output[2])
  13
  </code></pre>
  <p>Suppose we want to see the last 3 of the 13 layers:</p>
  <pre><code>&gt;&gt;&gt; [pair[0].shape for pair in output[2][-3:]]
  [torch.Size([4096, 768]), torch.Size([4096, 768]), torch.Size([4096, 768])]
  </code></pre>
  <p>So we see each of the 13 layers is shaped (4096 x 768), and they look like:</p>
  <pre><code>&gt;&gt;&gt; [pair[0] for pair in output[2][-3:]]
  [tensor([[-0.1494,  0.0190,  0.0389,  ..., -0.0470,  0.0259,  0.0609],
  </code></pre>
  <p>We still have a size of 4096, in that it corresponds to my token count:</p>
  <pre><code>&gt;&gt;&gt; np.mean(np.stack([pair[0].detach().numpy() for pair in output[2][-3:]]), axis=0).shape
  (4096, 768)
  </code></pre>
  <p>Averaging these together does not seem like it would give a valid embedding (for comparisons like cosine similarity).</p>