|
| 1 | +# Adaptive Layers |
| 2 | + |
| 3 | +Embedding models are often encoder models with numerous layers, such as 12 (e.g. [all-mpnet-base-v2](https://huggingface.co/sentence-transformers/all-mpnet-base-v2)) or 6 (e.g. [all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2)). To get embeddings, every single one of these layers must be traversed. [2D Matryoshka Sentence Embeddings](https://arxiv.org/abs/2402.14776) (2DMSE) revisits this concept by proposing an approach to train embedding models that will perform well when only using a selection of all layers. This results in faster inference speeds at relatively low performance costs. |
| 4 | + |
| 5 | +## Use Cases |
| 6 | + |
| 7 | +The 2DMSE paper mentions that using a few layers of a larger model trained using Adaptive Layers and Matryoshka Representation Learning can outperform a smaller model that was trained like a standard embedding model. |
| 8 | + |
| 9 | +## Results |
| 10 | + |
| 11 | +Let's look at the performance that we may be able to expect from an Adaptive Layer embedding model versus a regular embedding model. For this experiment, I have trained two models: |
| 12 | + |
| 13 | +* [tomaarsen/mpnet-base-nli-adaptive-layer](https://huggingface.co/tomaarsen/mpnet-base-nli-adaptive-layer): Trained by running [adaptive_layer_nli.py](adaptive_layer_nli.py) with [microsoft/mpnet-base](https://huggingface.co/microsoft/mpnet-base). |
| 14 | +* [tomaarsen/mpnet-base-nli](https://huggingface.co/tomaarsen/mpnet-base-nli): A near identical model as the former, but using only `MultipleNegativesRankingLoss` rather than `AdaptiveLayerLoss` on top of `MultipleNegativesRankingLoss`. I also use [microsoft/mpnet-base](https://huggingface.co/microsoft/mpnet-base) as the base model. |
| 15 | + |
| 16 | +Both of these models were trained on the AllNLI dataset, which is a concatenation of the [SNLI](https://huggingface.co/datasets/snli) and [MultiNLI](https://huggingface.co/datasets/multi_nli) datasets. I have evaluated these models on the [STSBenchmark](https://huggingface.co/datasets/mteb/stsbenchmark-sts) test set using multiple different embedding dimensions. The results are plotted in the following figure: |
| 17 | + |
| 18 | + |
| 19 | + |
| 20 | +The first figure shows that the Adaptive Layer model stays much more performant when reducing the number of layers in the model. This is also clearly shown in the second figure, which displays that 80% of the performance is preserved when the number of layers is reduced all the way to 1. |
| 21 | + |
| 22 | +Lastly, the third figure shows the expected speedup ratio for GPU & CPU devices in my tests. As you can see, removing half of the layers results in roughly a 2x speedup, at a cost of ~15% performance on STSB (~86 -> ~75 Spearman correlation). When removing even more layers, the performance benefit gets larger for CPUs, and between 5x and 10x speedups are very feasible with a 20% loss in performance. |
| 23 | + |
| 24 | +## Training |
| 25 | + |
| 26 | +Training with Adaptive Layer support is quite elementary: rather than applying some loss function on only the last layer, we also apply that same loss function on the pooled embeddings from previous layers. Additionally, we employ a KL-divergence loss that aims to make the embeddings of the non-last layers match that of the last layer. This can be seen as a fascinating approach of [knowledge distillation](../distillation/README.html#knowledge-distillation), but with the last layer as the teacher model and the prior layers as the student models. |
| 27 | + |
| 28 | +For example, with the 12-layer [microsoft/mpnet-base](https://huggingface.co/microsoft/mpnet-base), it will now be trained such that the model produces meaningful embeddings after each of the 12 layers. |
| 29 | + |
| 30 | +```python |
| 31 | +from sentence_transformers import SentenceTransformer |
| 32 | +from sentence_transformers.losses import CoSENTLoss, AdaptiveLayerLoss |
| 33 | + |
| 34 | +model = SentenceTransformer("microsoft/mpnet-base") |
| 35 | + |
| 36 | +base_loss = CoSENTLoss(model=model) |
| 37 | +loss = AdaptiveLayerLoss(model=model, loss=base_loss) |
| 38 | +``` |
| 39 | +* **Reference**: <a href="../../../docs/package_reference/losses.html#adaptivelayerloss"><code>AdaptiveLayerLoss</code></a> |
| 40 | + |
| 41 | +Note that training with `AdaptiveLayerLoss` is not notably slower than without using it. |
| 42 | + |
| 43 | +Additionally, this can be combined with the `MatryoshkaLoss` such that the resulting model can be reduced both in the number of layers, but also in the size of the output dimensions. See also the [Matryoshka Embeddings](../matryoshka/README.html) for more information on reducing output dimensions. In Sentence Transformers, the combination of these two losses is called `Matryoshka2dLoss`, and a shorthand is provided for simpler training. |
| 44 | + |
| 45 | +```python |
| 46 | +from sentence_transformers import SentenceTransformer |
| 47 | +from sentence_transformers.losses import CoSENTLoss, Matryoshka2dLoss |
| 48 | + |
| 49 | +model = SentenceTransformer("microsoft/mpnet-base") |
| 50 | + |
| 51 | +base_loss = CoSENTLoss(model=model) |
| 52 | +loss = Matryoshka2dLoss(model=model, loss=base_loss, matryoshka_dims=[768, 512, 256, 128, 64]) |
| 53 | +``` |
| 54 | + |
| 55 | +* **Reference**: <a href="../../../docs/package_reference/losses.html#matryoshka2dloss"><code>Matryoshka2dLoss</code></a> |
| 56 | + |
| 57 | +## Inference |
| 58 | + |
| 59 | +After a model has been trained using the Adaptive Layer loss, you can then truncate the model layers to your desired layer count. Note that this requires doing a bit of surgery on the model itself, and each model is structured a bit differently, so the steps are slightly different depending on the model. |
| 60 | + |
| 61 | +First of all, we will load the model & access the underlying `transformers` model like so: |
| 62 | + |
| 63 | +```python |
| 64 | +from sentence_transformers import SentenceTransformer |
| 65 | + |
| 66 | +model = SentenceTransformer("tomaarsen/mpnet-base-nli-adaptive-layer") |
| 67 | + |
| 68 | +# We can access the underlying model with `model[0].auto_model` |
| 69 | +print(model[0].auto_model) |
| 70 | +``` |
| 71 | +``` |
| 72 | +MPNetModel( |
| 73 | + (embeddings): MPNetEmbeddings( |
| 74 | + (word_embeddings): Embedding(30527, 768, padding_idx=1) |
| 75 | + (position_embeddings): Embedding(514, 768, padding_idx=1) |
| 76 | + (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True) |
| 77 | + (dropout): Dropout(p=0.1, inplace=False) |
| 78 | + ) |
| 79 | + (encoder): MPNetEncoder( |
| 80 | + (layer): ModuleList( |
| 81 | + (0-11): 12 x MPNetLayer( |
| 82 | + (attention): MPNetAttention( |
| 83 | + (attn): MPNetSelfAttention( |
| 84 | + (q): Linear(in_features=768, out_features=768, bias=True) |
| 85 | + (k): Linear(in_features=768, out_features=768, bias=True) |
| 86 | + (v): Linear(in_features=768, out_features=768, bias=True) |
| 87 | + (o): Linear(in_features=768, out_features=768, bias=True) |
| 88 | + (dropout): Dropout(p=0.1, inplace=False) |
| 89 | + ) |
| 90 | + (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True) |
| 91 | + (dropout): Dropout(p=0.1, inplace=False) |
| 92 | + ) |
| 93 | + (intermediate): MPNetIntermediate( |
| 94 | + (dense): Linear(in_features=768, out_features=3072, bias=True) |
| 95 | + (intermediate_act_fn): GELUActivation() |
| 96 | + ) |
| 97 | + (output): MPNetOutput( |
| 98 | + (dense): Linear(in_features=3072, out_features=768, bias=True) |
| 99 | + (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True) |
| 100 | + (dropout): Dropout(p=0.1, inplace=False) |
| 101 | + ) |
| 102 | + ) |
| 103 | + ) |
| 104 | + (relative_attention_bias): Embedding(32, 12) |
| 105 | + ) |
| 106 | + (pooler): MPNetPooler( |
| 107 | + (dense): Linear(in_features=768, out_features=768, bias=True) |
| 108 | + (activation): Tanh() |
| 109 | + ) |
| 110 | +) |
| 111 | +``` |
| 112 | +This output will differ depending on the model. We will look for the repeated layers in the encoder. For this MPNet model, this is stored under `model[0].auto_model.encoder.layer`. Then we can slice the model to only keep the first few layers to speed up the model: |
| 113 | + |
| 114 | +```python |
| 115 | +new_num_layers = 3 |
| 116 | +model[0].auto_model.encoder.layer = model[0].auto_model.encoder.layer[:new_num_layers] |
| 117 | +``` |
| 118 | + |
| 119 | +Then we can run inference with it using <a href="../../../docs/package_reference/SentenceTransformer.html#sentence_transformers.SentenceTransformer.encode"><code>SentenceTransformers.encode</code></a>. |
| 120 | + |
| 121 | +```python |
| 122 | +from sentence_transformers import SentenceTransformer |
| 123 | +from sentence_transformers.util import cos_sim |
| 124 | + |
| 125 | +model = SentenceTransformer("tomaarsen/mpnet-base-nli-adaptive-layer") |
| 126 | +new_num_layers = 3 |
| 127 | +model[0].auto_model.encoder.layer = model[0].auto_model.encoder.layer[:new_num_layers] |
| 128 | + |
| 129 | +embeddings = model.encode( |
| 130 | + [ |
| 131 | + "The weather is so nice!", |
| 132 | + "It's so sunny outside!", |
| 133 | + "He drove to the stadium.", |
| 134 | + ] |
| 135 | +) |
| 136 | +# Similarity of the first sentence with the other two |
| 137 | +similarities = cos_sim(embeddings[0], embeddings[1:]) |
| 138 | +# => tensor([[0.7761, 0.1655]]) |
| 139 | +# compared to tensor([[ 0.7547, -0.0162]]) for the full model |
| 140 | +``` |
| 141 | +As you can see, the similarity between the related sentences is much higher than the unrelated sentence, despite only using 3 layers. Feel free to copy this script locally, modify the `new_num_layers`, and observe the difference in similarities. |
| 142 | + |
| 143 | + |
| 144 | +## Code Examples |
| 145 | + |
| 146 | +See the following scripts as examples of how to apply the <a href="../../../docs/package_reference/losses.html#adaptivelayerloss"><code>AdaptiveLayerLoss</code></a> in practice: |
| 147 | + |
| 148 | +* **[adaptive_layer_nli.py](adaptive_layer_nli.py)**: This example uses the `MultipleNegativesRankingLoss` with `AdaptiveLayerLoss` to train a strong embedding model using Natural Language Inference (NLI) data. It is an adaptation of the [NLI](../nli/README) documentation. |
| 149 | +* **[adaptive_layer_sts.py](adaptive_layer_sts.py)**: This example uses the CoSENTLoss with AdaptiveLayerLoss to train an embedding model on the training set of the STSBenchmark dataset. It is an adaptation of the [STS](../sts/README) documentation. |
| 150 | + |
| 151 | +And the following scripts to see how to apply <a href="../../../docs/package_reference/losses.html#matryoshka2dloss"><code>Matryoshka2dLoss</code></a>: |
| 152 | +* **[2d_matryoshka_nli.py](../matryoshka/2d_matryoshka_nli.py)**: This example uses the `MultipleNegativesRankingLoss` with `Matryoshka2dLoss` to train a strong embedding model using Natural Language Inference (NLI) data. It is an adaptation of the [NLI](../nli/README) documentation. |
| 153 | +* **[2d_matryoshka_sts.py](../matryoshka/2d_matryoshka_sts.py)**: This example uses the `CoSENTLoss` with `Matryoshka2dLoss` to train an embedding model on the training set of the STSBenchmark dataset. It is an adaptation of the [STS](../sts/README) documentation. |
0 commit comments