Fair learning with frozen Gaussianization flows — design doc

Fair learning with frozen Gaussianization flows¶

1. TL;DR¶

Replace the CKA fairness penalty in keras-fairkl’s FairModelWrapper — a Keras port of the fair-kernel-learning idea of Pérez-Suay et al. (2017) — with a family of fairness losses built from a frozen Gaussianization flow. The flow is trained once on an auxiliary dataset, its weights are frozen, and it is then used as a differentiable Gaussian-space probe inside the downstream task’s optimisation loop. Gaussianisation lets us replace bandwidth-tuned RBF kernels (CKA, HSIC) with closed-form, parametric, scale-invariant penalties — the flow absorbs the kernel choice.

Three penalties, in order of strictness:

All three are differentiable w.r.t. the downstream model parameters and plug into FairModelWrapper via its fairness_loss=... argument.

2. The mental model — three pictures¶

2.1 Stage 1 (one-time): pretrain the probes¶

2.2 Stage 2 (every step): the fair training loop¶

The trick — and the load-bearing claim of this whole experiment — is that the flow’s weights are frozen but the flow’s input is the predictor’s output, so gradients still propagate from the loss back through $T_z$ to θ.

2.3 Where each loss penalises¶

3. Why Gaussianisation helps (the one-paragraph theory)¶

A Gaussianization flow $T: \mathbb{R}^d \to \mathbb{R}^d$ — in the lineage of Chen & Gopinath (2000), Laparra et al. (2011), and Meng et al. (2020) — is a smooth diffeomorphism with smooth inverse, so it preserves all statistical dependence: $T(Z) \perp T(Q) \iff Z \perp Q$. What it changes is the shape of the marginals. After training, each marginal of $T(X)$ is approximately $\mathcal{N}(0, 1)$. Three consequences:

1. Bandwidth-free dependence measures. CKA and HSIC need a kernel bandwidth; the “right” bandwidth depends on the scale of the data. In Gaussianised space the scale is fixed at 1, so a linear kernel (or a unit-bandwidth RBF) suffices. The flow absorbs the bandwidth choice into its mixture-CDF parameters during pretraining.

2. Gaussian-joint assumption becomes nearly free. Closed-form MI ($-\tfrac{1}{2}\log\det(I - CC^\top)$) requires assuming the joint $(Z, Q)$ is Gaussian. On raw data this is wildly wrong. After marginal Gaussianisation it is much closer to true — the marginals are exact, only the copula remains non-Gaussian — so the closed-form MI estimate becomes a usable surrogate.

3. Compatibility with frozen-flow autodiff. All flow components (MixtureCDFGaussianization, Householder) are smooth in their inputs. Stopping gradient on the flow’s weights does not stop gradient flow through the flow’s inputs, so the predictor’s parameters θ still receive a gradient signal from $\nabla_z \mathcal{L}_{\text{fair}}(T_z(z), T_q(q))$ via the chain rule.

The flow is therefore exactly the right thing to freeze: a fixed, smooth, scale-normalising, differentiable preprocessor that turns “measure non-linear dependence in the data” into “measure linear dependence between near-Gaussian variables.”

4. Mathematical formulation¶

Let $z = f_\theta(X) \in \mathbb{R}^{d_z}$ be the predictor output and $q \in \mathbb{R}^{d_q}$ the sensitive attribute. Define $Z = T_z(z)$, $Q = T_q(q)$ in Gaussianised space, and let $C = \widehat{\mathrm{Cov}}(Z, Q)$, $S_z = \widehat{\mathrm{Cov}}(Z, Z)$, $S_q = \widehat{\mathrm{Cov}}(Q, Q)$ denote sample (cross-)covariances on a batch of size $n$.

4.1 G-XCOV — linear-CKA in Gaussianised space¶

Equation ((1)) is exact linear CKA Cortes et al., 2012Kornblith et al., 2019 applied to the Gaussianised features. In $d_z = d_q = 1$ it collapses to $\rho^2$ where ρ is the Gaussianised cross-correlation. The un-normalised numerator $\lVert C \rVert_F^2$ is identically HSIC with linear kernels Gretton et al., 2005 on the Gaussianised features — so this single loss covers both “linear CKA in Gaussianised space” and “HSIC with linear kernels in Gaussianised space” depending on whether you toggle normalize.

Bounded, smooth, second-moment only. Gradient $\partial \rho^2 / \partial \rho = 2\rho$ is bounded — the loss is gentle near perfect dependence.

4.2 G-MI — closed-form Gaussian mutual information¶

If $(Z, Q)$ were jointly Gaussian with standardised marginals, mutual information has the Gel’fand–Yaglom closed form Gel'fand & Yaglom, 1957:

After Gaussianisation $S_z \approx I_{d_z}$ and $S_q \approx I_{d_q}$, so ((2)) simplifies to

In $d_z = d_q = 1$ ((3)) is $-\tfrac{1}{2}\log(1 - \rho^2)$. The gradient $\partial \mathcal{L} / \partial \rho = \rho / (1 - \rho^2)$ diverges as $\rho \to 1$, so G-MI is much sharper than G-XCOV at high dependence. We clip the eigenvalues of $I - CC^\top$ at a small $\varepsilon$ for numerical safety; this caps the loss at $-\tfrac{d}{2}\log\varepsilon$.

The closed-form requires the joint to be Gaussian after Gaussianisation. Marginal Gaussianisation gets us most of the way there, but the residual copula is still arbitrary — so G-MI underestimates true MI whenever the dependence has structure beyond second-order correlation (quadratic-in-$Z$, XOR-style, multi-modal). That gap is exactly what G-TC closes.

4.3 G-TC — total correlation under a frozen joint flow¶

Pretrain a joint flow $T_{zq}: \mathbb{R}^{d_z + d_q} \to \mathbb{R}^{d_z + d_q}$ on the empirical product distribution of the baseline data: draw $(z^{(0)}_i, q_{\pi(i)})$ where π is a random permutation. By construction these pairs are independent, so a well-fit $T_{zq}$ Gaussianises independent draws to $\mathcal{N}(0, I_{d_z + d_q})$. Freeze $T_{zq}$.

At downstream training time, evaluate the same frozen $T_{zq}$ on the actual (potentially dependent) pair $(z, q)$:

When $(z, q)$ is independent like the baseline, $T_{zq}(z, q) \sim \mathcal{N}(0, I)$ and ((4)) equals the entropy of $\mathcal{N}(0, I_{d_z + d_q})$ (a constant in θ). When $(z, q)$ carries dependence, $T_{zq}$ no longer Gaussianises the joint and the NLL is strictly larger.

By a change-of-variables argument, this NLL difference is exactly the KL divergence between $p(z, q)$ and $p(z)\,p(q)$ — i.e. the mutual information; see Watanabe (1960) for the total-correlation framing. Unlike G-MI it does not assume the joint is Gaussian: the flow itself learns the copula during pretraining. The price is needing a richer pretraining stage.

4.4 Comparison table¶

4.5 Deferred (stretch)¶

• G-HSIC-RBF. HSIC with a unit-bandwidth RBF kernel in Gaussianised space. Identical to existing CKA code but with $T_z(x)$ in place of $x$. Useful as an ablation to separate “flow as preprocessor” from “linear vs RBF after the flow”.
• DR-MI. True mutual information via Monte-Carlo marginalisation through the joint flow. Compare with the neural estimator of Belghazi et al. (2018). Expensive; G-TC already captures the same signal at lower cost.

5. Hypotheses we’re testing¶

Each loss family makes a falsifiable prediction about what kind of dependence it can suppress, and at what cost in accuracy.

6. Experiment design¶

6.1 Two-stage pipeline (reprise of the ASCII diagrams)¶

Stage 1 — pretrain + freeze.

• Train $T_z$ on the predictor-output distribution of an unconstrained baseline (e.g. for classification, sigmoid probabilities of the baseline MLP — not the raw binary labels, otherwise the flow lives on a 2-point support that is off-support of the actual predictions).
• Train $T_q$ on the marginal of $q$.
• Train $T_{zq}$ on the shuffled product distribution of the same baseline data — independent pairs by construction.
• Freeze all weights.

Stage 2 — fair downstream training.

• Drop in FairModelWrapper(base, mu=μ, fairness_loss=...) with any of the three new losses.
• compile(..., loss="...") as usual. The wrapper handles the dict packing {"x": X, "q": q}.

6.2 What gets penalised, where the gradient goes¶

A short walk through one optimiser step. Solid arrows are the forward pass; dashed arrows are the backward (autodiff) pass, each labelled with the chain-rule factor it carries.

The key chain-rule fact, read off the dashed path above: stop_gradient (or trainable=False) blocks gradients into the parameters of the flow, but does not block gradients into its inputs — the mixture-CDF Jacobian $\partial T_z(z)/\partial z$ is still smooth and non-zero, so the predictor still receives a fairness signal. Without this property the whole scheme collapses. $T_z$, $T_q$, and $T_{zq}$ themselves receive no gradient: they have zero trainable_weights.

6.3 File layout¶

projects/gaussianization/
├── src/gaussianization/fair/
│   ├── __init__.py          # public API
│   ├── losses.py            # GaussianizedXCovLoss / MutualInfoLoss / TotalCorrelationLoss
│   ├── pretrain.py          # fit_and_freeze, fit_and_freeze_joint
│   ├── freeze.py            # freeze_flow helper
│   └── metrics.py           # numpy fairness eval metrics
├── tests/test_fair.py       # 15 tests including closed-form checks
├── notebooks/fair_gauss/
│   ├── 05_fair_gauss_pretrain.ipynb
│   ├── 06_fair_gauss_synthetic.ipynb   # G-XCOV + G-MI + G-TC + CKA
│   ├── 07_fair_gauss_adult.ipynb       # same on UCI Adult
│   └── _style.py
└── docs/fair_gaussianization_experiment.md   # this file

7. Datasets¶

8. Evaluation plan¶

For every (dataset, loss, μ, seed) combination report:

• Predictive metrics. RMSE / R² (regression). Accuracy, ROC-AUC, log-loss (classification).
• Fairness metrics (numpy-side, neutral judge). Demographic-parity difference, equalized-odds difference, |Pearson($\hat y, q$)|. These are computed at evaluation time, not used as the training loss, so they are independent of the training penalty.
• Diagnostic: training-time loss value. Each loss’s own value through training (so we can see G-MI saturate at the eps clip if it does, etc.).

Comparison axes

1. Loss family: cka | g_xcov | g_mi | g_tc × $\mu \in \{0, 0.1, 0.5, 2, 10, 50, 200\}$.
2. Flow depth (num_blocks ∈ {2, 4, 8}) — does a deeper joint flow improve G-TC on the quadratic-dependence test (H3)?
3. Pretraining set — same-as-task vs held-out vs i.i.d. Gaussian (the latter collapses G-XCOV to vanilla cross-cov; control).
4. Batch size — G-MI and G-TC may be more batch-sensitive than G-XCOV (eigh + flow forward).

Statistical rigour. 3 seeds for the headline figures, 5 for any “X beats Y at matched fairness” claim, paired-bootstrap CIs.

Magnitude calibration follow-up. Each loss has a different natural scale. To make the loss-family axis a fair comparison, scale μ per-family by the loss value at the unconstrained baseline, so that effective μ puts everyone at the same relative pressure. Report both raw-μ and calibrated-μ Pareto fronts.

9. Risks & open questions¶

Open questions for follow-ups.

• Does fine-tuning the marginal flows during downstream training (a light EMA update) actually help, or does it leak into the predictor’s gradient and corrupt the fairness signal?
• G-TC requires a joint flow. Can a single $T_{zq}$ pretrained once serve every downstream task on the same data, or does it need to be re-trained per architecture?
• For multi-class sensitive attributes (race in COMPAS), does the one-hot encoding of $q$ work, or do we need a categorical flow head?

10. Milestones¶

Appendix A — Minimum working example¶

import os
os.environ.setdefault("KERAS_BACKEND", "jax")

import keras
import numpy as np

from fairkl.models import FairModelWrapper
from gaussianization.fair import (
    GaussianizedMutualInfoLoss,
    GaussianizedXCovLoss,
    GaussianizedTotalCorrelationLoss,
    fit_and_freeze,
    fit_and_freeze_joint,
)

# Synthetic data with q as a feature
rng = np.random.default_rng(0)
n = 4000
q = rng.binomial(1, 0.5, n).astype("float32")
x = rng.standard_normal((n, 3)).astype("float32")
y = (np.tanh(x[:, 0]) + 0.5 * x[:, 1] + 3 * q
     + 0.1 * rng.standard_normal(n)).astype("float32")
X = np.concatenate([x, q.reshape(-1, 1)], axis=1)

# Stage 1: pretrain + freeze probes
flow_y, _  = fit_and_freeze(y.reshape(-1, 1), num_blocks=4, epochs=40, seed=0)
flow_q, _  = fit_and_freeze(q.reshape(-1, 1), num_blocks=2, epochs=40, seed=0)
flow_yq, _ = fit_and_freeze_joint(y, q, num_blocks=4, epochs=40, seed=0)

# Stage 2: pick a loss and train
mlp = keras.Sequential([keras.layers.Dense(32, "relu"),
                        keras.layers.Dense(32, "relu"),
                        keras.layers.Dense(1)])
fair = FairModelWrapper(
    mlp, mu=2.0,
    fairness_loss=GaussianizedMutualInfoLoss(flow_z=flow_y, flow_q=flow_q),
    # or: GaussianizedXCovLoss(flow_z=flow_y, flow_q=flow_q),
    # or: GaussianizedTotalCorrelationLoss(joint_flow=flow_yq),
)
fair.compile(optimizer="adam", loss="mse")
fair.fit(X, y, q=q, epochs=40, batch_size=256)