Since the advent of the Transformer architecture (Vaswani et al. 2017) and of BERT models (Devlin et al. 2019), Transformer models have become ubiquitous in NLP, achieving SOTA results on most NLP datasets.

Before Sesame Street puppets flooded on ArXiv, the de-facto method to train an NLP model leveraged word embeddings pre-trained using Glove or word2vec. These word embeddings were used to initialize the first embedding layer of your model, and you just had to plug the rest of your architecture above this first layer.

Nowadays, it’s increasingly common to use BERT-like language models trained on vast amounts of unlabeled text, and fine-tune the model on task-specific data. This change of paradigm led to a new subfield in NLP called BERTology which studies the ins and outs of BERT-like models.

Many hyper-parameters of BERT — such as the number of attention heads or the learning rate to use — have been studied, but few papers looked at how layer freezing affects performances. Freezing layers means disabling gradient computation and backpropagation for the weights of these layers. This is a common technique in NLP or Computer Vision when working on small datasets as it usually reduces the chances of your model to overfit. Usually, models are frozen from bottom to top, as the first layers of neural networks are thought to encode universal features while the top layers are more task-specific. Another advantage of freezing is that it provides a nice speed boost and reduced memory usage during training.

Most machine learning teams in industry don’t have access to high-quality and domain-specific datasets. As the costs of building custom datasets are high, these datasets are usually way smaller than those often used in academia. Learning how to train better models with fewer data is a major concern, especially in industry. While the effect of freezing on performance has been explored (Lee, Tang, and Lin 2019), the low-resource setting hasn’t attracted much attention yet. Luckily, it’s what this post is all about 😎

📝 Methods

To test how freezing affects performances when data is scarce, we will use the now popular 🤗 Transformers library from Hugging Face to train a RoBERTa model (base version) on two standard binary classification datasets from the GLUE benchmark: SST-2 (Stanford Sentiment Treebank) and CoLA (Corpus of Linguistic Acceptability).

SST-2 is a sentiment analysis task and CoLA a dataset containing sentences labeled with respect to their grammatical acceptability (i.e whether each sentence is grammatically correct or not). SST-2 and CoLA contain ~51% and ~69% of positive samples respectively in their validation set, which gives us a majority class baseline to compare to. We will use accuracy as the metric of evaluation for both datasets.

We will train models with a subset of the training dataset and evaluate on the whole validation dataset. For CoLA and SST-2, we generate subsets of size 50, 100, 200, 500, and 1000, by:

• taking care that the fraction of positive samples is the same as in the original train dataset (stratify splitting)
• generating the splits by increasing size (50, 100, 200, 500, 1000) and reusing all samples from the previous split. By doing so, we limit the variability coming from using different samples in each split.

To assess the impact of layer freezing on performances, 0, 2, 4, 8, 10, or 12 encoder layers were frozen. I also tested to only freeze the embedding layers (frozen layers = -1).

Similar to the original RoBERTa paper (Liu et al. 2019), I fine-tuned for 10 epochs with a linear warmup of 1 epoch followed by a linear decay to zero. As my computational budget was limited, I didn’t perform any hyper-parameter search and stuck to the default values (lr = 5e-5, AdamW optimizer).

On both datasets, for each layer freezing value (-1, 0, 2, 4, 8, 10, 12) and for each train size (50, 100, 200, 500, 1000), 5 replicates were trained with a different random seed.

I decided to use the recent (and pretty awesome!) 🤗Datasets library — also developed by folks from Hugging Face — to easily fetch and process these two datasets:

DATASET_MAP = {"sst2": ("glue", "sst2"), "cola": ("glue", "cola")}

def get_dataset(tokenizer, dataset_name: str, split: str, split_path: Optional[Path] = None):
    ds = datasets.load_dataset(*DATASET_MAP[dataset_name], split=split)
    ds = ds.shuffle(seed=42)

    if split_path is not None:
	      # split_path is a npy file containing indexes of samples to keep
        split_ids = set(np.load(split_path).tolist())
        ds = ds.filter(lambda idx: idx in split_ids, input_columns="idx")

    ds = ds.map(lambda e: tokenizer(e["sentence"], padding=False, truncation=True), batched=True)
    ds.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
    return ds

The get_dataset function returns a Dataset object that integrates seamlessly with the Transformers library. The sentences are tokenized (and cached) on the fly using Dataset.map.

The training function is pretty simple, thanks to the use of the Trainer class that takes care of most of the boilerplate code for us:

def get_random_seed():
    return int.from_bytes(os.urandom(4), "big")

def train(
    output_dir: str,
    dataset_name: str,
    train_size: Optional[int] = None,
    freeze_layer_count: int = 0,
):
    # get_split_path returns the path of a npy file containing sample indexes
    # to include in the training set
    split_path = get_split_path(dataset_name, train_size) if train_size is not None else None
    args_dict = {
        "evaluation_strategy": "steps",
        "per_device_train_batch_size": 16,
        "per_device_eval_batch_size": 16,
        "learning_rate": 5e-5,
        "num_train_epochs": 10,
        "logging_first_step": True,
        "save_total_limit": 1,
        "fp16": True,
        "dataloader_num_workers": 1,
        "load_best_model_at_end": True,
        "metric_for_best_model": "accuracy",
        # we need to generate a random seed manually, as otherwise 
        # the same constant random seed is used during training for each run
        "seed": get_random_seed(),
    }
    model = AutoModelForSequenceClassification.from_pretrained("roberta-base", return_dict=True)

    if freeze_layer_count:
	      # We freeze here the embeddings of the model
        for param in model.roberta.embeddings.parameters():
            param.requires_grad = False

        if freeze_layer_count != -1:
	          # if freeze_layer_count == -1, we only freeze the embedding layer
	          # otherwise we freeze the first `freeze_layer_count` encoder layers
            for layer in model.roberta.encoder.layer[:freeze_layer_count]:
                for param in layer.parameters():
                    param.requires_grad = False

    tokenizer = AutoTokenizer.from_pretrained("roberta-base")
    train_ds = get_dataset(tokenizer, dataset_name, split="train", split_path=split_path)
    val_ds = get_dataset(tokenizer, dataset_name, split="validation")

    epoch_steps = len(train_ds) / args_dict["per_device_train_batch_size"]
    args_dict["warmup_steps"] = math.ceil(epoch_steps)  # 1 epoch
    args_dict["logging_steps"] = max(1, math.ceil(epoch_steps * 0.5))  # 0.5 epoch
    args_dict["save_steps"] = args_dict["logging_steps"]

    training_args = TrainingArguments(output_dir=str(output_dir), **args_dict)
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_ds,
        eval_dataset=val_ds,
        compute_metrics=compute_metrics,
        tokenizer=tokenizer,
    )
    trainer.train()

The complete code is also available as a gist.

🔎 Results

Here comes the most interesting part!

Below is displayed for both datasets the best accuracy obtained during training on the development set for all train sizes and all numbers of frozen layers. Each red dot on the plot represents a single run.

The complete results can also be found as tables at the end of the post.

SST-2: Best accuracy during training on development set

CoLA: Best accuracy during training on development set

As expected, having more data always help: we gain +0.15 (0.75 → 0.90) in accuracy by increasing dataset size from 50 to 1000 for SST-2.

Freezing all (12) layers of the encoder is definitely not a good idea when dealing with limited data: there is always a significant performance drop, and the resulting models are barely better than a majority class classifier. Freezing 10 layers is also associated with degraded accuracy, even if the effect isn’t as strong. Aside from 10- and 12-frozen layers, there are no clear impact of layer freezing on performances.

We also observe high variability in accuracy between runs with the same settings (but with different random seeds). This instability is known since the release of BERT. While catastrophic forgetting and the small sizes of datasets were first suspected as the causes of this instability, more recent work (Mosbach et al. 2020) suggests that optimization difficulties leading to vanishing gradients are the actual reasons. For both SST-2 and CoLA datasets, there are no clear evidence that freezing reduces these instabilities.

As a final note, the performance gap between SST-2 and CoLA with 50 examples is interesting: while best models for SST-2 classify 82% of samples correctly (with 51% of positive samples), best run are barely better than the majority class baseline with 70% of samples correctly classified. This is indicative of the inductive bias of these models, which are more prone to detect polarity than grammatical correctness.

Conclusion

And that’s all for this second post 😊!

The results of these experiments were surprising to me, I initially thought layer freezing would have a more visible impact on performances.

I hope you’ve enjoyed the post, feel free to share your feedbacks or your experience with layer freezing (my email address is available on the blog).

Bibliography

Devlin, Jacob, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. “BERT: Pre-Training of Deep Bidirectional Transformers for Language Understanding.” ArXiv:1810.04805 [Cs], May. http://arxiv.org/abs/1810.04805.

Lee, Jaejun, Raphael Tang, and Jimmy Lin. 2019. “What Would Elsa Do? Freezing Layers During Transformer Fine-Tuning.” ArXiv:1911.03090 [Cs], November. http://arxiv.org/abs/1911.03090.

Liu, Yinhan, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019. “RoBERTa: A Robustly Optimized BERT Pretraining Approach.” ArXiv:1907.11692 [Cs], July. http://arxiv.org/abs/1907.11692.

Mosbach, Marius, Maksym Andriushchenko, and Dietrich Klakow. 2020. “On the Stability of Fine-Tuning BERT: Misconceptions, Explanations, and Strong Baselines.” ArXiv:2006.04884 [Cs, Stat], October. http://arxiv.org/abs/2006.04884.

Vaswani, Ashish, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. “Attention Is All You Need.” ArXiv:1706.03762 [Cs], December. http://arxiv.org/abs/1706.03762.

Appendix

SST-2

CoLA