Answers.Body: <p>During pre-training DistilBERT tries to mimic soft target distributions for <em>all</em> tokens (not just masked tokens since by transformer's self-attention nature DistilBERT still computes logits for all positions same as its teacher) of an input sequence via the <em>required</em> distillation CE loss and the additional optional cosine embedding loss and masked language modeling (MLM) loss. Therefore, like the distillation loss, the cosine loss aligns all tokens' hidden states output by DistilBERT with their corresponding tokens' hidden states output by BERT across each position by default implicitly, unless it's explicitly claimed to only focus on masked tokens as the apparent MLM loss. See the original <a href="https://arxiv.org/pdf/1910.01108" rel="nofollow noreferrer">DistilBERT paper</a>.</p>
  <blockquote>
  <p>To leverage the
  inductive biases learned by larger models during pre-training, we introduce a triple loss combining language modeling, distillation and cosine-distance losses... The student is trained with a distillation loss over the soft target probabilities of the teacher... This objective results in a rich training signal by leveraging the full teacher
  distribution... We found it beneficial to add a cosine embedding loss which will tend to align the directions of the student and teacher hidden states vectors.</p>
  </blockquote>
  <p>We can cross-check Hugging Face transformer's distilbert model's <a href="https://github.com/huggingface/transformers/blob/main/src/transformers/models/distilbert/modeling_distilbert.py" rel="nofollow noreferrer">codebase</a> where there're only CE/BCE/MSE losses, unfortunately. Hugging Face focuses on practical efficient implementations and they might have modified the distillation process to omit the optional cosine embedding loss described in the paper for simplicity or performance, though apparently still keeps the optional MLM CE loss. We can see only MLM loss use <em>labels</em> which are essentially masked tokens from below code snippet.</p>
  <blockquote>
  <p>labels (<code>torch.LongTensor</code> of shape <code>(batch_size, sequence_length)</code>, <em>optional</em>):
  Labels for computing the masked language modeling loss. Indices should be in <code>[-100, 0, ..., config.vocab_size]</code> (see <code>input_ids</code> docstring) Tokens with indices set to <code>-100</code> are ignored (masked), the
  loss is only computed for the tokens with labels in <code>[0, ..., config.vocab_size]</code>.</p>
  </blockquote>
  <p><code>loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))</code></p>
  <p>But for the distillation loss, say for Q&amp;A start/end position prediction class, we can see there's no masked token filtering at all.</p>
  <pre><code>            ignored_index = start_logits.size(1)
              start_positions = start_positions.clamp(0, ignored_index)
              loss_fct = nn.CrossEntropyLoss(ignore_index=ignored_index)
              start_loss = loss_fct(start_logits, start_positions)
  </code></pre>
  <p>So I believe we can now safely assume that the optional cosine loss of DistilBERT recommended by the original paper also align hidden states across all tokens of an input sequence to have better (soft) prediction performance for all kinds of downstream NLP tasks which may need prediction beyond masked tokens, not just masked tokens.</p>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2024-12-08T20:01:43.440Z
Answers.Text: Thanks. Do you know how the distillation process ensures that [CLS] token's hidden state produced by DistilBERT aligns with that from BERT, given that the [CLS] token is never masked and there's no next sentence prediction in the distillation process? I started a separate post for this: https://ai.stackexchange.com/questions/47458/how-does-distilberts-training-ensure-that-the-cls-tokens-embedding-from-dist
Answers.UserDisplayName: Fijoy Vadakkumpadan
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2024-12-08T23:27:37.957Z
Answers.Text: Thanks for your followup question about the important [CLS] token in BERT/DistilBERT, which prompted me to review further the Hugging Face implementation which shows the required distillation loss is applied to all tokens and similarly I believe now the optional omitted cosine loss recommended in the original paper must be also applied to all positions to further refine important token alignment such as [CLS]. I've also replied to your follow up question. Hope it completely clarifies now.
Answers.UserDisplayName: cinch
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2024-12-08T06:55:18.353Z
Answers.LastActivityDate: 2024-12-08T23:09:30.767Z
Answers.LastEditDate: 2024-12-08T23:09:30.767Z
Answers.UserDisplayName: cinch
Body: <p>Does the cosine embedding loss used when pre-training DistilBERT align all hidden states output by DistilBERT with those output by BERT? Or, does it align only the hidden states of the masked tokens?</p>
  <p>Can anyone point me to a reference or DistilBERT implementation code snippet for this?</p>
  
Tags: <bert>
Title: Cosine embedding loss in DistilBERT
UserDisplayName: Fijoy Vadakkumpadan
ThreadText: # Cosine embedding loss in DistilBERT
  
  (Asked by: username_0 on 2024-12-08)
  
  Does the cosine embedding loss used when pre-training DistilBERT align all hidden states output by DistilBERT with those output by BERT? Or, does it align only the hidden states of the masked tokens?
  
  Can anyone point me to a reference or DistilBERT implementation code snippet for this?
  
  **Tags:** bert
  
  
  ---
  
  ## 1 answers
  
  ### Answer by username_1 on 2024-12-08. Score: 1
  
  During pre-training DistilBERT tries to mimic soft target distributions for *all* tokens (not just masked tokens since by transformer's self-attention nature DistilBERT still computes logits for all positions same as its teacher) of an input sequence via the *required* distillation CE loss and the additional optional cosine embedding loss and masked language modeling (MLM) loss. Therefore, like the distillation loss, the cosine loss aligns all tokens' hidden states output by DistilBERT with their corresponding tokens' hidden states output by BERT across each position by default implicitly, unless it's explicitly claimed to only focus on masked tokens as the apparent MLM loss. See the original [DistilBERT paper](https://arxiv.org/pdf/1910.01108).
  
  > To leverage the
  > inductive biases learned by larger models during pre-training, we introduce a triple loss combining language modeling, distillation and cosine-distance losses... The student is trained with a distillation loss over the soft target probabilities of the teacher... This objective results in a rich training signal by leveraging the full teacher
  > distribution... We found it beneficial to add a cosine embedding loss which will tend to align the directions of the student and teacher hidden states vectors.
  
  We can cross-check Hugging Face transformer's distilbert model's [codebase](https://github.com/huggingface/transformers/blob/main/src/transformers/models/distilbert/modeling_distilbert.py) where there're only CE/BCE/MSE losses, unfortunately. Hugging Face focuses on practical efficient implementations and they might have modified the distillation process to omit the optional cosine embedding loss described in the paper for simplicity or performance, though apparently still keeps the optional MLM CE loss. We can see only MLM loss use *labels* which are essentially masked tokens from below code snippet.
  
  > labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
  > Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
  > loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
  
  `loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))`
  
  But for the distillation loss, say for Q&A start/end position prediction class, we can see there's no masked token filtering at all.
  
  ```
              ignored_index = start_logits.size(1)
              start_positions = start_positions.clamp(0, ignored_index)
              loss_fct = nn.CrossEntropyLoss(ignore_index=ignored_index)
              start_loss = loss_fct(start_logits, start_positions)
  
  ```
  
  So I believe we can now safely assume that the optional cosine loss of DistilBERT recommended by the original paper also align hidden states across all tokens of an input sequence to have better (soft) prediction performance for all kinds of downstream NLP tasks which may need prediction beyond masked tokens, not just masked tokens.
  
  > **Comments:**
  > **username_0**: Thanks. Do you know how the distillation process ensures that [CLS] token's hidden state produced by DistilBERT aligns with that from BERT, given that the [CLS] token is never masked and there's no next sentence prediction in the distillation process? I started a separate post for this: https://ai.stackexchange.com/questions/47458/how-does-distilberts-training-ensure-that-the-cls-tokens-embedding-from-dist  
  > **username_1**: Thanks for your followup question about the important [CLS] token in BERT/DistilBERT, which prompted me to review further the Hugging Face implementation which shows the required distillation loss is applied to all tokens and similarly I believe now the optional omitted cosine loss recommended in the original paper must be also applied to all positions to further refine important token alignment such as [CLS]. I've also replied to your follow up question. Hope it completely clarifies now.  
  
  ---
  
  