Training classifier with frozen DistilBERT embeddings
Hello everyone,I’m trying to do some sentiment analysis on the IMDB movie reviews dataset. First I trained a model based on GloVe embeddings followed by an LSTM layer, and then a fully connected feedforward layer. I implemented it with Pytorch and it works like a charm. Now I’m trying to replace the GloVe + LSTM by some transformer based model. I managed to do it and I chose DistilBERT as it is supposed to be lightweight (I’m training the model on my laptop which has no GPU). I kept the DistilBERT embedding frozen, in order, again, to minimize the computational cost. But the results are really bad. It basically looks like it is oscillating around random guess.Here are a few metrics after 1 epoch:Accuracy = 0.5201
F1 score binary = 0.6723
Recall score = 0.9898
Average precision score = 0.5090
Air under the ROC = 0.7421

Loss on validation set: 1040.9420166015625Here are the same metrics after 5 epochs:Accuracy = 0.6855
F1 score binary = 0.6725
Recall score = 0.6493
Average precision score = 0.6974
Air under the ROC = 0.7497

Loss on validation set: 480.90106201171875and here are the same metrics after 11 epochs (after that I stopped):Accuracy = 0.5939
F1 score binary = 0.3630
Recall score = 0.2327
Average precision score = 0.8251
Air under the ROC = 0.7330

Loss on validation set: 738.0289306640625So not the magical bump in accuracy I was hoping forHere is also the model I used:class DistilBERTWithPoolingClf(nn.Module):
    """
    Classifier based on HuggingFace Transformers implementation of DistillBERT,
    using a basic DistillBERT layer with maxpooling on top of it.
    """
__name__ = "DistilBERTbase"

def __init__(self, keep_prob, seq_length):
    super(DistilBERTWithPoolingClf, self).__init__()
    self.DistilBERT = DistilBertModel.from_pretrained("distilbert-base-uncased")
    self.DistilBERT.requires_grad_(False)  # Embeddings are frozen
    self.maxpool = nn.MaxPool1d(seq_length)
    self.dropout = nn.Dropout(1 - keep_prob)
    self.hidden2bin = nn.Linear(768, 2)  # For Bi-LSTM

def forward(self, ids, mask, token_type_ids):
    batch_size = ids.shape[0]
    # Unlike for BERT (Hugging Face implementation), the forward method returns
    # the embedding of every input token and there is no embedding of the CLS token
    # (as far as I know)
    hidden = self.DistilBERT(ids, attention_mask=mask, return_dict=False)
    hidden = hidden[0]
    hidden = hidden.permute(0, 2, 1)
    hidden = self.maxpool(hidden)
    hidden = self.dropout(hidden)
    logits = self.hidden2bin(hidden.view(batch_size, 768))
    return logitsI have a few questions:Am I naive trying to use these models without GPU? I was hoping that by freezing them, I’ll have only to train the feedforward layer and that this part would be accessible?Could it be that the results are that bad because the DistilBERT layer is frozen?Do you see some obvious mistake in my definition of the model? I am a complete beginner when it comes to Huggingface Transformers so I wouldn’t be surprised if there was any.Any other suggestion?
2021-02-25T10:25:23Z
Hi@abercher, regarding your questions:You can certainly use Transformers as feature extractors on a CPU - since the weights are frozen, you only need the forward pass which is relatively quick to computeMy experience has generally been that youcanget significantly worse results when using the last hidden states as features vs fine-tuning end-to-end (in some cases > 20 F1 points!). But this varies depending on the dataset / task, so I am not sure if it’s also true for IMDB.One thing that seems a bit odd is thehidden.permute(0, 2, 1)part of your forward pass - why do you do this?You might get better results by using the average of the unmasked hidden state, e.g. innumpycode:input_ids = torch.tensor(batch["input_ids"]).to(device)
  attention_mask = torch.tensor(batch["attention_mask"]).to(device)
  with torch.no_grad():
    last_hidden_state = model(input_ids, attention_mask).last_hidden_state
    last_hidden_state = last_hidden_state.cpu().numpy()
  # Use average of unmasked hidden states for classification
  lhs_shape = last_hidden_state.shape
  boolean_mask = ~np.array(batch["attention_mask"]).astype(bool)
  boolean_mask = np.repeat(boolean_mask, lhs_shape[-1], axis=-1)
  boolean_mask = boolean_mask.reshape(lhs_shape)
  masked_mean = np.ma.array(last_hidden_state, mask=boolean_mask).mean(axis=1)
  batch["hidden_state"] = masked_mean.dataYou could also see what the performance looks like with something simpler like logistic regression, e.g. first extract the features:def forward_pass(batch):
  input_ids = torch.tensor(batch["input_ids"]).to(device)
  attention_mask = torch.tensor(batch["attention_mask"]).to(device)
  with torch.no_grad():
    last_hidden_state = model(input_ids, attention_mask, return_dict=True).last_hidden_state
    last_hidden_state = last_hidden_state.cpu().numpy()
  # Use average of unmasked hidden states for classification
  lhs_shape = last_hidden_state.shape
  boolean_mask = ~np.array(batch["attention_mask"]).astype(bool)
  boolean_mask = np.repeat(boolean_mask, lhs_shape[-1], axis=-1)
  boolean_mask = boolean_mask.reshape(lhs_shape)
  masked_mean = np.ma.array(last_hidden_state, mask=boolean_mask).mean(axis=1)
  batch["hidden_state"] = masked_mean.data
  return batchand then extract the hidden states from your tokenized dataset, e.g.imdb_enc:imdb_enc = imdb_enc.map(forward_pass, batched=True, batch_size=16)
X_train = np.array(imdb_enc["train"]["hidden_state"])
X_test = np.array(imdb_enc["test"]["hidden_state"])
y_train = np.array(imdb_enc["train"]["label"])
y_test= np.array(imdb_enc["test"]["label"])and then train a classifierfrom sklearn.linear_model import LogisticRegression
lr_clf = LogisticRegression(n_jobs=-1, penalty="none")
lr_clf.fit(X_train, y_train)
lr_clf.score(X_valid, y_valid)If the result is still bad then you might have to try something more elaborate like averaging over certain layers as was done in the BERT paper.HTH!
2021-02-25T11:35:25Z
Hi lewtun,Thank you so much for your answer! It’s very kind of you to take the time to write a detailed answer. I highly appreciate your help.I’m working now, but I’ll look at your answer in detail this evening when I’m done with my job.Thanks again
2021-02-25T20:40:13Z
Hi@abercher, curiosity got the better of me so I ran the above suggestion on IMDB and got an accuracy of 86.6% on the test setYou can find the hacky codehereand it shouldn’t be too hard to adapt to your MLP classifier.
2021-02-25T22:20:13Z
Hi again@lewtun,Thanks for your second answer. I’m looking forward to trying it. Probably this weekend because today was too tight for me to do it.I looked at the code concerning this swapping of dimension inhidden = self.DistilBERT(ids, attention_mask=mask, return_dict=False)
        hidden = hidden[0]
        hidden = hidden.permute(0, 2, 1)
        hidden = self.maxpool(hidden)and the reason is that the DistilBERT layer outputs a tensor of dimension (n_batch, input_sequence_length, 768) where 768 is just the size of individual token embeddings as DistilBERT outputs. As the nn.MaxPool1d layer takes the max over the last dimension, I swapped the dimension because I want the max pooling to be done on the “temporal” axis so in other words take the max over the embeddings obtained for the different input tokens.Concerning this averaging of unmasked hidden states, I’m a bit confused. I know what masking is, at least the general idea, but I don’t understand why we wouldn’t want to have the embeddings of the (randomly) masked tokens into our average.Also, is averaging a more common practice than taking the max for these kinds of models (the transformer based ones)?Anyway, thanks a lot for your support. It was really great. I’ll try all this and come back to you. Sorry for the delay.
2025-01-24T22:15:45Z
hi@lewtun, i was just going through the code you provided. first of all it really helped me get the grasp of the embeddings. i am basically an AI student working on a project and i was intending to use BERT embeddings to train an LSTM model. it would be a great help if you could let me know that whether this technique of extracting embeddings and then averaging it can work with LSTMS as i was reading through other blogs and what i got to know was that lstms typically expects embeddings in this format (batch_size, num_timesteps, num_features) but this method you mentioned returns the embeddings in the shape (batch_size, num_features). is it possible to train an LSTM with this technique? for further clarification i am intending to train the model for sentiment analysis. it would really mean alot if you can take out the time and reply to me. thanks in advance.