Comment affiner BERT pour une tâche de classification ? Que prendre de la sortie ?

Pour le fine-tuning de BERT pour une tâche de classification, on prend généralement la sortie du token spécial [CLS], qui est utilisé par le modèle pour agréger l'information de toute la séquence.

Étapes principales:

1. Ajouter une couche de classification au-dessus de BERT — généralement une couche linéaire qui reçoit le vecteur de la couche de sortie de BERT correspondant au token [CLS].

2. Les données d'entrée sont passées à BERT avec le token [CLS] au début.

3. Extraire l'embedding [CLS] de la dernière couche cachée de BERT. Ce vecteur, par exemple, de taille 768 (pour le modèle de base), contient l'information condensée de toute la séquence.

4. Passer ce vecteur au classificateur (par exemple, linear + softmax pour une classification multiclasse).

5. Entraîner tout le modèle (BERT + classificateur) conjointement sur la tâche cible, en utilisant une fonction de perte appropriée (par exemple, CrossEntropyLoss).

Exemple en PyTorch:

from transformers import BertModel, BertTokenizer
import torch.nn as nn

class BertClassifier(nn.Module):
    def __init__(self, bert_model_name, num_classes):
        super().__init__()
        self.bert = BertModel.from_pretrained(bert_model_name)
        self.classifier = nn.Linear(self.bert.config.hidden_size, num_classes)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        cls_output = outputs.last_hidden_state[:, 0, :]  # token [CLS]
        logits = self.classifier(cls_output)
        return logits

Ainsi, pour la classification, on prend précisément la sortie du token [CLS] du dernier couche de BERT.