noise-injection 0.1.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.

Installation

# From GitHub
pip install git+https://github.com/theoryzhenkov/noise-injection.git

# With optional dependencies
pip install "noise-injection[mlx] @ git+https://github.com/theoryzhenkov/noise-injection.git"

# Local development
git clone https://github.com/theoryzhenkov/noise-injection.git
cd noise-injection
pip install -e .

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(),      # 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()                                  # 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

Zero-overhead memory efficiency: Uses seeded RNG to generate identical noise for add/subtract operations. No parameter copying needed - the same noise can be regenerated and subtracted to restore original weights.

For each sigma:
  1. Generate noise with seeded RNG
  2. Add noise to parameters in-place
  3. Evaluate model
  4. Regenerate same noise (same seed)
  5. Subtract to restore original weights

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}
}