Skip to content

model

Initialization for core model modules.

Modules:

Name Description
noisy_encoder_transformer_masking_model

Module containing the NoiseMaskedNoisyEncoderTransformerModel class, a subclass of NoiseMaskedNoisyTransformerModel that supports

noisy_model

Module for creating noisy models by integrating a noise layer.

noisy_transformer_masking_model

Module for a noisy transformer model with masking capabilities.

peft_noisy_encoder_transformer_masking_model

Module for PEFT-enabled noisy transformer models with masking capabilities.

peft_noisy_transformer_masking_model

Module for PEFT-enabled noisy transformer models with masking capabilities.

privacy_encoder

Optional client-side privacy encoder applied after a Stained Glass Transform noise layer.

truncated_module

Module for a truncated module that interrupts forward passes.

Classes:

Name Description
NoisyModel

Applies a BaseNoiseLayer to a model input Tensor or a submodule output Tensor.

PrivacyEncoder

Client-side MLP Encoder: z -> y.

TruncatedModule

A module that wraps another module that interrupts the forward pass when a specified truncation point is reached.

NoisyModel

Bases: Module, ABC, Generic[ModuleT, NoiseLayerP, NoiseLayerT_co]

Applies a BaseNoiseLayer to a model input Tensor or a submodule output Tensor.

Parameters:

Name Type Description Default

noise_layer_class

Callable[NoiseLayerP, NoiseLayerT_co]

The type of BaseNoiseLayer to apply.

required

base_model

ModuleT

The model to apply the BaseNoiseLayer to.

required

*args

args

Positional arguments to noise_layer_class.

()

target_layer

str | None

The name of the base_model submodule (e.g. 'features.0.conv.1.2') whose output Tensor to transform. If provided, target_parameter must be None.

None

target_parameter

str | None

The name of the base_model input Tensor argument to transform. If provided, target_layer must be None.

None

**kwargs

kwargs

Keyword arguments to noise_layer_class.

{}

Raises:

Type Description
ValueError

If both target_layer and target_parameter are None.

ValueError

If neither target_layer nor target_parameter are None.

Methods:

Name Description
__getstate__

Serialize the model to a dictionary.

__setstate__

Deserialize the model from a dictionary.

clean_and_noisy_forward

Return both the clean and the noisy outputs of base_model for the same input.

deserialize_base_model

Deserialize the base model from its serialized representation.

disable_noise_layer

Temporarily skip the noise layer for calls into self.forward.

distillation_context

Prepare the base model to facilitate distillation training by applying losses over the transformed and non-transformed

forward

Call the base_model, applying the noise_layer to the target_parameter or target_layer output.

reset_parameters

Reinitialize parameters and buffers.

serialize_base_model

Serialize the base model.

serialize_init_passthrough_kwargs

Serialize the keyword arguments needed to re-initialize the NoisyModel from its serialized base model

Attributes:

Name Type Description
target_layer Module

The base_model submodule whose output Tensor to transform.

target_parameter str | None

The name of the base_model input Tensor argument to transform when target_layer is None.

target_parameter_index int

The index of the base_model input Tensor argument to transform when target_layer is None.

target_layer property

target_layer: Module

The base_model submodule whose output Tensor to transform.

Raises:

Type Description
ValueError

If _target_layer cannot be found as a submodule of base_model.

target_parameter property

target_parameter: str | None

The name of the base_model input Tensor argument to transform when target_layer is None.

target_parameter_index cached property

target_parameter_index: int

The index of the base_model input Tensor argument to transform when target_layer is None.

__getstate__

__getstate__() -> dict[str, Any]

Serialize the model to a dictionary.

Returns:

Type Description
dict[str, Any]

A dictionary containing the model's state, including the base model, noise layer, and state dict.

Added in version v3.18.0. Add support for serializing NoisyModel instances.

__setstate__

__setstate__(
    state: Mapping[str, Any],
    trust_remote_code: bool = False,
    third_party_model_path: (
        str | PathLike[str] | None
    ) = None,
) -> None

Deserialize the model from a dictionary.

Warning

The state_dict key is considered optional. If it is not present, or is incomplete, the missing parameters will be initialized to the meta device. Allowing this to be optional enables the NoisyModel parameters to be restored as part of a larger model.

Parameters:

Name Type Description Default

state

Mapping[str, Any]

A dictionary containing the model's state, including the base model, noise layer, and possibly state dict.

required

trust_remote_code

bool

Whether to trust remote code when loading from the Hugging Face Hub.

False

third_party_model_path

str | PathLike[str] | None

The path or huggingface reference to a third-party model to load. This is useful when loading SGTs whose internal structure depends on transformers which are not importable directly through transformers, but are present on the Hugging Face Hub.

None

Added in version v3.18.0. Add support for serializing NoisyModel instances.

Changed in version v3.78.0: A post-processor recorded in the serialized noise layer state is now reconstructed and passed to the noise layer's constructor.

clean_and_noisy_forward

clean_and_noisy_forward(
    *args: Any,
    frozen_overrides: Mapping[str, Module] | None = None,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> tuple[Any, Any]

Return both the clean and the noisy outputs of base_model for the same input.

When frozen_overrides is None or empty the model is run on the default training path. When frozen_overrides is provided the teacher and student are evaluated separately: the teacher runs with the given replacements installed at the indicated dotted paths, under torch.no_grad, and with the noise layer disabled via disable_noise_layer; the student runs normally with the live submodules in place. Use the two-branch form when one or more submodules of base_model are also being trained and the teacher must reference frozen pretrained weights rather than the live ones.

noise_mask, when provided, is passed only to the student branch — the teacher does not apply the noise layer in either path.

Output structure: each of clean_output and noisy_output mirrors the base model's output structure. Mapping containers are reconstructed as plain dict (so subclasses such as HuggingFace ModelOutput are not preserved); list / tuple are reconstructed in kind; tensor leaves are sliced along axis 0; non-tensor leaves are shared between the two outputs (no copy).

Parameters:

Name Type Description Default

*args

Any

Positional arguments forwarded to base_model.

required

frozen_overrides

Mapping[str, Module] | None

Optional map from dotted submodule path to its frozen replacement. Construct each replacement with FrozenSnapshot or any equivalent clone-and-freeze utility. Passing an empty mapping is equivalent to None.

None

noise_mask

Tensor | None

Optional mask passed to the student branch only (never to the teacher).

None

**kwargs

Any

Keyword arguments forwarded to base_model.

required

Returns:

Type Description
tuple[Any, Any]

Tuple (clean_output, noisy_output) matching the base model's output structure. The clean output is detached from the

tuple[Any, Any]

autograd graph; the noisy output carries gradients into the cloak and into any live trainable submodules of the base model.

Raises:

Type Description
TypeError

If noise_mask is provided but is not a torch.Tensor or None.

Added in version v3.40.0. Dispatch entrypoint that returns `(clean_output, noisy_output)`. Picks between the default training path and a separate teacher + student forward depending on whether `frozen_overrides` is supplied.

deserialize_base_model classmethod

deserialize_base_model(
    state: Mapping[str, Any],
    trust_remote_code: bool = False,
    third_party_model_path: (
        str | PathLike[str] | None
    ) = None,
) -> nn.Module

Deserialize the base model from its serialized representation.

This is used by __setstate__ to reconstruct the base model from its serialized representation.

In general, this does not load the state dict, since the noisy model itself will handle loading the state dict. However, the deserialization may need to involve loading model weights from a pretrained checkpoint, e.g. when the base model is a Hugging Face Transformers model and the serialization includes a Hugging Face config.

The default implementation of deserialize_base_model is to import the class from the base_model_type_str field in the serialization, then call __setstate__ on an instance of that class with the base_state field in the serialization. Subclasses can override this method to customize how the base model is deserialized, e.g. by using the Hugging Face from_pretrained method to load a model from a config.

Parameters:

Name Type Description Default

state

Mapping[str, Any]

The serialized representation of the base model.

required

trust_remote_code

bool

Whether to allow executing remote code when deserializing third-party models. This should only be set to True when deserializing models from trusted sources, as executing remote code can be dangerous.

False

third_party_model_path

str | PathLike[str] | None

An optional path to a local directory containing the files needed to deserialize a third-party model, which may include custom code. This is used when deserializing third-party models that require custom code, and should only be used with models from trusted sources.

None

Returns:

Type Description
nn.Module

The deserialized base model.

Raises:

Type Description
TypeError

If the deserialized base model is not an instance of nn.Module.

disable_noise_layer

disable_noise_layer() -> Generator[None]

Temporarily skip the noise layer for calls into self.forward.

Inside the block, self(*args, **kwargs) calls self.base_model directly: no noise forward hook is registered, the noise layer's forward is not invoked, the noise-layer-called sentinel is not enforced, and the noise_mask kwarg is silently ignored. Restores normal noised forward on exit, even if the wrapped block raises.

Nesting is supported (inner exits restore to the outer state). The context manager is not thread-safe; concurrent forwards on the same NoisyModel from multiple threads may observe an inconsistent enabled/disabled state.

Useful for: - the teacher branch of clean_and_noisy_forward, - logging a clean baseline at validation time, - computing a non-private reference metric.

Examples:

>>> with model.disable_noise_layer():
...     clean_logits = model(input_ids=batch.input_ids)

Added in version v3.40.0. Context manager to temporarily bypass the noise layer; `forward` returns the un-noised base model output for the duration of the block. Used by the teacher branch of `clean_and_noisy_forward` and for clean-baseline logging.

distillation_context

distillation_context() -> contextlib.ExitStack

Prepare the base model to facilitate distillation training by applying losses over the transformed and non-transformed activations.

Note

This context manager assumes that the output of the base_model is a mutable mapping with a logits key.

Returns:

Type Description
contextlib.ExitStack

A context manager that detaches the hooks when exited.

Added in version v2.6.0.

forward

forward(
    *args: Any,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> Any

Call the base_model, applying the noise_layer to the target_parameter or target_layer output.

When wrapped in a disable_noise_layer context, the noise layer hooks are not attached and the base_model is called with the original arguments. noise_mask is silently ignored in that case — the model returns the un-noised output.

Parameters:

Name Type Description Default

*args

Any

Positional arguments to base_model.

required

noise_mask

Tensor | None

An optional mask that selects the elements of the target_parameter or target_layer output to transform. Where the mask is False, the original values of the target are used. If None, the entire target is transformed.

None

**kwargs

Any

Keyword arguments to base_model.

required

Returns:

Type Description
Any

The result of base_model with the noise_layer applied to the target_parameter or target_layer output.

reset_parameters

reset_parameters() -> None

Reinitialize parameters and buffers.

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

serialize_base_model

serialize_base_model() -> dict[str, Any]

Serialize the base model.

This is used by getstate to get a JSON-serializable representation of the base model. This does not, in general, have to include a copy of the state dict, since the noisy model itself will store the base model state dict within its own state dict.

The default implementation of serialize_base_model is to remove any hooks, then call base_model.__getstate__(), which works for some models, but some models may need to override this method to return a more JSON-serializable representation of the base model, e.g. by returning the base model's Hugging Face config.

Subclasses can override this method to customize how the base model is serialized, e.g. by returning a Hugging Face config instead of the instance dictionary.

More common than subclassing NoisyModel, however, would be to wrap the base model in a custom class that implements its own JSON-serializable/loadable __getstate__/__setstate__ methods, and then pass that custom wrapper instance in as the base model.

Returns:

Type Description
dict[str, Any]

A serialized representation of the base model.

Raises:

Type Description
NotImplementedError

If the base model's __getstate__ method is the default torch.nn.Module.__getstate__, since the default __getstate__ has undefined behavior when serializing a NoisyModel. (The default __getstate__ just copies the __dict__ of the module, which, in general, is not JSON-serializable).

ValueError

If the base model's state contains non-JSON-serializable values, which cannot be serialized as part of the NoisyModel.

Added in version v3.18.0. Add support for serializing NoisyModel instances.

serialize_init_passthrough_kwargs

serialize_init_passthrough_kwargs() -> dict[str, Any]

Serialize the keyword arguments needed to re-initialize the NoisyModel from its serialized base model and noise layer.

This is used by __getstate__ to get a JSON-serializable representation of the keyword arguments needed to re-initialize the NoisyModel from its serialized base model and noise layer.

Returns:

Type Description
dict[str, Any]

A dictionary containing the keyword arguments needed to re-initialize the NoisyModel from its serialized

dict[str, Any]

form.

Added in version v3.18.0. Add support for serializing NoisyModel instances.

PrivacyEncoder

Bases: PostProcessor

Client-side MLP Encoder: z -> y.

Expands the latent space and applies residual blocks before projecting down. use_skip and identity_init are independent switches; all four combinations are supported:

  • use_skip=True, identity_init=True (the default): proj_down is zero-initialized, so the global residual connection (z + ...) makes the encoder compute the identity at initialization regardless of proj_up/blocks.
  • use_skip=True, identity_init=False: proj_down is left at its plain (random) nn.Linear init, so the encoder computes z + MLP(z) with a non-trivial MLP from the start.
  • use_skip=False, identity_init=True: proj_down is set to the exact left-inverse of proj_up (via pseudo-inverse), giving proj_down(proj_up(z)) == z at initialization -- the residual blocks are already identity at init on their own, via their internal skip connections. This is only mathematically possible when hidden_dim >= dim (a rank constraint); otherwise identity_init raises.
  • use_skip=False, identity_init=False: proj_down is left at its plain (random) nn.Linear init, so the encoder is a genuinely non-trivial function of its input from the start.

Intended to run immediately after a Stained Glass Transform noise layer, applied only to the embedding positions the noise layer actually transformed.

Methods:

Name Description
__init__

Build the encoder's layers (proj_up, the residual blocks, and proj_down) and initialize their parameters.

deserialize

Reconstruct a post-processor from the dictionary produced by serialize.

forward

Apply the privacy encoder to input.

from_pretrained

Load a privacy encoder previously saved with save_pretrained.

get_config

Get the configuration needed to reconstruct this encoder's architecture before loading weights.

reset_parameters

Reinitialize every parameter in the encoder, then set proj_down so the encoder starts as an identity map (or as

save_pretrained

Save this encoder's configuration and weights to a single file.

serialize

Serialize this post-processor's type and configuration to a JSON-serializable dictionary.

__init__

__init__(
    dim: int,
    hidden_dim: int,
    n_layers: int = 4,
    seed: int | None = None,
    use_skip: bool = True,
    identity_init: bool = True,
) -> None

Build the encoder's layers (proj_up, the residual blocks, and proj_down) and initialize their parameters.

proj_up, the residual blocks, and proj_down are constructed and then reinitialized by reset_parameters within the same isolated RNG fork (seeded by seed if given), so the encoder starts as an identity map (or as close to one as possible); see reset_parameters for details.

Parameters:

Name Type Description Default

dim

int

The dimensionality of the embeddings this encoder operates on. Must match the base model's hidden size.

required

hidden_dim

int

The dimensionality of the encoder's internal representation.

required

n_layers

int

The number of residual MLP blocks.

4

seed

int | None

Optional seed used only for initialization, isolated from the global RNG stream. Has no effect once weights are loaded from a checkpoint.

None

use_skip

bool

Whether the encoder adds a global residual connection (z + MLP(z)) or returns the MLP's output directly (MLP(z)). This changes the encoder's forward computation, not just its initialization, and is not recoverable from a loaded state dict — it must match how the encoder was trained.

True

identity_init

bool

Whether the encoder starts as an identity map.

True

Raises:

Type Description
ValueError

If identity_init is requested with use_skip=False and hidden_dim < dim.

deserialize classmethod

deserialize(state: Mapping[str, Any]) -> Self

Reconstruct a post-processor from the dictionary produced by serialize.

Calling this on PostProcessor accepts any post-processor, while calling it on a subclass accepts only that subclass or its descendants.

The reconstructed post-processor's parameters hold whatever the class initializes them to; the containing noise layer is responsible for loading the trained weights from its own state dictionary.

Parameters:

Name Type Description Default

state

Mapping[str, Any]

The serialized post-processor type and configuration.

required

Returns:

Type Description
Self

The reconstructed post-processor.

Raises:

Type Description
TypeError

If the serialized type is not a subclass of the class this is called on.

forward

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

Apply the privacy encoder to input.

Parameters:

Name Type Description Default

input

Tensor

Input embeddings of shape (..., dim).

required

noise_mask

Tensor | None

Unused. This encoder operates pointwise, so it computes the same value for a position regardless of which other positions the noise layer transformed.

None

Returns:

Type Description
torch.Tensor

A tensor of the same shape and dtype as input.

from_pretrained classmethod

from_pretrained(
    path: str | Path,
    map_location: device | str | None = None,
) -> PrivacyEncoder

Load a privacy encoder previously saved with save_pretrained.

Parameters:

Name Type Description Default

path

str | Path

The file path to load from.

required

map_location

device | str | None

The device to load the encoder's weights onto. If None, weights are loaded onto the device they were saved from.

None

Returns:

Type Description
PrivacyEncoder

The loaded privacy encoder.

get_config

get_config() -> PrivacyEncoderConfig

Get the configuration needed to reconstruct this encoder's architecture before loading weights.

Returns:

Type Description
PrivacyEncoderConfig

The encoder's configuration.

reset_parameters

reset_parameters() -> None

Reinitialize every parameter in the encoder, then set proj_down so the encoder starts as an identity map (or as close to one as possible), if identity_init.

Reinitializes proj_up and the residual blocks too, not just proj_down, so this method alone is sufficient to initialize the encoder after materializing it from the meta device (e.g. via to_empty) -- construction-time initialization cannot be relied on there, since meta tensors have no storage to initialize into. A no-op while the encoder's parameters are still on the meta device, matching the convention used throughout noise_layer/: the initialization below is real (dispatched) computation, not a free no-op, so it must be re-run via an explicit call to this method after materializing real storage.

Must be called directly on the encoder, not relied upon to run via a generic per-parameter-owning-module meta-initialization walk (e.g. FSDP's default meta materialization): neither PrivacyEncoder nor ResidualMLPBlock owns any parameter directly (they all live on child nn.Linear/nn.LayerNorm submodules), so such a walk would reset those leaves directly and skip this method entirely, silently leaving the encoder at plain random init instead of identity.

Raises:

Type Description
ValueError

If identity_init is requested with use_skip=False and hidden_dim < dim.

save_pretrained

save_pretrained(path: str | Path) -> None

Save this encoder's configuration and weights to a single file.

This is the checkpoint format from_pretrained expects. Deliberately independent of StainedGlassTransform's .sgt serialization: a privacy encoder is not tied to any particular SGT and is meant to be trained, saved, and loaded on its own.

Parameters:

Name Type Description Default

path

str | Path

The file path to save to.

required

serialize

serialize() -> dict[str, Any]

Serialize this post-processor's type and configuration to a JSON-serializable dictionary.

Returns:

Type Description
dict[str, Any]

This post-processor's fully qualified class name alongside its configuration.

Raises:

Type Description
ValueError

If this post-processor's configuration uses the reserved post_processor_type key.

TruncatedModule

Bases: Module, Generic[ModuleT]

A module that wraps another module that interrupts the forward pass when a specified truncation point is reached.

This truncation happens by temporarily adding a hook to the truncation point that raises a TruncationExecutionFinished exception which is then caught by the TruncatedModule forward and the output of the truncation point is returned.

Examples:

Instantiating a TruncatedModule with a Binary Classification model and a truncation point:

>>> model = torch.nn.Sequential(
...     torch.nn.Linear(10, 20),
...     torch.nn.ReLU(),
...     torch.nn.Linear(20, 30),
...     torch.nn.ReLU(),
...     torch.nn.Linear(30, 40),
...     torch.nn.ReLU(),
...     torch.nn.Linear(40, 2),
... )
>>> truncation_layer = model[1]
>>> truncated_model = TruncatedModule(model, truncation_layer)

Using the TruncatedModule to get the output of the truncation point:

>>> input = torch.randn(1, 10)
>>> output = truncated_model(input)
>>> # Note that shape of the output has the output_shape of the truncation point, not the full model
>>> assert output.shape == (1, 20)

The base model of the TruncatedModule is completely unaffected by the truncation:

>>> base_output = model(input)
>>> assert base_output.shape == (1, 2)  # Binary classification output shape

The base model is also accessible directly through the module attribute of the TruncatedModule:

>>> base_output = truncated_model.module(input)
>>> assert base_output.shape == (1, 2)  # Binary classification output shape

Methods:

Name Description
__init__

Initialize the TruncatedModule with the provided module and truncation point.

forward

Forward pass of the TruncatedModule that interrupts the forward pass when the truncation point is reached.

lazy_register_truncation_hook

Create a prehook that will be added to the truncation point to interrupt the forward pass when the truncation point is reached.

truncation_hook

Intercept the output of the truncation point and raise a TruncationExecutionFinished exception containing that output.

__init__

__init__(module: ModuleT, truncation_point: Module) -> None

Initialize the TruncatedModule with the provided module and truncation point.

Parameters:

Name Type Description Default

module

ModuleT

The module to wrap.

required

truncation_point

Module

The submodule of the provided module at which to interrupt the forward pass.

required

Raises:

Type Description
ValueError

If the truncation point is not a submodule of the provided module.

forward

forward(*args: Any, **kwargs: Any) -> Any

Forward pass of the TruncatedModule that interrupts the forward pass when the truncation point is reached.

Parameters:

Name Type Description Default

*args

Any

The positional arguments to pass to the wrapped module.

required

**kwargs

Any

The keyword arguments to pass to the wrapped module.

required

Returns:

Type Description
Any

The output of the truncation point submodule.

Raises:

Type Description
HookNotCalledError

If the truncation hook is not called, meaning the truncation point was not reached.

lazy_register_truncation_hook

lazy_register_truncation_hook() -> _HandlerWrapper

Create a prehook that will be added to the truncation point to interrupt the forward pass when the truncation point is reached.

Returns:

Type Description
_HandlerWrapper

A handler wrapper that contains the hook that was added to the truncation point.

truncation_hook staticmethod

truncation_hook(
    truncation_point: Module, args: Any, output: Tensor
) -> NoReturn

Intercept the output of the truncation point and raise a TruncationExecutionFinished exception containing that output.

Parameters:

Name Type Description Default

truncation_point

Module

The truncation point submodule. Unused.

required

args

Any

The arguments passed to the truncation point. Unused.

required

output

Tensor

The output of the truncation point. This is the output that will be returned by the TruncatedModule.

required

Raises:

Type Description
TruncationExecutionFinished

Always, in order to interrupt the wrapped model's forward method.