transformer_cloak

Classes:

Name	Description
`TransformerCloak`	Applies a stochastic transformation to a causal language model embedding `Tensor` using `TransformerBlockEstimator`,

TransformerBlockEstimator ¶

Bases: Module, Generic[TransformerT]

Estimates components of sequence dependent noise using a single layer transformer model.

Parameters:

Name	Type	Description	Default
`transformer_type` ¶	`type[TransformerT]`	The type of transformer model to build a single layer estimator of, e.g. `transformers.LlamaModel` or `transformers.MistralModel`.	required
`config_path` ¶	`str`	The path to the transformers config.	required
`initial` ¶	`float`	Initial value of the final `Linear` bias.	`0.0`
`dropout` ¶	`float`	Dropout probability of the transformer model output.	`0.1`
`use_causal_mask` ¶	`bool`	Whether to use a causal or a non-causal attention mask.	`True`
`initialization_scale` ¶	`float`	The scale factor to multiply the initial values of `linear.weight` by.	`0.05`
`**kwargs` ¶	`Any`	Additional keyword arguments to [`transformers.PretrainedConfig.from_pretrained`] used to create the underlying transformer model.	required

Raises:

Type	Description
`ValueError`	If the PyTorch version is <2.0.0 and `use_causal_mask` is `False`.

Methods:

Name	Description
`__init__`
`forward`	Compose the transformer block with a dropout and a linear adapter layer.
`reset_parameters`	Reinitialize parameters and buffers.

init ¶

__init__(
    transformer_type: type[TransformerT],
    config_path: str,
    initial: float = 0.0,
    dropout: float = 0.1,
    use_causal_mask: bool = True,
    initialization_scale: float = 0.05,
    **kwargs: Any,
) -> None

Changed in version 0.85.0: Passing param path to transformer cloak was highly error prone and unreasonable for the typical user.

forward ¶

forward(*args: Any, **kwargs: Any) -> torch.Tensor

Compose the transformer block with a dropout and a linear adapter layer.

Parameters:

Name	Type	Description	Default
`*args` ¶	`Any`	Positional arguments to the transformer model.	required
`**kwargs` ¶	`Any`	Keyword arguments to the transformer model.	required

Returns:

Type	Description
`torch.Tensor`	The output of the transformer parameter model.

Changed in version 0.75.1: The noise mask should always be non-None when using TransformerCloak.

reset_parameters ¶

reset_parameters() -> None

Reinitialize parameters and buffers.

This method is useful for initializing tensors created on the meta device.

TransformerCloak ¶

Bases: BaseNoiseLayer[TransformerBlockEstimator[TransformerT], Union[CloakStandardDeviationParameterization, DirectStandardDeviationParameterization], Optional[PercentMasker]]

Applies a stochastic transformation to a causal language model embedding Tensor using TransformerBlockEstimator, with standard deviations parameterized by either CloakStandardDeviationParameterization or DirectStandardDeviationParameterization, and optional standard deviation-based input masking using PercentMasker.

Parameters:

Name	Type	Description	Default
`scale` ¶	`tuple[float, float]`	The range of standard deviations of the noise.	required
`transformer_type` ¶	`type[TransformerT]`	The type of the transformer to build a single layer estimator from.	required
`config_path` ¶	`str`	Path to transformer config.	required
`percent_to_mask` ¶	`float \| None`	The percentage of the input to mask.	`None`
`shallow` ¶	`float`	A fixed temperature like parameter which alters the scale of the standard deviation of the noise.	`1.0`
`seed` ¶	`int \| None`	Seed for the random number generator used to generate noise.	`None`
`rho_init` ¶	`float`	Initial values for rhos.	`-3.0`
`std_dropout` ¶	`float`	Dropout ratio for std parameter model.	`0.0`
`mean_dropout` ¶	`float`	Dropout ratio for mean parameter model.	`0.0`
`directly_learn_stds` ¶	`bool`	Whether or not the rhos estimator is used to learn rhos (values in R) or standard deviations directly (values in R^+).	`False`
`mean_num_experts` ¶	`int`	The number of experts to use for the multilayer perceptron after the attention layer for mean_estimator. The value zero corresponds to not using mixture of experts.	`0`
`std_num_experts` ¶	`int`	The number of experts to use for the multilayer perceptron after the attention layer for std_estimator. The value zero corresponds to not using mixture of experts.	`0`
`use_causal_mask` ¶	`bool`	Whether to use a causal or a non-causal attention mask in the llama estimator.	`True`
`kwargs` ¶	`Any`	Keyword arguments used to define the transformer parameter models.	required

Raises:

Type	Description
`ValueError`	If `shallow` is not `1.0` when `directly_learn_stds` is `True`.
`ValueError`	If `rho_init` is not `0.0` when `directly_learn_stds` is `True`.

Methods:

Name	Description
`__call__`	Transform the input data.
`__getstate__`	Prepare a serializable copy of `self.__dict__`.
`__init__`
`__init_subclass__`	Set the default dtype to `torch.float32` inside all subclass `__init__` methods.
`__setstate__`	Restore from a serialized copy of `self.__dict__`.
`forward`	Transform the input data.
`get_applied_transform_components_factory`	Create a function that returns the elements of the transform components (`'mean'` and `'std'`) applied during the most recent
`get_transformed_output_factory`	Create a function that returns the transformed output from the most recent forward pass.
`initial_seed`	Return the initial seed of the CPU device's random number generator.
`manual_seed`	Seed each of the random number generators.
`reset_parameters`	Reinitialize parameters and buffers.
`seed`	Seed each of the random number generators using a non-deterministic random number.

call ¶

__call__(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any,
) -> torch.Tensor

Transform the input data.

Parameters:

Name	Type	Description	Default
`input` ¶	`Tensor`	The input to transform.	required
`noise_mask` ¶	`Tensor \| None`	An optional mask that selects the elements of `input` to transform. Where the mask is `False`, the original `input` value is returned. Also used to select the elements of the sampled standard deviations to use to mask the `input`. If `None`, the entire `input` is transformed.	`None`
`**kwargs` ¶	`Any`	Additional keyword arguments to the estimator modules.	required

getstate ¶

__getstate__() -> dict[str, Any]

Prepare a serializable copy of self.__dict__.

init ¶

__init__(
    scale: tuple[float, float],
    transformer_type: type[TransformerT],
    config_path: str,
    percent_to_mask: float | None = None,
    shallow: float = 1.0,
    seed: int | None = None,
    rho_init: float = -3.0,
    std_dropout: float = 0.0,
    mean_dropout: float = 0.0,
    directly_learn_stds: bool = False,
    mean_num_experts: int = 0,
    std_num_experts: int = 0,
    use_causal_mask: bool = True,
    **kwargs: Any,
) -> None

Changed in version 0.85.0: Passing param path to transformer cloak was highly error prone and unreasonable for the typical user.

Changed in version 0.105.0: The std_loss_type argument is deprecated and no longer has any effect.

__init_subclass__ ¶

__init_subclass__() -> None

Set the default dtype to torch.float32 inside all subclass __init__ methods.

setstate ¶

__setstate__(state: dict[str, Any]) -> None

Restore from a serialized copy of self.__dict__.

forward ¶

forward(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any,
) -> torch.Tensor

Transform the input data.

Parameters:

Name	Type	Description	Default
`input` ¶	`Tensor`	The input to transform.	required
`noise_mask` ¶	`Tensor \| None`	A mask that selects the elements of `input` to transform. Where the mask is `False`, the original `input` value is returned. Also used to select the elements of the sampled standard deviations to use to mask the `input`.	`None`
`**kwargs` ¶	`Any`	Additional keyword arguments to the estimator modules.	required

Returns:

Type	Description
`torch.Tensor`	The transformed input data.

Raises:

Type	Description
`ValueError`	If the `noise_mask` is `None`.

Changed in version 0.74.0: The `noise_token_mask` was renamed to `noise_mask` to create a uniform interface everywhere.

get_applied_transform_components_factory ¶

get_applied_transform_components_factory() -> Callable[
    [], dict[str, torch.Tensor]
]

Create a function that returns the elements of the transform components ('mean' and 'std') applied during the most recent forward pass.

Specifically, the applied elements are those selected by the noise mask (if supplied) and standard deviation mask (if std_estimator.masker is not None). If no masks are used, all elements are returned.

The applied transform components are returned flattened.

This function is intended to be used to log histograms of the transform components.

Returns:

Type	Description
`Callable[[], dict[str, torch.Tensor]]`	A function that returns the the elements of the transform components applied during the most recent forward pass.

Examples:

>>> from torch import nn
>>> from stainedglass_core import model as sg_model, noise_layer as sg_noise_layer
>>> base_model = nn.Linear(20, 2)
>>> noisy_model = sg_model.NoisyModel(
...     sg_noise_layer.CloakNoiseLayer1,
...     base_model,
...     target_parameter="input",
... )
>>> get_applied_transform_components = (
...     noisy_model.noise_layer.get_applied_transform_components_factory()
... )
>>> input = torch.ones(1, 20)
>>> noise_mask = torch.tensor(5 * [False] + 15 * [True])
>>> output = noisy_model(input, noise_mask=noise_mask)
>>> applied_transform_components = get_applied_transform_components()
>>> applied_transform_components
{'mean': tensor(...), 'std': tensor(...)}
>>> {
...     component_name: component.shape
...     for component_name, component in applied_transform_components.items()
... }
{'mean': torch.Size([15]), 'std': torch.Size([15])}

get_transformed_output_factory ¶

get_transformed_output_factory() -> Callable[
    [], torch.Tensor
]

Create a function that returns the transformed output from the most recent forward pass.

If super batching is active, only the transformed half of the super batch output is returned.

Returns:

Type	Description
`Callable[[], torch.Tensor]`	A function that returns the transformed output from the most recent forward pass.

Examples:

>>> from stainedglass_core import noise_layer as sg_noise_layer
>>> noise_layer = sg_noise_layer.CloakNoiseLayer1()
>>> get_transformed_output = noise_layer.get_transformed_output_factory()
>>> input = torch.ones(2, 3, 32, 32)
>>> output = noise_layer(input)
>>> transformed_output = get_transformed_output()
>>> assert output.equal(transformed_output)

initial_seed ¶

initial_seed() -> int

Return the initial seed of the CPU device's random number generator.

manual_seed ¶

manual_seed(seed: int | None) -> None

Seed each of the random number generators.

Setting seed to None will destroy any existing generators.

Parameters:

Name	Type	Description	Default
`seed` ¶	`int \| None`	The seed to set.	required

reset_parameters ¶

reset_parameters() -> None

Reinitialize parameters and buffers.

This method is useful for initializing tensors created on the meta device.

seed ¶

seed() -> None

Seed each of the random number generators using a non-deterministic random number.

transformer_parameter_model ¶

transformer_parameter_model(
    transformer_type: type[TransformerT],
    config_path: str,
    **kwargs: Any,
) -> TransformerT

Create a single block of a transformers.PreTrainedModel and loads the weights from the parameter path.

Parameters:

Name	Type	Description	Default
`transformer_type` ¶	`type[TransformerT]`	The type of the transformer to use to construct the transformer parameter model.	required
`config_path` ¶	`str`	Path to transformer config.	required
`**kwargs` ¶	`Any`	The keyword arguments to pass to transformers.PreTrainedModel.from_pretrained.	required

Returns:

Type	Description
`TransformerT`	A transformer that can be used to estimate rhos/locs.

Changed in version 0.85.0: Passing param path to transformer cloak was highly error prone and unreasonable for the typical user.

transformer_cloak

TransformerBlockEstimator ¶

transformer_type ¶

config_path ¶

initial ¶

dropout ¶

use_causal_mask ¶

initialization_scale ¶

**kwargs ¶

__init__ ¶

forward ¶

*args ¶

**kwargs ¶

reset_parameters ¶

TransformerCloak ¶

scale ¶

transformer_type ¶

config_path ¶

percent_to_mask ¶

shallow ¶

seed ¶

rho_init ¶

std_dropout ¶

mean_dropout ¶

directly_learn_stds ¶

mean_num_experts ¶

std_num_experts ¶

use_causal_mask ¶

kwargs ¶

__call__ ¶

input ¶

noise_mask ¶

**kwargs ¶

__getstate__ ¶

__init__ ¶

__init_subclass__ ¶

__setstate__ ¶

forward ¶

input ¶

noise_mask ¶

**kwargs ¶

get_applied_transform_components_factory ¶

get_transformed_output_factory ¶

initial_seed ¶

manual_seed ¶

seed ¶

reset_parameters ¶

seed ¶

transformer_parameter_model ¶

transformer_type ¶

config_path ¶

**kwargs ¶

`transformer_type` ¶

`config_path` ¶

`initial` ¶

`dropout` ¶

`use_causal_mask` ¶

`initialization_scale` ¶

`**kwargs` ¶

init ¶

`*args` ¶

`**kwargs` ¶

`scale` ¶

`transformer_type` ¶

`config_path` ¶

`percent_to_mask` ¶

`shallow` ¶

`seed` ¶

`rho_init` ¶

`std_dropout` ¶

`mean_dropout` ¶

`directly_learn_stds` ¶

`mean_num_experts` ¶

`std_num_experts` ¶

`use_causal_mask` ¶

`kwargs` ¶

call ¶

`input` ¶

`noise_mask` ¶

`**kwargs` ¶

getstate ¶

init ¶

setstate ¶

`input` ¶

`noise_mask` ¶

`**kwargs` ¶

`seed` ¶

`transformer_type` ¶

`config_path` ¶

`**kwargs` ¶