¿Cómo ajustar finamente BERT para una tarea de clasificación? ¿Qué tomar de la salida?

Para el ajuste fino de BERT para tareas de clasificación, generalmente se toma la salida del token especial [CLS], que la modelo utiliza para agregar información sobre toda la secuencia.

Pasos principales:

1. Agregar una capa de clasificación sobre BERT — generalmente una capa lineal que recibe el vector de la capa de salida de BERT correspondiente al token [CLS].

2. Los datos de entrada se ingresan en BERT con el token [CLS] al principio.

3. Extraer la incrustación [CLS] de la última capa oculta de BERT. Este vector, por ejemplo, de tamaño 768 (para el modelo base), contiene la información resumida de toda la secuencia.

4. Pasar este vector al clasificador (por ejemplo, linear + softmax para clasificación multiclase).

5. Entrenar todo el modelo (BERT + clasificador) conjuntamente en la tarea objetivo, usando una función de pérdida adecuada (por ejemplo, CrossEntropyLoss).

Ejemplo 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

Por lo tanto, para clasificación, se toma específicamente la salida del token [CLS] del último capa de BERT.