https://github.com/zeekmartin/sal-torch
Structurally Adaptive Learning - training-time sparsification for PyTorch
https://github.com/zeekmartin/sal-torch
machine-learning model-compression pruning pytorch sparsification transformers
Last synced: 25 days ago
JSON representation
Structurally Adaptive Learning - training-time sparsification for PyTorch
- Host: GitHub
- URL: https://github.com/zeekmartin/sal-torch
- Owner: zeekmartin
- License: other
- Created: 2026-06-25T15:53:42.000Z (about 1 month ago)
- Default Branch: main
- Last Pushed: 2026-06-29T20:32:07.000Z (27 days ago)
- Last Synced: 2026-07-01T14:34:31.591Z (26 days ago)
- Topics: machine-learning, model-compression, pruning, pytorch, sparsification, transformers
- Language: Python
- Homepage: https://pypi.org/project/sal-torch/
- Size: 226 KB
- Stars: 1
- Watchers: 1
- Forks: 0
- Open Issues: 1
-
Metadata Files:
- Readme: README.md
- Changelog: CHANGELOG.md
Awesome Lists containing this project
README
# sal-torch
    
**Structurally Adaptive Learning for PyTorch**
Training-time sparsification that makes neural networks structurally resilient to compression.
## Install
```bash
pip install sal-torch # core
pip install sal-torch[hf] # + HuggingFace Trainer
pip install sal-torch[all] # everything
```
```python
from sal import SALConfig, SALCallback
config = SALConfig.auto(model)
trainer = Trainer(model=model, callbacks=[SALCallback(config)])
trainer.train()
```
Three lines. Any transformer. Compression-resilient.
## Know your model before you touch it
### PlasticityScanner — where can a model absorb compression?
FI tells you how fragile a model *is*. `PlasticityScanner` tells you how much
room it has to *reorganize*, so you know where it is safe to compress. It scores
three complementary axes per layer — routing flexibility (attention entropy),
inter-layer redundancy (linear CKA), and intra-layer redundancy (an MI proxy) —
and folds them into an **absorption map** that labels each layer `ELASTIC`
(safe), `SATURATED` (bottleneck), or `HUB` (compensates when others are pruned).
```python
from sal import PlasticityScanner
pmap = PlasticityScanner(model, probe_dataset).scan()
print(pmap.summary) # "3 elastic, 1 saturated, 2 hub | mean routing=0.61 ..."
rec = pmap.recommend(target_compression=0.33)
rec.safe_to_prune # [(layer, head), ...] — prune these first
rec.never_touch # heads in hub layers — leave alone
rec.expected_impact # heuristic accuracy delta
pmap.save("plasticity.json") # raw scores
pmap.save("plasticity.pdf") # visual report (needs sal-torch[reports])
```
### sal.compare() — SAL vs. other pruning methods
Benchmark SAL against post-hoc baselines at a matched compression level and see
which keeps the most accuracy (or lowest loss) after heads are removed.
```python
from sal import compare
result = compare(model, train_dataset, eval_dataset,
methods=["sal", "magnitude", "random_posthoc"],
compression=0.33, sal_epochs=3, metric="accuracy")
print(result.table) # method | score | pruned_heads | time
print(result.winner)
result.save("comparison.pdf") # bar chart + table
# plug in your own method
compare.register_method("my_pruner", lambda model, ds, eval_ds, ctx: my_score)
```
## Continual learning without replay buffers
### StructuralGuard — protect what matters when you fine-tune
When you fine-tune a trained model on a new task, it quietly overwrites the
structure that carried the old one. `StructuralGuard` reads the model's
structural map and **freezes the critical attention heads** (hub layers,
structural bottlenecks, and the functionally unique heads) while leaving the
redundant heads free to absorb the new task. No EWC, no replay buffer, no
distillation — the topology itself decides what to protect.
```python
from sal import StructuralGuard
# After training on task A, build a guard from the model's structure.
guard = StructuralGuard.from_model(model, probe_dataset, protection_level=0.5)
print(guard.protected_heads) # [(layer, head), ...] frozen during fine-tuning
print(guard.trainable_heads) # [(layer, head), ...] free to absorb task B
print(guard.protection_map) # {layer: [protected head indices]}
guard.protect(model) # zero gradients for protected heads (backward hooks)
trainer.train() # fine-tune on task B with ANY training loop
guard.release(model)
drift = guard.measure_drift(model, probe_dataset=probe_dataset)
print(drift.forgetting_score) # 0 = nothing forgot, 1 = total reorganization
print(drift.protected_integrity) # ~1.0 if the protected heads held
guard.save("model_guard.json") # serialize; reload before task C, D, ...
guard = StructuralGuard.load("model_guard.json")
```
Protection is at the **head level** — some heads in a layer can be frozen while
others in the same layer keep learning. `protection_level` (0.0–1.0) sets the
fraction of the most critical heads to protect.
HuggingFace `Trainer`? Use the callback — it applies protection on
`train_begin`, measures drift on `train_end`:
```python
from sal import StructuralGuardCallback
guard = StructuralGuard.from_model(model, probe_dataset)
callback = StructuralGuardCallback(guard)
trainer = Trainer(model=model, callbacks=[callback])
trainer.train()
print(callback.drift_report.summary)
```
### DriftMonitor — measure structural forgetting after any fine-tuning
`DriftMonitor` quantifies how much a model's structure moved, guarded or not.
Snapshot before and after, then compare.
```python
from sal import DriftMonitor
monitor = DriftMonitor(model, probe_dataset)
monitor.snapshot("before_task_b")
trainer.train()
monitor.snapshot("after_task_b")
drift = monitor.compare("before_task_b", "after_task_b")
print(drift.summary)
print(drift.layer_drift) # per-layer activation retention (1 = identical)
print(drift.classification_changes) # layers whose fragility class flipped
drift.save("drift_report.json")
drift.save("drift_report.pdf") # visual before/after comparison
```
Snapshots are keyed, so you can track drift across many sequential tasks and
compare any pair.
## Examples
- [`examples/quickstart.py`](examples/quickstart.py) — 3-line SAL training on DistilBERT
- [`examples/standalone_fi.py`](examples/standalone_fi.py) — Fragility Index scan, no training
- [`examples/full_control.py`](examples/full_control.py) — manual config + standalone trainer
- [`examples/compare_with_without_sal.py`](examples/compare_with_without_sal.py) — SAL vs. baseline under compression
New here? Start with [docs/getting_started.md](docs/getting_started.md).
## License
BSL 1.1 — free for research and evaluation. Commercial production requires a license.
Built by [Cognitive Engineering](https://cognitive-engineering.dev) in Switzerland.