from __future__ import annotations import random import warnings from collections.abc import Iterable from typing import Any import torch from torch import Tensor, nn from torch.nn import functional as F from sentence_transformers import SentenceTransformer from sentence_transformers.losses.CachedGISTEmbedLoss import CachedGISTEmbedLoss from sentence_transformers.losses.CachedMultipleNegativesRankingLoss import CachedMultipleNegativesRankingLoss from sentence_transformers.models import Transformer class TransformerDecorator: """ Decorator that caches the embeddings of all layers of the transformer. When `layer_idx` is set, it returns the cached embeddings of that layer instead. This is meant to override the forward function of the Transformer. """ def __init__(self, transformer: Transformer, original_forward) -> None: self.transformer = transformer self.original_forward = original_forward self.embeddings: list[tuple[Tensor]] = [] self.last_embeddings: list[Tensor] = [] self.features: list[dict[str, Tensor]] = [] self.layer_idx = None self.call_idx = 0 def set_layer_idx(self, layer_idx) -> None: self.layer_idx = layer_idx self.call_idx = 0 def get_layer_embeddings(self) -> Tensor: return torch.concat([embedding[self.layer_idx] for embedding in self.embeddings], dim=1) def __call__(self, features) -> dict[str, Tensor]: if self.layer_idx is None: output = self.call_grow_cache(features) else: output = self.call_use_cache(features) self.call_idx += 1 return output def call_grow_cache(self, features: dict[str, Tensor]) -> dict[str, Tensor]: """ Temporarily sets the output_hidden_states to True, runs the model, and then restores the original setting. Use the all_layer_embeddings to get the embeddings of all layers. """ original_output_hidden_states = self.transformer.auto_model.config.output_hidden_states self.transformer.auto_model.config.output_hidden_states = True output = self.original_forward(features) # We ignore the first layer, as it is the input embeddings # and the last layer, as we already computed the loss over it self.num_layers = len(output["all_layer_embeddings"]) - 1 self.embeddings.append(output["all_layer_embeddings"][1:-1]) self.last_embeddings.append(output["token_embeddings"]) self.features.append( {key: value for key, value in output.items() if key not in ["all_layer_embeddings", "token_embeddings"]} ) # Restore original setting self.transformer.auto_model.config.output_hidden_states = original_output_hidden_states if original_output_hidden_states: del output["all_layer_embeddings"] return output def call_use_cache(self, features: dict[str, Tensor]) -> dict[str, Tensor]: return {**self.features[self.call_idx], "token_embeddings": self.embeddings[self.call_idx][self.layer_idx]} class ForwardDecorator: """ Decorator that caches the embeddings after all modules (e.g. pooling) of the model. Required to get the embeddings after all modules for the KL-divergence loss. This is meant to override the forward function of the SentenceTransformer. """ def __init__(self, fn) -> None: self.fn = fn self.embeddings = [] def __call__(self, features: dict[str, Tensor]) -> dict[str, Tensor]: output = self.fn(features) self.embeddings.append(output["sentence_embedding"]) return output def get_embeddings(self) -> Tensor: embeddings = torch.concat(self.embeddings, dim=0) self.embeddings = [] return embeddings class AdaptiveLayerLoss(nn.Module): def __init__( self, model: SentenceTransformer, loss: nn.Module, n_layers_per_step: int = 1, last_layer_weight: float = 1.0, prior_layers_weight: float = 1.0, kl_div_weight: float = 1.0, kl_temperature: float = 0.3, ) -> None: """ The AdaptiveLayerLoss can be seen as a loss *modifier* that allows you to use other loss functions at non-final layers of the Sentence Transformer model. This is useful for when you want to train a model where users have the option to lower the number of layers used to improve their inference speed and memory usage. Args: model: SentenceTransformer model loss: The loss function to be used, e.g. :class:`MultipleNegativesRankingLoss`, :class:`CoSENTLoss`, etc. n_layers_per_step: The number of layers to use per step. If -1, then all layers are used. If > 0, then a random sample of `n_layers_per_step` layers are used per step, separate from the final layer, which is always used. The 2DMSE paper uses `n_layers_per_step=1`. The default value is 1. last_layer_weight: The weight to use for the loss of the final layer. Increase this to focus more on the performance when using all layers. The default value is 1.0. prior_layers_weight: The weight to use for the loss of the prior layers. Increase this to focus more on the performance when using fewer layers. The default value is 1.0. kl_div_weight: The weight to use for the KL-divergence loss that is used to make the prior layers match that of the last layer. Increase this to focus more on the performance when using fewer layers. The default value is 1.0. kl_temperature: The temperature to use for the KL-divergence loss. If 0, then the KL-divergence loss is not used. The default value is 1.0. References: - The concept was inspired by the 2DMSE paper: https://arxiv.org/abs/2402.14776 - `Adaptive Layers <../../examples/training/adaptive_layer/README.html>`_ Requirements: 1. The base loss cannot be :class:`CachedMultipleNegativesRankingLoss` or :class:`CachedGISTEmbedLoss`. Inputs: +---------------------------------------+--------+ | Texts | Labels | +=======================================+========+ | any | any | +---------------------------------------+--------+ Relations: - :class:`Matryoshka2dLoss` uses this loss in combination with :class:`MatryoshkaLoss` which allows for output dimensionality reduction for faster downstream tasks (e.g. retrieval). Example: :: from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses from datasets import Dataset model = SentenceTransformer("microsoft/mpnet-base") train_dataset = Dataset.from_dict({ "anchor": ["It's nice weather outside today.", "He drove to work."], "positive": ["It's so sunny.", "He took the car to the office."], }) loss = losses.MultipleNegativesRankingLoss(model=model) loss = losses.AdaptiveLayerLoss(model, loss) trainer = SentenceTransformerTrainer( model=model, train_dataset=train_dataset, loss=loss, ) trainer.train() """ super().__init__() self.model = model self.loss = loss self.n_layers_per_step = n_layers_per_step self.last_layer_weight = last_layer_weight self.prior_layers_weight = prior_layers_weight self.kl_div_weight = kl_div_weight self.kl_temperature = kl_temperature assert isinstance(self.model[0], Transformer) if isinstance(loss, CachedMultipleNegativesRankingLoss): warnings.warn("MatryoshkaLoss is not compatible with CachedMultipleNegativesRankingLoss.", stacklevel=2) if isinstance(loss, CachedGISTEmbedLoss): warnings.warn("MatryoshkaLoss is not compatible with CachedGISTEmbedLoss.", stacklevel=2) def forward(self, sentence_features: Iterable[dict[str, Tensor]], labels: Tensor) -> Tensor: # Decorate the forward function of the transformer to cache the embeddings of all layers original_transformer_forward = self.model[0].forward transformer_decorator = TransformerDecorator(self.model[0], original_transformer_forward) self.model[0].forward = transformer_decorator # Decorate the forward function of the model to get the embeddings after all modules (e.g. pooling) original_forward = self.model.forward forward_decorator = ForwardDecorator(original_forward) self.model.forward = forward_decorator # Run the loss normally: i.e. the final layer, but 1) use the transformers decorator to cache # the embeddings of all layers and 2) use the forward decorator to get the embeddings after all modules # for the KL-divergence loss loss = self.loss(sentence_features, labels) * self.last_layer_weight if self.kl_temperature > 0: final_embeddings = forward_decorator.get_embeddings() final_embeddings = F.softmax(final_embeddings / self.kl_temperature, dim=-1) num_layers = transformer_decorator.num_layers layer_indices = range(num_layers - 1) if self.n_layers_per_step > 0 and self.n_layers_per_step < num_layers - 1: layer_indices = random.sample(layer_indices, self.n_layers_per_step) # This loop is over `num_layer - 1` layers because we already computed the loss over the final layer for layer_idx in layer_indices: # Add regular loss for each layer by using the cached embeddings of that layer transformer_decorator.set_layer_idx(layer_idx) layer_loss = self.loss(sentence_features, labels) loss = loss + layer_loss / (1 + layer_idx) / len(layer_indices) * self.prior_layers_weight # and KL-divergence loss between the current layer and the final layer # Note: we use "batchmean" reduction as that aligns with the mathematical definition if self.kl_temperature > 0: embeddings = forward_decorator.get_embeddings() kl_div_loss = F.kl_div( F.log_softmax(embeddings / self.kl_temperature, dim=-1), final_embeddings, reduction="batchmean", ) loss = loss + kl_div_loss * self.kl_temperature * self.kl_div_weight self.model[0].forward = original_transformer_forward self.model.forward = original_forward return loss def get_config_dict(self) -> dict[str, Any]: return { "loss": self.loss.__class__.__name__, "n_layers_per_step": self.n_layers_per_step, "last_layer_weight": self.last_layer_weight, "prior_layers_weight": self.prior_layers_weight, "kl_div_weight": self.kl_div_weight, "kl_temperature": self.kl_temperature, } @property def citation(self) -> str: return """ @misc{li20242d, title={2D Matryoshka Sentence Embeddings}, author={Xianming Li and Zongxi Li and Jing Li and Haoran Xie and Qing Li}, year={2024}, eprint={2402.14776}, archivePrefix={arXiv}, primaryClass={cs.CL} } """