<p>Yes, this is directly related to the way that BERT is trained. Specifically, I encourage you to have a look at the <a href="https://arxiv.org/pdf/1810.04805.pdf" rel="nofollow noreferrer">original BERT paper</a>, in which the authors introduce the meaning of the <code>[CLS]</code> token:</p>
<blockquote>
<p><code>[CLS]</code> is a special symbol added in front of every input example [...].</p>
</blockquote>
<p>Specifically, it is used for classification purposes, and therefore the first and simplest choice for any fine-tuning for classification tasks. What your relevant code fragment is doing, is basically just extracting this <code>[CLS]</code> token.</p>
<p>Unfortunately, the DistilBERT documentation of Huggingface's library does not explicitly refer to this, but you rather have to check out their <a href="https://huggingface.co/transformers/model_doc/bert.html" rel="nofollow noreferrer">BERT documentation</a>, where they also highlight some issues with the <code>[CLS]</code> token, analogous to your concerns:</p>
<blockquote>
<p>Alongside MLM, BERT was trained using a next sentence prediction (NSP)
objective using the [CLS] token as a sequence approximate. The user
may use this token (the first token in a sequence built with special
tokens) to get a sequence prediction rather than a token prediction.
However, averaging over the sequence may yield better results than
using the [CLS] token.</p>
</blockquote>
why take the first hidden state for sequence classification (DistilBertForSequenceClassification) by HuggingFace
<p>In the last few layers of sequence classification by <a href="https://github.com/huggingface/transformers/blob/33d3072e1c54bcd235447b98c6dea1b4cb71234c/src/transformers/modeling_distilbert.py#L634" rel="noreferrer">HuggingFace</a>, they took the first hidden state of the sequence length of the transformer output to be used for classification. </p>

<pre><code>hidden_state = distilbert_output[0]  # (bs, seq_len, dim) &lt;-- transformer output
pooled_output = hidden_state[:, 0]  # (bs, dim)           &lt;-- first hidden state
pooled_output = self.pre_classifier(pooled_output)  # (bs, dim)
pooled_output = nn.ReLU()(pooled_output)  # (bs, dim)
pooled_output = self.dropout(pooled_output)  # (bs, dim)
logits = self.classifier(pooled_output)  # (bs, dim)
</code></pre>

<p>Is there any benefit to taking the first hidden state over the last, average, or even the use of a Flatten layer instead?</p>