noise-injection 0.3.0

Memory-efficient weight noise injection for LLMs

Noise Injection

Memory-efficient noise injection for LLM weights. Useful for sandbagging detection, robustness testing, and interpretability research.

Documentation

Installation

pip install noise-injection

# With MLX support (Apple Silicon)
pip install "noise-injection[mlx]"

Quick Start

from noise_injection import inject_noise, with_noise, generators as G, selectors as S

# Context manager (auto-cleanup)
with with_noise(model, G.gaussian(sigma=0.01)):
    accuracy = evaluate(model)
# Noise automatically removed

# Manual control
layer = inject_noise(model, G.gaussian(sigma=0.01))
accuracy = evaluate(model)
layer.remove(model)

# Target specific parameters
with with_noise(model, G.gaussian(sigma=0.01), selector=S.by_name("layers.0.*")):
    ...

API

Noise Injection

# Inject noise and get handle for removal
layer = inject_noise(
    model,
    generator=G.gaussian(sigma=0.01, mu=0.0),
    selector=S.all_params(),      # Which parameters to affect
    seed=42,               # For reproducibility
    store_tensors=False    # False = regenerate from seed (zero memory overhead)
)

# Remove when done
layer.remove(model)

Context Manager

with with_noise(model, G.gaussian(sigma=0.01)) as layer:
    # Model has noise injected
    accuracy = evaluate(model)
# Noise automatically removed

Sigma Sweep

from noise_injection import sweep_sigma

results = sweep_sigma(
    model,
    eval_fn=lambda: evaluate(model),
    sigma_start=0.0,
    sigma_end=0.01,
    sigma_step=0.001,
    mu=0.0,
)
# Returns: [(0.0, 0.65), (0.001, 0.63), ...]

Generators

from noise_injection import generators as G

G.gaussian(sigma=0.01, mu=0.0)           # Gaussian noise
G.uniform(low=-0.01, high=0.01)          # Uniform noise
G.custom(lambda shape, gen: ...)         # Custom generator

Selectors

from noise_injection import selectors as S

S.all_params()                                  # All parameters
S.by_name("layers.0.*")                  # Regex pattern
S.by_name("attn") & S.by_name("weight")  # Composable

Examples

See examples/ for complete usage examples:

• sandbagging-detection-mmlu - Replicate sandbagging detection from arXiv:2412.01784

# MMLU sweep for sandbagging detection
python examples/sandbagging-detection-mmlu/mmlu_sweep.py --num-samples 100

How It Works

Snapshot-restore (default for sweeps): Clone selected parameters once, restore from the clone between iterations. Bitwise-exact weight restoration regardless of dtype, including bfloat16/float16.

1. Snapshot selected parameters
2. For each sigma:
   a. Restore from snapshot
   b. Inject noise → evaluate
3. Final restore from snapshot

Seeded subtract (for single injection, or memory-constrained sweeps): Uses seeded RNG to regenerate identical noise for add/subtract. O(max single parameter) memory overhead instead of O(total parameters).

1. Generate noise with seeded RNG
2. Add noise in-place → evaluate
3. Regenerate same noise (same seed) → subtract to restore

Module Structure

noise-injection/
├── src/noise_injection/
│   ├── __init__.py         # Main exports
│   ├── api/
│   │   ├── functional.py   # inject_noise, remove_noise
│   │   ├── context.py      # with_noise context manager
│   │   └── sweep.py        # sweep_sigma utility
│   └── core/
│       ├── generators.py   # Noise generators (gaussian, uniform, custom)
│       ├── selectors.py    # Parameter selectors
│       └── noise_layer.py  # NoiseLayer, NoiseStack
├── examples/
│   └── sandbagging-detection-mmlu/  # Replicate arXiv:2412.01784
└── pyproject.toml

Citation

If you use this tool in your research, please cite the original paper on noise injection for sandbagging detection:

@article{tice2024noise,
  title={Noise Injection Reveals Hidden Capabilities of Sandbagging Language Models},
  author={Tice, Cameron and Kreer, Philipp Alexander and others},
  journal={arXiv preprint arXiv:2412.01784},
  year={2024}
}