sal-torch 0.3.0

Structurally Adaptive Learning — training-time sparsification for robust neural networks

sal-torch

Structurally Adaptive Learning for PyTorch

Training-time sparsification that makes neural networks structurally resilient to compression.

Install

pip install sal-torch            # core
pip install sal-torch[hf]        # + HuggingFace Trainer
pip install sal-torch[all]       # everything

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).

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.

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.

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:

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.

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 — 3-line SAL training on DistilBERT
• examples/standalone_fi.py — Fragility Index scan, no training
• examples/full_control.py — manual config + standalone trainer
• examples/compare_with_without_sal.py — SAL vs. baseline under compression

New here? Start with docs/getting_started.md.

License

BSL 1.1 — free for research and evaluation. Commercial production requires a license.

Built by Cognitive Engineering in Switzerland.