privacy_encoder

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

Classes:

Name	Description
`PrivacyEncoder`	Client-side MLP Encoder: z -> y.
`PrivacyEncoderConfig`	Configuration needed to reconstruct a `PrivacyEncoder` before loading
`ResidualMLPBlock`	A standard MLP block with a skip connection.

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.

PrivacyEncoderConfig ¶

Bases: TypedDict

Configuration needed to reconstruct a PrivacyEncoder before loading weights.

hidden_dim and n_layers parametrize the ResidualMLPBlock stack directly: hidden_dim is passed as every block's hidden_dim, and n_layers is how many blocks are chained in sequence. dim never reaches the blocks -- it only fixes the input/output size of the surrounding proj_up (dim -> hidden_dim) and proj_down (hidden_dim -> dim) projections, so the blocks' width and count are independent of dim.

Attributes:

Name	Type	Description
`dim`	`int`	The dimensionality of the embeddings the encoder operates on. Sets only the input/output size of `proj_up`/`proj_down`, not the residual blocks' internal width.
`hidden_dim`	`int`	The dimensionality of the encoder's internal representation. Passed directly as the `hidden_dim` of every `ResidualMLPBlock` in the stack.
`n_layers`	`int`	The number of residual MLP blocks. Equal to the number of `ResidualMLPBlock` instances chained between `proj_up` and `proj_down`.
`use_skip`	`bool`	Whether the encoder adds a global residual connection (`z + MLP(z)`) or returns the MLP's output directly (`MLP(z)`).
`identity_init`	`bool`	Whether the encoder starts as an identity map.

ResidualMLPBlock ¶

Bases: Module

A standard MLP block with a skip connection.

The final linear layer is zero-initialized so the block starts as an identity map.

Methods:

Name	Description
`__init__`	Build the block's Linear-LayerNorm-GELU-Linear MLP, zero-initializing the final layer so the block starts as an identity map.
`forward`	Apply the block to `x`.
`reset_parameters`	Initialize the block's MLP, then zero-initialize the final linear layer so the block starts as an identity map.

init ¶

__init__(hidden_dim: int) -> None

Build the block's Linear-LayerNorm-GELU-Linear MLP, zero-initializing the final layer so the block starts as an identity map.

Parameters:

Name	Type	Description	Default
`hidden_dim` ¶	`int`	The dimensionality of the block's input, output, and internal representation.	required

forward ¶

forward(x: Tensor) -> torch.Tensor

Apply the block to x.

Parameters:

Name	Type	Description	Default
`x` ¶	`Tensor`	Input tensor of shape `(..., hidden_dim)`.	required

Returns:

Type	Description
`torch.Tensor`	A tensor of the same shape as `x`.

reset_parameters ¶

reset_parameters() -> None

Initialize the block's MLP, then zero-initialize the final linear layer so the block starts as an identity map.

A no-op while the block's parameters are still on the meta device: the initialization below is real (dispatched) computation, not a free no-op, so it is skipped there and must be re-run via an explicit call to this method after materializing real storage (e.g. via to_empty), matching the convention used throughout noise_layer/.

privacy_encoder

PrivacyEncoder ¶

init ¶

`dim` ¶

`hidden_dim` ¶

`n_layers` ¶

`seed` ¶

`use_skip` ¶

`identity_init` ¶

deserialize `classmethod` ¶

`state` ¶

forward ¶

`input` ¶

`noise_mask` ¶

from_pretrained `classmethod` ¶

`path` ¶

`map_location` ¶

get_config ¶

reset_parameters ¶

save_pretrained ¶

`path` ¶

serialize ¶

PrivacyEncoderConfig ¶

ResidualMLPBlock ¶

init ¶

`hidden_dim` ¶

forward ¶

`x` ¶

reset_parameters ¶

privacy_encoder

PrivacyEncoder ¶

__init__ ¶

dim ¶

hidden_dim ¶

n_layers ¶

seed ¶

use_skip ¶

identity_init ¶

deserialize classmethod ¶

state ¶

forward ¶

input ¶

noise_mask ¶

from_pretrained classmethod ¶

path ¶

map_location ¶

get_config ¶

reset_parameters ¶

save_pretrained ¶

path ¶

serialize ¶

PrivacyEncoderConfig ¶

ResidualMLPBlock ¶

__init__ ¶

hidden_dim ¶

forward ¶

x ¶

reset_parameters ¶

init ¶

`dim` ¶

`hidden_dim` ¶

`n_layers` ¶

`seed` ¶

`use_skip` ¶

`identity_init` ¶

deserialize `classmethod` ¶

`state` ¶

`input` ¶

`noise_mask` ¶

from_pretrained `classmethod` ¶

`path` ¶

`map_location` ¶

`path` ¶

init ¶

`hidden_dim` ¶

`x` ¶