step-decomposed-influence 0.1.1

Step-Decomposed Influence (SDI) and fast TracIn for (looped) transformers.

Step-Decomposed Influence (SDI)

Step-Decomposed Influence (SDI) decomposes TracIn into a step-resolved influence trajectory for looped/weight-tied models. This reference implementation computes per-step influence while avoiding per-example gradient materialization by sketching during backprop (TensorSketch for 2D weights, CountSketch for 1D params).

Paper: https://arxiv.org/abs/2602.10097

Install (UV-first)

# From PyPI
uv add step-decomposed-influence

# From source (editable, in a UV project)
uv add --editable .

CPU vs CUDA:

• This package is device-agnostic. CUDA works automatically if you install a CUDA-enabled PyTorch.
• Follow the official PyTorch install instructions to choose the right wheel for your system.

Quickstart

import torch
import torch.nn.functional as F
from sdi import ProjectedTracInSDI, CheckpointSpec

# model = your looped transformer
# target_modules = the recurrent core used at each loop step

def loss_fn(model, batch):
    logits = model(batch["tokens"])  # (B, T, vocab)
    targets = batch["tokens"][:, 1:]
    logits = logits[:, :-1, :]
    per_pos = F.cross_entropy(
        logits.reshape(-1, logits.size(-1)),
        targets.reshape(-1),
        reduction="none",
    ).view(targets.size(0), -1)
    return per_pos.mean(dim=1)  # per-example loss (B,)

est = ProjectedTracInSDI(
    model=model,
    target_modules=target_modules,
    projection_size=2048,
    loss_reduction="sum",
)

out = est.influence_across_checkpoints(
    checkpoints=[CheckpointSpec("checkpoints/ckpt_0001.pt")],
    train_loader=train_loader,
    query_loader=query_loader,
    loss_fn=loss_fn,
    mode="sdi",
    train_chunk_size=128,
    query_chunk_size=128,
)

# out.sdi:    (N_train, N_query, steps_query)
# out.tracin: (N_train, N_query)

Example

Run the toy looped-transformer demo:

uv run python examples/toy_looped_transformer_sdi.py --device cuda

This script:

• trains a tiny looped transformer on random data,
• saves a few checkpoints,
• computes projected SDI (sketch-during-backprop),
• and shows fast TracIn (scalar-only) via mode="tracin".

API (minimal)

• ProjectedTracInSDI: projected SDI/TracIn using TensorSketch + CountSketch.
• FullGradientTracInSDI: exact baseline (materializes per-sample gradients; small models only).
• CheckpointSpec: path + optional checkpoint weight (eta).
• InfluenceOutputs: sdi (None when mode="tracin"), tracin, and optional sdi_matrix.

Assumptions & limitations

• Steps are inferred from how often the selected parameters are used in a single forward/backward. For looped transformers, this equals the loop horizon when you pass only the looped core.
• Projected estimator supports:
  • 2D params: nn.Linear.weight, nn.Embedding.weight (TensorSketch)
  • 1D params: CountSketch of per-sample grads via Opacus grad samplers (plus analytic nn.Linear.bias)
• Custom modules without Opacus grad-sampler coverage are not supported in strict mode.
• Checkpoint loading uses torch.load. Do not load untrusted checkpoints.

Citation

For citing this software repository, use GitHub's "Cite this repository" button.

For citing the paper, use:

@article{kaissis2026step,
  title         = {Step-resolved data attribution for looped transformers},
  author        = {Georgios Kaissis and David Mildenberger and Juan Felipe Gomez and Martin J. Menten and Eleni Triantafillou},
  year          = {2026},
  journal       = {arXiv preprint arXiv:2602.10097},
  eprint        = {2602.10097},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2602.10097},
}

Authors

• Georgios Kaissis (Hasso-Plattner Institute)
• David Mildenberger (Technical University of Munich)
• Juan Felipe Gomez (Harvard University)
• Martin J. Menten (Technical University of Munich)
• Eleni Triantafillou (Google DeepMind)