# Training classifier with frozen DistilBERT embeddings
**URL:**
**Author:**  [@abercher](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fabercher)%5C
**Post date:** [February 25, 2021, 9:21am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ftraining-classifier-with-frozen-distilbert-embeddings%2F4000%2F1 "2021-02-25T09:21:22Z")
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.9420166015625
```
Here 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.90106201171875
```
and 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.0289306640625
```
So not the magical bump in accuracy I was hoping for 😅
Here 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 logits
```
I have a few questions:
1. 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?
2. Could it be that the results are that bad because the DistilBERT layer is frozen?
3. 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.
4. Any other suggestion?
---
**Author:**  [@lewtun](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Flewtun)%5C
**Post date:** [February 25, 2021, 10:25am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ftraining-classifier-with-frozen-distilbert-embeddings%2F4000%2F2 "2021-02-25T10:25:23Z")
Hi @abercher, regarding your questions:
1. 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 compute
2. My experience has generally been that you _can_ get 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.
3. One thing that seems a bit odd is the `hidden.permute(0, 2, 1)` part of your forward pass - why do you do this?
4. You might get better results by using the average of the unmasked hidden state, e.g. in `numpy` code:
```auto
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.data
```
You could also see what the performance looks like with something simpler like logistic regression, e.g. first extract the features:
```auto
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 batch
```
and then extract the hidden states from your tokenized dataset, e.g. `imdb_enc`:
```auto
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 classifier
```auto
from 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!
---
**Author:**  [@abercher](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fabercher)%5C
**Post date:** [February 25, 2021, 11:35am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ftraining-classifier-with-frozen-distilbert-embeddings%2F4000%2F3 "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
---
**Author:**  [@lewtun](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Flewtun)%5C
**Post date:** [February 25, 2021, 8:40pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ftraining-classifier-with-frozen-distilbert-embeddings%2F4000%2F4 "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 set 🙂
You can find the hacky code [here](https://colab.research.google.com/drive/1VJlA-tdqGvXRelCu-PkHzbGO95A7phY6?usp=sharing) and it shouldn’t be too hard to adapt to your MLP classifier.
---
**Author:**  [@abercher](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fabercher)%5C
**Post date:** [February 25, 2021, 10:20pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ftraining-classifier-with-frozen-distilbert-embeddings%2F4000%2F5 "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 in
```
hidden = 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.
---
**Author:**  [@faraz4020](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Ffaraz4020)%5C
**Post date:** [January 24, 2025, 10:15pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ftraining-classifier-with-frozen-distilbert-embeddings%2F4000%2F6 "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.