distillation
distillation_loss_factory(noisy_model: NoiseMaskedNoisyTransformerModel[CausalModelT, ..., TransformerCloak[Any]], distillation_layer_index: int, alpha: float, std_log_ratio_loss_weight: float, input_embedding_similarity_loss_weight: float, distillation_layer_distance_loss_weight: float) -> tuple[Callable[[torch.Tensor], torch.Tensor], Callable[[], dict[str, torch.Tensor]], Callable[[], dict[str, Any]]]
Create a loss function to perform distillation training on a Stained Glass Transform causal language model using a teacher causal language model.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
noisy_model |
NoiseMaskedNoisyTransformerModel[CausalModelT, ..., TransformerCloak[Any]]
|
The model containing both the causal language base model and the Stained Glass Transform. |
required |
distillation_layer_index |
int
|
The index of the decoder layer in the base model at which to perform distillation. |
required |
alpha |
float
|
The interpolation factor between the distillation loss (maximizing model similarity) and the Stained Glass Transform loss (maximizing transformation strength). Should be in the range [0, 1], where 0 corresponds to higher model similarity and 1 corresponds to higher transformation strength. |
required |
std_log_ratio_loss_weight |
float
|
The weight of the loss component which aims to maximize the Stained Glass Transform's standard deviations (transformation strength). |
required |
input_embedding_similarity_loss_weight |
float
|
The weight of the loss component which aims to minimize the similarity of the input embeddings. |
required |
distillation_layer_distance_loss_weight |
float
|
The weight of the loss component which aims to maximize the similarity of the distillation layer embeddings. |
required |
Returns:
| Type | Description |
|---|---|
tuple[Callable[[torch.Tensor], torch.Tensor], Callable[[], dict[str, torch.Tensor]], Callable[[], dict[str, Any]]]
|
A tuple of 3 functions: the loss function, a function to retrieve the loss components, and a function to retrieve the |
tuple[Callable[[torch.Tensor], torch.Tensor], Callable[[], dict[str, torch.Tensor]], Callable[[], dict[str, Any]]]
|
hyperparameters. These functions may be called at most once each after a forward pass through both models. |