image_similarity
Module for image-similarity privacy losses.
Provides differentiable losses that quantify how much structural and linear-dependence information remains between a noisy (Stained Glass-transformed) image and its clean counterpart, after a battery of non-learned probe transforms (clamp, min-max, Sobel, Gaussian blur, gamma, wavelet denoise, z-score-sigmoid). Use these as a privacy regularizer when training image cloaks.
Design notes — alternatives considered.
Pixel-space metrics (L1 / L2). Direct |noisy - clean| or (noisy - clean)^2
were considered first. They are trivially defeated by an attacker that learns a linear
de-noising filter, and they treat structural and textural information identically. We
chose structural similarity metrics instead because they capture perceptual differences
the way a downstream classifier or vision encoder would.
Single-scale MS-SSIM only. Earlier prototypes compared the bare denormalized
noisy_01 against clean_01 with a single MS-SSIM call. In experiments this term
drove MS-SSIM to very small values while the obfuscation remained weak — an attacker
could still recover recognizable structure from the noisy image via simple
post-processing (e.g. min-max normalization, gamma correction, denoising). The current
multi-probe formulation evaluates MS-SSIM (and squared Pearson) after applying
each candidate probe transform on both sides, so the loss penalizes residue that
survives those probes. The squared-Pearson family then catches residual linear
dependence in directions that MS-SSIM's structural measure only partially penalizes
(Pearson is affine-invariant, so it covers a complementary failure mode).
Other perceptual metrics (LPIPS, DISTS, FSIM, PSNR). Not included to keep the loss free of learned components (no learned reference model whose weights could be reverse-engineered) and lightweight. They remain candidates for future additions if a use case appears.
Functions:
| Name | Description |
|---|---|
multi_attacker_privacy_loss |
Compute a 15-component privacy loss against non-learned image probes. |
multi_probe_privacy_loss |
Compute a 15-component privacy loss against non-learned image probes. |
pearson_sq_mean |
Mean across the batch of the squared Pearson correlation between |
safe_ms_ssim |
Compute functional MS-SSIM between |
Attributes:
| Name | Type | Description |
|---|---|---|
DEFAULT_AGGREGATION_TEMPERATURE |
Final[float]
|
Default softmax temperature for |
DEFAULT_GAMMA_ATTACK_VALUES |
Final[tuple[float, ...]]
|
Gamma values used by the gamma-correction probe family (covers both dark- and bright-region recovery). |
DEFAULT_GAUSSIAN_BLUR_SIGMA |
Final[float]
|
Sigma of the Gaussian-blur probe, in pixels of the [0, 1]-denormalized image space. |
DEFAULT_WAVELET_SOFT_THRESHOLD |
Final[float]
|
Soft-threshold magnitude applied to the detail wavelet coefficients in the wavelet-denoise probe. |
MS_SSIM_COMPONENT_NAMES |
Final[tuple[str, ...]]
|
Names of the nine MS-SSIM privacy components produced by |
MS_SSIM_STACK_KEY |
Final[str]
|
Key of the optional un-aggregated MS-SSIM score stack, present only when |
MS_SSIM_TOTAL_KEY |
Final[str]
|
Key of the aggregated MS-SSIM total in the |
PEARSON_COMPONENT_NAMES |
Final[tuple[str, ...]]
|
Names of the six squared-Pearson privacy components produced by |
PEARSON_STACK_KEY |
Final[str]
|
Key of the optional un-aggregated squared-Pearson score stack, present only when |
PEARSON_TOTAL_KEY |
Final[str]
|
Key of the aggregated squared-Pearson total in the |
PearsonChannelReduction |
How |
|
ProbeAggregation |
How |
DEFAULT_AGGREGATION_TEMPERATURE
module-attribute
¶
Default softmax temperature for "softmax" probe aggregation.
Chosen to be on the order of the spread of the per-probe scores rather than of their range. Because every component lies in
[0, 1], a temperature of 1.0 makes the softmax weights nearly uniform and the aggregate indistinguishable from "mean"; 0.1
concentrates most of the weight on the leading probes while keeping every weight strictly positive.
DEFAULT_GAMMA_ATTACK_VALUES
module-attribute
¶
Gamma values used by the gamma-correction probe family (covers both dark- and bright-region recovery).
DEFAULT_GAUSSIAN_BLUR_SIGMA
module-attribute
¶
Sigma of the Gaussian-blur probe, in pixels of the [0, 1]-denormalized image space.
DEFAULT_WAVELET_SOFT_THRESHOLD
module-attribute
¶
Soft-threshold magnitude applied to the detail wavelet coefficients in the wavelet-denoise probe.
MS_SSIM_COMPONENT_NAMES
module-attribute
¶
MS_SSIM_COMPONENT_NAMES: Final[tuple[str, ...]] = (
"rgb_clamp",
"gray_clamp",
"rgb_minmax",
"gray_minmax",
"sobel_gray",
"blur_gray",
"standardize_gray",
"gamma_gray",
"wavelet_gray",
)
Names of the nine MS-SSIM privacy components produced by multi_probe_privacy_loss.
Also fixes the row order of the optional ms_ssim_stack tensor.
MS_SSIM_STACK_KEY
module-attribute
¶
Key of the optional un-aggregated MS-SSIM score stack, present only when return_component_stack=True.
MS_SSIM_TOTAL_KEY
module-attribute
¶
Key of the aggregated MS-SSIM total in the multi_probe_privacy_loss result.
PEARSON_COMPONENT_NAMES
module-attribute
¶
PEARSON_COMPONENT_NAMES: Final[tuple[str, ...]] = (
"pearson_rgb",
"pearson_gray",
"pearson_sobel",
"pearson_blur",
"pearson_gamma",
"pearson_wavelet",
)
Names of the six squared-Pearson privacy components produced by multi_probe_privacy_loss.
Also fixes the row order of the optional pearson_stack tensor.
PEARSON_STACK_KEY
module-attribute
¶
Key of the optional un-aggregated squared-Pearson score stack, present only when return_component_stack=True.
PEARSON_TOTAL_KEY
module-attribute
¶
Key of the aggregated squared-Pearson total in the multi_probe_privacy_loss result.
PearsonChannelReduction
module-attribute
¶
PearsonChannelReduction = Literal['pooled', 'max', 'mean']
How pearson_sq_mean treats the channel axis.
"max": correlate each channel separately and keep the worst channel. The default for image inputs, and the only mode that detects a transform mapping channels differently."mean": correlate each channel separately and average. Channel-aware, but a single leaking channel is diluted by the others."pooled": flatten every non-batch axis together and take one correlation per image.
"pooled" is not a recommended setting and is retained only for continuity. It is rank-agnostic, so it remains the default of the
low-level pearson_sq_mean where a channel axis may not exist, but as a privacy statistic it is known-broken: it reports 0.0556 on an
exactly invertible per-channel transform that "max" correctly reports as 1.0000. It survives so that numbers from checkpoints trained
before this distinction existed can still be reproduced. Once that comparison is no longer needed it should be removed rather than left
as a footgun — a caller reaching for it on image data is almost always making a mistake.
ProbeAggregation
module-attribute
¶
ProbeAggregation = Literal['mean', 'softmax', 'logsumexp']
How multi_probe_privacy_loss reduces a per-probe score
family to a single scalar total.
"mean": the unweighted arithmetic mean. Every probe receives an equal1 / nshare of the gradient."softmax": a detached softmax-weighted mean —sum(softmax(scores.detach() / temperature) * scores). Concentrates the gradient on the probes currently scoring highest, i.e. the ones the cloak is failing against. Because the weights are detached, the value returned is not a function whose true gradient this is — the value's own minimum is generally at a point where this gradient is nonzero, so gradient descent on it does not converge to a stationary point of the value it is minimizing."logsumexp":temperature * logsumexp(scores / temperature) - temperature * log(n). Its true (non-detached) gradient is exactlysoftmax(scores / temperature)— the same weights"softmax"uses — so the two modes agree pointwise on the gradient signal everywhere. Unlike"softmax", this is a genuine differentiable function of the scores, so the reported gradient is the derivative of the reported value. Shares"softmax"'s probe-concentrating behavior without its detached-gradient inconsistency.
multi_attacker_privacy_loss
¶
multi_attacker_privacy_loss(
noisy_01: Tensor,
clean_01: Tensor,
*,
gamma_values: Sequence[float] = (0.5, 2.0),
blur_sigma: float = 2.0,
wavelet_threshold: float = 0.1,
compute_ms_ssim: bool = True,
compute_pearson: bool = True,
aggregation: ProbeAggregation = "mean",
aggregation_temperature: float = 0.1,
return_component_stack: bool = True,
pearson_channel_reduction: PearsonChannelReduction = "max"
) -> dict[str, torch.Tensor]
Compute a 15-component privacy loss against non-learned image probes.
Alias retained for backwards compatibility. Forwards every argument unchanged to
multi_probe_privacy_loss,
which is the supported entry point.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
|
Tensor
|
Stained-Glass-transformed image tensor of shape |
required |
|
Tensor
|
Clean image tensor of shape |
required |
|
Sequence[float]
|
Gamma values used for the |
(0.5, 2.0)
|
|
float
|
Sigma of the Gaussian blur applied for the |
2.0
|
|
float
|
Soft-threshold magnitude applied to wavelet detail bands for the
|
0.1
|
|
bool
|
If |
True
|
|
bool
|
If |
True
|
|
ProbeAggregation
|
How to reduce each probe family to its aggregate key. |
'mean'
|
|
float
|
Softmax temperature used when |
0.1
|
|
bool
|
If |
True
|
|
PearsonChannelReduction
|
How the squared-Pearson family reduces across colour channels. |
'max'
|
Returns:
| Type | Description |
|---|---|
dict[str, torch.Tensor]
|
The dict returned by |
dict[str, torch.Tensor]
|
Raises:
| Type | Description |
|---|---|
ValueError
|
If |
Added in version v3.39.0. 15-component multi-probe privacy loss.
Deprecated since version v3.70.1. Renamed to `multi_probe_privacy_loss`; the transforms are non-learned probes, not adversaries.
multi_probe_privacy_loss
¶
multi_probe_privacy_loss(
noisy_01: Tensor,
clean_01: Tensor,
*,
gamma_values: Sequence[float] = (0.5, 2.0),
blur_sigma: float = 2.0,
wavelet_threshold: float = 0.1,
compute_ms_ssim: bool = True,
compute_pearson: bool = True,
aggregation: ProbeAggregation = "mean",
aggregation_temperature: float = 0.1,
return_component_stack: bool = True,
pearson_channel_reduction: PearsonChannelReduction = "max"
) -> dict[str, torch.Tensor]
Compute a 15-component privacy loss against non-learned image probes.
Both inputs must already be in the [0, 1] range. The height and width may be arbitrary: both
tensors are first centered and reflect-padded (see
reflect_pad_to_size) up to
the minimum size MS-SSIM requires (176px) and up to even parity (required by the
single-level Haar DWT used by the wavelet probes). This is a no-op for inputs already at least 176px and even on both
axes. The noisy image is compared against the
clean image under fifteen different probe transforms — nine MS-SSIM terms that
measure structural similarity and six squared-Pearson terms that measure residual
linear dependence. Each family is reduced to one aggregate key by aggregation.
For "lossy" transforms (blur, sobel, gamma, wavelet, standardize), the same
transform is applied to the clean side too so the loss floor stays at a clean 1.0
and the gradient signal is purely driven by noise residue.
Structural similarity terms (MS-SSIM). Contrast-recovery probes compared
directly against clean/clean_gray:
rgb_clamp: display-level RGB — clamp noisy RGB to[0, 1].gray_clamp: display-level luma — clamp noisy luma.rgb_minmax: tensor-level RGB — min-max normalize noisy RGB.gray_minmax: tensor-level luma — reduce to luma then min-max normalize.
Lossy-transform probes (T applied to both sides, compared in the transformed space):
sobel_gray: Sobel gradient magnitude of luma (min-max normalized).blur_gray: Gaussian low-pass of luma atblur_sigmapixels.standardize_gray: per-image z-score + sigmoid of luma.gamma_gray: mean MS-SSIM acrossgamma_values.wavelet_gray: single-level Haar wavelet soft-threshold denoising.
Linear-dependence terms (squared Pearson). Pearson is affine-invariant, so applying it to min-max / other affine transforms is redundant with applying it to the raw luma — those variants are intentionally omitted. The included transforms all break linearity:
pearson_rgb: flattened RGB, usingclamp(noisy_rgb)vsclean_rgb.pearson_gray: flattened luma, usingclamp(noisy_gray)vsclean_gray.pearson_sobel: Sobel magnitudes of both sides.pearson_blur: Gaussian-blurred luma on both sides.pearson_gamma: mean overgamma_valuesof squared Pearson after gamma correction on both sides.pearson_wavelet: wavelet-denoised luma on both sides.
If noisy_01 contains NaN/Inf values, returns a dict of zeros (and emits a warning)
rather than propagating invalid values through the backward pass. The root cause
(e.g. an unbounded mean estimator output) should be addressed at the noise layer.
Privacy: this loss handles user-provided image data. Callers should ensure that the
invocation context does not log or persist clean_01 or noisy_01.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
|
Tensor
|
Stained-Glass-transformed image tensor of shape |
required |
|
Tensor
|
Clean image tensor of shape |
required |
|
Sequence[float]
|
Gamma values used for the |
(0.5, 2.0)
|
|
float
|
Sigma of the Gaussian blur applied for the |
2.0
|
|
float
|
Soft-threshold magnitude applied to wavelet detail bands for the
|
0.1
|
|
bool
|
If |
True
|
|
bool
|
If |
True
|
|
ProbeAggregation
|
How to reduce each probe family to its aggregate key. |
'mean'
|
|
float
|
Softmax temperature used when |
0.1
|
|
PearsonChannelReduction
|
How the squared-Pearson family treats the colour-channel axis. Defaults to |
'max'
|
|
bool
|
If |
True
|
Returns:
| Type | Description |
|---|---|
dict[str, torch.Tensor]
|
A dict with one key per component listed above, plus two aggregate keys: |
dict[str, torch.Tensor]
|
|
dict[str, torch.Tensor]
|
|
dict[str, torch.Tensor]
|
All of these are scalar tensors in |
dict[str, torch.Tensor]
|
dict also holds |
Raises:
| Type | Description |
|---|---|
ValueError
|
If |
Example
clean = torch.rand(2, 3, 176, 176) noisy = (clean + torch.randn_like(clean)).clamp(0.0, 1.0) components = multi_probe_privacy_loss(noisy, clean, return_component_stack=True) components["ms_ssim_stack"].shape, components["pearson_stack"].shape (torch.Size([9]), torch.Size([6])) bool( ... torch.allclose( ... components["ms_ssim_stack"].mean(), components["ms_ssim_total"] ... ) ... ) True
"softmax"leans toward the strongest probe, so its total is never below the mean.¶weighted = multi_probe_privacy_loss(noisy, clean, aggregation="softmax") bool(weighted["ms_ssim_total"] >= components["ms_ssim_total"]) True
Changed in version v3.72.0: Added ``aggregation="logsumexp"``
Changed in version v3.68.0: Added the `aggregation`, `aggregation_temperature`, and `return_component_stack` arguments. Defaults preserve the previous unweighted-mean behavior.
Changed in version v3.68.0: The squared-Pearson family now correlates each colour channel separately and keeps the worst, via the new `pearson_channel_reduction` argument defaulting to `"max"`. The previous pooled correlation was structurally blind to any transform that maps channels differently: on an exactly invertible per-channel compression it reports 0.0556 where the per-channel form reports 1.0000. `pearson_rgb` and `pearson_total` therefore change value and are **not** comparable across this boundary; the five single-channel components are unaffected. Pass `pearson_channel_reduction="pooled"` to restore the old numbers.
Added in version v3.39.0. 15-component multi-probe privacy loss.
pearson_sq_mean
¶
pearson_sq_mean(
pred: Tensor,
target: Tensor,
*,
eps: float = 1e-12,
channel_reduction: PearsonChannelReduction = "pooled"
) -> torch.Tensor
Mean across the batch of the squared Pearson correlation between pred and target.
Pearson correlation is affine-invariant (invariant to shift/scale of either argument),
so this term captures residual linear dependence in directions that MS-SSIM's
structural measure only partially penalizes. Squaring maps the result into [0, 1]
with a smooth minimum at 0 (no linear dependence) and a smooth maximum at 1 (perfect
linear dependence up to sign), keeping gradients finite across the entire range —
unlike |pearson| which is non-differentiable at zero.
Both the mean-subtraction and the variance terms are kept attached to the graph (not
detached), so the gradient with respect to pred reflects the true Pearson
derivative.
The channel axis matters, and pooling it away opens a hole. Affine invariance holds only for a single affine map. Flattening
the channels together asks whether one shared map relates the two tensors — so a transform applying a different map per channel
breaks the pooled correlation even though every channel remains perfectly recoverable on its own. Measured on a transform that
compresses red and green into [0.961, 1.0] while scaling blue to near zero — exactly invertible, so an ideal statistic reads
1.0 — "pooled" reports 0.0556 while "max" reports 1.0000. That is not a corner case: per-channel curves are the
first thing any photo editor offers, and a cloak that compresses channels unequally is invisible to "pooled".
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
|
Tensor
|
Tensor of shape |
required |
|
Tensor
|
Tensor of the same shape as |
required |
|
float
|
Floor applied to the product of variances before the sqrt in the Pearson denominator, so the sqrt's gradient stays finite for near-constant inputs. |
1e-12
|
|
PearsonChannelReduction
|
How to treat the channel axis. Defaults to |
'pooled'
|
Returns:
| Type | Description |
|---|---|
torch.Tensor
|
Scalar tensor in |
Raises:
| Type | Description |
|---|---|
ValueError
|
If |
Example
torch.manual_seed(0) # doctest: +ELLIPSIS
clean = torch.rand(1, 3, 32, 32) An exactly invertible transform that maps each channel differently.¶
noisy = torch.stack( ... [ ... 0.961 + 0.039 * clean[:, 0], ... 0.961 + 0.039 * clean[:, 1], ... 0.02 * clean[:, 2], ... ], ... dim=1, ... )
Pooling the channels together hides it; correlating per channel does not.¶
bool(pearson_sq_mean(noisy, clean, channel_reduction="pooled") < 0.2) True bool(pearson_sq_mean(noisy, clean, channel_reduction="max") > 0.99) True
Added in version v3.39.0. Image-similarity privacy loss family.
safe_ms_ssim
¶
safe_ms_ssim(
pred: Tensor,
target: Tensor,
*,
data_range: float = 1.0,
normalize: Literal["relu", "simple"] | None = "relu"
) -> torch.Tensor
Compute functional MS-SSIM between pred and target, replacing NaN/Inf with zero.
Wraps torchmetrics.functional.image.multiscale_structural_similarity_index_measure.
Returns 0 (and emits a warning) on NaN/Inf so a single bad batch does not poison
training via a NaN gradient.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
|
Tensor
|
Predicted image tensor of shape |
required |
|
Tensor
|
Target image tensor of shape |
required |
|
float
|
The dynamic range of the inputs, forwarded to torchmetrics. The default of |
1.0
|
|
Literal['relu', 'simple'] | None
|
Normalization mode for negative MS-SSIM contributions, forwarded to torchmetrics. |
'relu'
|
Returns:
| Type | Description |
|---|---|
torch.Tensor
|
Scalar MS-SSIM tensor with the same dtype/device as |
Added in version v3.39.0. Image-similarity privacy loss family.