I need to finetune BERT model (from the huggingface repository) on a sentence classification task. However, my dataset is really small.I have 12K sentences and only 10% of them are from positive classes. Does anyone here have any experience on finetuning bert in small datasets? Do you have any suggestions for learning rate, batch size, epoch number warmup steps etc.?


2 Answers 2


Data augumentation

If the number of text data is small, text data argumentations may be applicable e.g. nlpaug. Applying text summarization, removing stopwords or punctuations would be a simple way to create variations of data.

Learning Rate

How to Fine-Tune BERT for Text Classification? pointed out the learning rate is the key to avoid Catastrophic Forgetting where the pre-trained knowledge is erased during learning of new knowledge.

We find that a lower learning rate, such as 2e-5, is necessary to make BERT overcome the catastrophic forgetting problem. With an aggressive learn rate of 4e-4, the training set fails to converge.
enter image description here

Probably this is the reason why the BERT paper used 5e-5, 4e-5, 3e-5, and 2e-5 for fine-tuning.

We use a batch size of 32 and fine-tune for 3 epochs over the data for all GLUE tasks. For each task, we selected the best fine-tuning learning rate (among 5e-5, 4e-5, 3e-5, and 2e-5) on the Dev set

Note that the base model pre-training itself used higher learning rate.

The model was trained on 4 cloud TPUs in Pod configuration (16 TPU chips total) for one million steps with a batch size of 256. The sequence length was limited to 128 tokens for 90% of the steps and 512 for the remaining 10%. The optimizer used is Adam with a learning rate of 1e-4, β1=0.9 and β2=0.999, a weight decay of 0.01, learning rate warmup for 10,000 steps and linear decay of the learning rate after.


The number of epochs would be fairly small. The original paper fine-tuning experiments indicated the amount of time/epochs required were small e.g. 3 epochs for GLUE tasks. In my personal experience, 5 is more or less sufficient for text classification tasks, although it will depends on your data. The number of epochs as well would depends on the metrics to monitor.

Batch size

The original paper used 32 for fine tuning but it depends on the maximum sequence length too.

We use a batch size of 32 and fine-tune for 3 epochs over the data for all GLUE tasks.

Each word is encoded into a floating point vector of size 768 and there are 12 layers for the BERT/base. If the max 512 length is used, the data may not fit into GPU memory with the batch size 32. Then reduce to 16. If the max length is 128 or 256, then 32 would be a good number. Please check the available GPU memory.

Model for fine tuning

You can add multiple classification layers on top of the BERT base model but the original paper indicates only one output layer to convert 768 outputs into the number of labels you have, and apparently it is the way widely used when fine-tuning is done on BERT.

As a result, the pre-trained BERT model can be fine tuned with just one additional output layer to create state-of-the-art models for a wide range of tasks, such as question answering and language inference, without substantial task specific architecture modifications.

All the parameters in the BERT model will be fine-tuned, but you can try freezing the base model and add more classification layers on top of the BERT base model.

For each task, we simply plug in the task specific inputs and outputs into BERT and finetune all the parameters end-to-end.


The original paper also used Adam with weight decay. Huggingface provides AdamWeightDecay (TensorFlow) or AdamW (PyTorch). Keep using the same optimizer would be sensible although different ones can be tried. The default learning rate is set to the value used at pre-training. Hence need to set to the value for fine-tuning.



How many classes do you have? Bert can handle a high-quality 12k dataset for binary classification. I recommend duplicating your positive test case 4x and sampling a 5k test cases from your negative class. This will give you a balanced dataset. Then implement BERT in google colab using the original GitHub repository from the google Bert team. You can find a functional collab here: https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/bert_finetuning_with_cloud_tpus.ipynb On collab, TPUs takes about 10 to 15 minutes to run. All you have to do is use write your data processor. The following parameters are a good place to start.

# Warmup is a period of time where hte learning rate 
# is small and gradually increases--usually helps training.
# Model configs

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.