Skip to content

Stabilize and achieve excellent performance with transformers

License

Notifications You must be signed in to change notification settings

flowerpot-ai/stabilizer

Repository files navigation

Stabilizer

Stabilize and achieve excellent performance with transformers.

The stabilizer library offer solutions to tackle one of the biggest challenges that comes along with training State of the art Transformer models, Unstable training

Unstable training

Unstable training is the phenomenon in which training large transformer models with trivial changes such as changing the random seed drastically changes the performance of the model. Here is a screenshot of finetuning the CoLA dataset from GLUE tasks with two different random seeds applied only to the dropout of the transformer model.

dropout_random_seed

Installation

pip install stabilizer

Techniques currently implemented in this library

  1. Reinitialization
  2. Layerwise Learning Rate Decay

Reinitialization

Reinitialize the last n layers of the transformer encoder. This technique works well because we reinitialize the task specific parameters that the pretrained models have learnt specific to the pretraining task.

from stabilizer.reinitialize import reinit_autoencoder_model
from transformers import AutoModel

transformer = AutoModel.from_pretrained(
    pretrained_model_name_or_path="bert-base-uncased",
    hidden_dropout_prob=0.1,
    attention_probs_dropout_prob=0.1,
)
transformer.encoder = reinit_autoencoder_model(
    transformer.encoder, reinit_num_layers=1
)

Here is the result of the same model but reinitialized last 4 layers applied on the CoLA dataset. You can see that the model has converged to almost the same performance with reinitialization. reinit_random_seed

Layerwise Learning Rate Decay

Apply layerwise learning rate to the transformer layers. Starting from the task specific layer every layer before it gets an exponentially decreasing learning rate.

from stabilizer.llrd import get_optimizer_parameters_with_llrd
from stabilizer.model import PoolerClassifier

from transformers import AdamW, AutoModel


    transformer = AutoModel.from_pretrained(
        pretrained_model_name_or_path=config["pretrained_model_name_or_path"],
        hidden_dropout_prob=config["dropout_prob"],
        attention_probs_dropout_prob=config["dropout_prob"],
    )
    
    model = PoolerClassifier(
        transformer=transformer,
        transformer_output_size=transformer.config.hidden_size,
        transformer_output_dropout_prob=config["dropout_prob"],
        num_classes=config["num_classes"],
        task_specific_layer_seed=config["layer_initialization_seed"],
    )

    model_parameters = get_optimizer_parameters_with_llrd(
        model=model,
        peak_lr=config["lr"],
        multiplicative_factor=config["multiplicative_factor"],
    )
    optimizer = AdamW(params=model_parameters, lr=config["lr"])

Here is the result of the same model but with LLRD applied on the CoLA dataset. Here you can see that the model has diverged quite a lot by applying LLRD. Therefore as we discussed earlier their is no universal remedy yet but some techniques work well on some datasets llrd_random_seed

Conbtributing to this repository

There are three possible ways by which people can contribute to this repository:

  1. Adding more techniques for stabilizing the training
  2. Write example scripts to train new datasets with the already implemented stabilization techniques
  3. Adding tensorflow/jax equivalent codes for the existing Pytorch code

About

Stabilize and achieve excellent performance with transformers

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages