Metadata-Version: 2.4
Name: noise-injection
Version: 0.3.0
Summary: Memory-efficient weight noise injection for LLMs
Keywords: ai-safety,sandbagging,model-evaluation,noise-injection,llm
Requires-Python: >=3.9
Description-Content-Type: text/markdown
Requires-Dist: torch>=2.0.0
Requires-Dist: tqdm>=4.65.0
Provides-Extra: examples
Requires-Dist: transformers>=4.30.0; extra == "examples"
Requires-Dist: datasets>=2.0.0; extra == "examples"
Requires-Dist: accelerate>=0.20.0; extra == "examples"
Requires-Dist: numpy>=1.20.0; extra == "examples"
Requires-Dist: python-dotenv>=1.0.0; extra == "examples"
Requires-Dist: pyyaml>=6.0; extra == "examples"
Requires-Dist: matplotlib>=3.9.0; extra == "examples"
Provides-Extra: mlx
Requires-Dist: mlx>=0.20.0; extra == "mlx"
Requires-Dist: mlx-lm>=0.19.0; extra == "mlx"
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: build>=1.4.0; extra == "dev"
Requires-Dist: twine>=6.2.0; extra == "dev"
Provides-Extra: docs
Requires-Dist: mkdocs-material>=9.0.0; extra == "docs"

# Noise Injection

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

**[Documentation](https://theoryzhenkov.github.io/noise-injection/)**

## Installation

```bash
pip install noise-injection

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

## Quick Start

```python
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

```python
# 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

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

### Sigma Sweep

```python
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

```python
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

```python
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](examples/sandbagging-detection-mmlu/)** - Replicate sandbagging detection from [arXiv:2412.01784](https://arxiv.org/abs/2412.01784)

```bash
# 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:

```bibtex
@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}
}
```
