Reinit layers of pretrained transformer in cached_transformers.get()

allennlp/modules/token_embedders/pretrained_transformer_embedder.py

@@ -1,6 +1,6 @@
 import logging
 import math
-from typing import Optional, Tuple, Dict, Any
+from typing import Optional, Tuple, Dict, Any, Union, List, cast
 
 
 import torch
@@ -48,6 +48,12 @@ class PretrainedTransformerEmbedder(TokenEmbedder):
         When `True` (the default), only the final layer of the pretrained transformer is taken
         for the embeddings. But if set to `False`, a scalar mix of all of the layers
         is used.
+    reinit_layers: `Optional[Union[int, List[int]]]`, optional (default = `None`)
+        If this is an integer, the last `reinit_layers` layers of the transformer will be
+        re-initialized. If this is a list, the layers indexed by `reinit_layers` will be
+        re-initialized. Re-initializing the last few layers of a pretrained transformer can reduce
+        the instability of fine-tuning on small datasets and may improve performance
+        (https://arxiv.org/abs/2006.05987v3). Has no effect if `load_weights` is `False`.
     override_weights_file: `Optional[str]`, optional (default = `None`)
         If set, this specifies a file from which to load alternate weights that override the
         weights from huggingface. The file is expected to contain a PyTorch `state_dict`, created
@@ -82,6 +88,7 @@ def __init__(
         train_parameters: bool = True,
         eval_mode: bool = False,
         last_layer_only: bool = True,
+        reinit_layers: Optional[Union[int, List[int]]] = None,
         override_weights_file: Optional[str] = None,
         override_weights_strip_prefix: Optional[str] = None,
         load_weights: bool = True,
@@ -119,6 +126,22 @@ def __init__(
             self._scalar_mix = ScalarMix(self.config.num_hidden_layers)
             self.config.output_hidden_states = True
 
+        # Optionally, re-initialize the parameters of certain layers.
+        self._reinit_layers = cast(List[int], reinit_layers)
+        if self._reinit_layers and load_weights:
+            num_layers = len(self.transformer_model.encoder.layer)
+            if isinstance(reinit_layers, int):
+                self._reinit_layers = list(range(num_layers - reinit_layers, num_layers))
+            if any(layer_idx < 0 or layer_idx > num_layers for layer_idx in self._reinit_layers):
+                raise ValueError(
+                    f"A layer index in reinit_layers ({self._reinit_layers}) is invalid. Must be"
+                    f" between 0 and the maximum layer index ({num_layers - 1}.)"
+                )
+            for layer_idx in self._reinit_layers:
+                self.transformer_model.encoder.layer[layer_idx].apply(
+                    self.transformer_model._init_weights
+                )
+
         tokenizer = PretrainedTransformerTokenizer(
             model_name,
             tokenizer_kwargs=tokenizer_kwargs,

tests/modules/token_embedders/pretrained_transformer_embedder_test.py

@@ -341,6 +341,41 @@ def test_embeddings_resize(self):
             == 28997
         )
 
+    def test_reinit_layers(self):
+        regular_token_embedder = PretrainedTransformerEmbedder("bert-base-cased")
+        assert regular_token_embedder._reinit_layers is None
+        # Test the case when reinit_layers is a valid int. Comparing all weights of the model is
+        # rather complicated, so arbitrarily compare the weights of attention module.
+        preinit_weights = torch.cat(
+            [
+                layer.attention.output.dense.weight
+                for layer in regular_token_embedder.transformer_model.encoder.layer
+            ]
+        )
+        reinit_token_embedder = PretrainedTransformerEmbedder("bert-base-cased", reinit_layers=2)
+        postinit_weights = torch.cat(
+            [
+                layer.attention.output.dense.weight
+                for layer in reinit_token_embedder.transformer_model.encoder.layer
+            ]
+        )
+        assert reinit_token_embedder._reinit_layers == [10, 11]
+        assert torch.equal(postinit_weights[:10], preinit_weights[:10])
+        assert not torch.equal(postinit_weights[10:], preinit_weights[10:])
+        # Test the case when reinit_layers is a valid list of integers.
+        reinit_token_embedder = PretrainedTransformerEmbedder(
+            "bert-base-cased", reinit_layers=[10, 11]
+        )
+        assert reinit_token_embedder._reinit_layers == [10, 11]
+        assert torch.equal(postinit_weights[:10], preinit_weights[:10])
+        assert not torch.equal(postinit_weights[10:], preinit_weights[10:])
+        # Should raise a ValueError because reinit_layers contains at least one index that is
+        # greater than the models maximum number of layers
+        with pytest.raises(ValueError):
+            _ = PretrainedTransformerEmbedder("bert-base-cased", reinit_layers=1000)
+        with pytest.raises(ValueError):
+            _ = PretrainedTransformerEmbedder("bert-base-cased", reinit_layers=[1, 1000])
+
     def test_eval_mode(self):
         token_embedder = PretrainedTransformerEmbedder("epwalsh/bert-xsmall-dummy", eval_mode=True)
         assert token_embedder.training and not token_embedder.transformer_model.training