Merge pull request #10 from EleutherAI/cdf

Add QuantileNormalizer class

.gitignore

@@ -158,3 +158,7 @@ cython_debug/
 # and can be added to the global gitignore or merged into this file. For a more nuclear
 # option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# Weights & Biases logs
+*.ckpt
+wandb/

concept_erasure/__init__.py

@@ -3,12 +3,15 @@
 from .leace import ErasureMethod, LeaceEraser, LeaceFitter
 from .oracle import OracleEraser, OracleFitter
 from .quadratic import QuadraticEditor, QuadraticEraser, QuadraticFitter
+from .quantile import QuantileNormalizer, cdf, icdf
 from .shrinkage import optimal_linear_shrinkage
 from .utils import assert_type
 
 __all__ = [
  "assert_type",
+ "cdf",
  "groupby",
+ "icdf",
  "optimal_linear_shrinkage",
  "ConceptScrubber",
  "ErasureMethod",
@@ -20,4 +23,5 @@
  "QuadraticEditor",
  "QuadraticEraser",
  "QuadraticFitter",
+ "QuantileNormalizer",
 ]

concept_erasure/quadratic.py

@@ -51,8 +51,11 @@ class QuadraticEditor:
 
  def transport(self, x: Tensor, source_z: int, target_z: int) -> Tensor:
  """Transport `x` from class `source_z` to class `target_z`"""
+ x_ = x.flatten(1)
+
  T = self.ot_maps[source_z, target_z]
- return (x - self.class_means[source_z]) @ T.mH + self.class_means[target_z]
+ x_ = (x_ - self.class_means[source_z]) @ T.mH + self.class_means[target_z]
+ return x_.view_as(x)
 
  def __call__(self, x: Tensor, source_z: Tensor, target_z: int) -> Tensor:
  """Transport `x` from classes `source_z` to class `target_z`."""
@@ -146,7 +149,7 @@ def update_single(self, x: Tensor, z: int) -> "QuadraticFitter":
 
  return self
 
- def editor(self, device: str | None = None) -> QuadraticEditor:
+ def editor(self, device: str | torch.device | None = None) -> QuadraticEditor:
  """Quadratic editor for the concept."""
  sigma = self.sigma_xx
  device = device or sigma.device

concept_erasure/quantile.py

@@ -0,0 +1,105 @@
+import torch
+from torch import Tensor
+
+from .groupby import groupby
+
+
+def cdf(x: float | Tensor, q: Tensor) -> Tensor:
+ """Evaluate empirical CDF defined by quantiles `q` on `x`.
+
+ Args:
+ x: `[...]` Scalar or tensor of data points of arbitrary shape.
+ q: `[..., num_quantiles]` batch of quantiles. Must be sorted, and
+ should broadcast with `x` except for the last dimension.
+
+ Returns:
+ `[...]` Empirical CDF evaluated for each element of `x`.
+ """
+ n = q.shape[-1]
+ assert n > 2, "Must have at least two quantiles to interpolate."
+
+ # Approach used by SciPy interp1d with kind='previous'
+ # Shift x toward +inf by epsilon to appropriately handle ties
+ x = torch.nextafter(torch.as_tensor(x), q.new_tensor(torch.inf))
+ return torch.searchsorted(q, x, out_int32=True) / n
+
+
+def icdf(p: Tensor, q: Tensor) -> Tensor:
+ """(Pseudo-)inverse of the ECDF defined by quantiles `q`.
+
+ Returns the *smallest* `x` such that the ECDF of `x` is greater than or
+ equal to `p`.
+
+ NOTE: Strictly speaking, this function should return `-inf` when `p` is exactly
+ zero, because there is no smallest `x` such that `p(x) = 0`. But in practice we
+ want this function to always return a finite value, so we clip to the minimum
+ value in `q`.
+
+ Args:
+ x: `[...]` Tensor of data points of arbitrary shape.
+ q: `[..., num_quantiles]` batch of quantiles. Must be sorted, and
+ should broadcast with `x` except for the last dimension.
+
+ Returns:
+ `[...]` Empirical CDF evaluated for each element of `x`.
+ """
+ n = q.shape[-1]
+ assert n > 2, "Must have at least two quantiles to interpolate."
+
+ soft_ranks = torch.nextafter(p * n, p.new_tensor(0.0))
+ return q.gather(-1, soft_ranks.long())
+
+
+class QuantileNormalizer:
+ """Componentwise quantile normalization."""
+
+ lut: Tensor
+ """`[k, ..., num_bins]` batch of lookup tables."""
+
+ dim: int
+ """Dimension along which to group the data."""
+
+ def __init__(
+ self,
+ x: Tensor,
+ z: Tensor,
+ num_bins: int = 256,
+ dim: int = 0,
+ ):
+ # Efficiently get a view onto each class
+ grouped = groupby(x, z, dim=dim)
+ self.dim = dim
+
+ k = len(grouped.labels)
+ self.lut = x.new_empty([k, *x.shape[1:], num_bins])
+
+ grid = torch.linspace(0, 1, num_bins, device=x.device)
+ for i, grp in grouped:
+ self.lut[i] = grp.quantile(grid, dim=dim).movedim(0, -1)
+
+ @property
+ def num_bins(self) -> int:
+ return self.lut.shape[-1]
+
+ def cdf(self, z: int, x: Tensor) -> Tensor:
+ return cdf(x.movedim(0, -1), self.lut[z]).movedim(-1, 0)
+
+ def sample(self, z: int, n: int) -> Tensor:
+ lut = self.lut[z]
+
+ # Sample p from uniform distribution, then apply inverse CDF
+ p = torch.rand(*lut[..., 0].shape, n, device=lut.device)
+ return icdf(p, lut).movedim(-1, 0)
+
+ def transport(self, x: Tensor, source_z: Tensor, target_z: int) -> Tensor:
+ """Transport `x` from class `source_z` to class `target_z`"""
+ return (
+ groupby(x, source_z, dim=self.dim)
+ .map(
+ # Probability integral transform, followed by inverse for target class
+ lambda z, x: icdf(
+ cdf(x.movedim(0, -1), self.lut[z]), self.lut[target_z]
+ ).movedim(-1, 0)
+ )
+ .coalesce()
+ )

experiments/prediction_steering/models.py

@@ -0,0 +1,173 @@
+from itertools import pairwise
+from typing import Literal
+
+import pytorch_lightning as pl
+import torch
+import torchmetrics as tm
+import torchvision as tv
+from torch import nn
+from torch.optim import RAdam
+from torch.optim.lr_scheduler import CosineAnnealingLR
+
+
+class Mlp(pl.LightningModule):
+ def __init__(self, k, h=512, **kwargs):
+ super().__init__()
+ self.save_hyperparameters()
+
+ self.build_net()
+ self.train_acc = tm.Accuracy("multiclass", num_classes=k)
+ self.val_acc = tm.Accuracy("multiclass", num_classes=k)
+ self.test_acc = tm.Accuracy("multiclass", num_classes=k)
+
+ def build_net(self):
+ sizes = [3 * 32 * 32] + [self.hparams["h"]] * 4
+
+ self.net = nn.Sequential(
+ *[
+ MlpBlock(
+ in_dim,
+ out_dim,
+ device=self.device,
+ dtype=self.dtype,
+ residual=True,
+ act="gelu",
+ )
+ for in_dim, out_dim in pairwise(sizes)
+ ]
+ )
+ self.net.append(nn.Linear(self.hparams["h"], self.hparams["k"]))
+
+ def forward(self, x):
+ return self.net(x)
+
+ def training_step(self, batch, batch_idx):
+ x, y = batch
+
+ y_hat = self(x)
+ loss = torch.nn.functional.cross_entropy(y_hat, y)
+ self.log("train_loss", loss)
+
+ self.train_acc(y_hat, y)
+ self.log("train_acc", self.train_acc, on_epoch=True, on_step=False)
+ # Log the norm of the weights
+ fc = self.net[-1] if isinstance(self.net, nn.Sequential) else None
+ if isinstance(fc, nn.Linear):
+ self.log("weight_norm", fc.weight.data.norm())
+
+ return loss
+
+ def validation_step(self, batch, batch_idx):
+ x, y = batch
+
+ y_hat = self(x)
+ loss = torch.nn.functional.cross_entropy(y_hat, y)
+
+ self.val_acc(y_hat, y)
+ self.log("val_loss", loss)
+ self.log("val_acc", self.val_acc, prog_bar=True)
+ return loss
+
+ def test_step(self, batch, batch_idx):
+ x, y = batch
+
+ y_hat = self(x)
+ loss = torch.nn.functional.cross_entropy(y_hat, y)
+
+ self.test_acc(y_hat, y)
+ self.log("test_loss", loss)
+ self.log("test_acc", self.test_acc, prog_bar=True)
+ return loss
+
+ def configure_optimizers(self):
+ opt = RAdam(self.parameters(), lr=1e-4)
+ return [opt], [CosineAnnealingLR(opt, T_max=200)]
+
+
+class MlpMixer(Mlp):
+ def build_net(self):
+ from mlp_mixer_pytorch import MLPMixer
+
+ self.net = MLPMixer(
+ image_size=32,
+ channels=3,
+ patch_size=self.hparams.get("patch_size", 4),
+ num_classes=self.hparams["k"],
+ dim=512,
+ depth=6,
+ dropout=0.1,
+ )
+
+
+class ResNet(Mlp):
+ def build_net(self):
+ self.net = tv.models.resnet18(pretrained=False, num_classes=self.hparams["k"])
+
+
+class ViT(MlpMixer):
+ def build_net(self):
+ from vit_pytorch import ViT
+
+ self.net = ViT(
+ image_size=32,
+ patch_size=self.hparams.get("patch_size", 4),
+ num_classes=self.hparams["k"],
+ dim=512,
+ depth=6,
+ heads=8,
+ mlp_dim=512,
+ dropout=0.1,
+ emb_dropout=0.1,
+ )
+
+
+class MlpBlock(nn.Module):
+ def __init__(
+ self,
+ in_features: int,
+ out_features: int,
+ device=None,
+ dtype=None,
+ residual: bool = True,
+ *,
+ act: Literal["relu", "gelu"] = "relu",
+ norm: Literal["batch", "layer"] = "batch",
+ ):
+ super().__init__()
+
+ self.linear1 = nn.Linear(
+ in_features, out_features, bias=False, device=device, dtype=dtype
+ )
+ self.linear2 = nn.Linear(
+ out_features, out_features, bias=False, device=device, dtype=dtype
+ )
+ self.act_fn = nn.ReLU() if act == "relu" else nn.GELU()
+
+ norm_cls = nn.BatchNorm1d if norm == "batch" else nn.LayerNorm
+ self.bn1 = norm_cls(out_features, device=device, dtype=dtype)
+ self.bn2 = norm_cls(out_features, device=device, dtype=dtype)
+ self.downsample = (
+ nn.Linear(in_features, out_features, bias=False, device=device, dtype=dtype)
+ if in_features != out_features
+ else None
+ )
+ self.residual = residual
+
+ def forward(self, x):
+ identity = x
+
+ out = self.linear1(x)
+ out = self.bn1(out)
+ out = self.act_fn(out)
+
+ out = self.linear2(out)
+ out = self.bn2(out)
+
+ if self.downsample is not None:
+ identity = self.downsample(identity)
+
+ if self.residual:
+ out += identity
+
+ out = self.act_fn(out)
+ return out