sal-torch : training-time head masking for compression-resilient transformers

projects
1

Hi everyone,

I’ve been working on a different approach to model compression and wanted to share the resulting package.

The idea: instead of pruning a trained model and hoping it holds up, what if we trained the model to tolerate pruning from the start? sal-torch does this by randomly deactivating attention heads during training. The model is forced to redistribute its learned functions across surviving heads. When you later prune 33% of heads permanently, it barely notices.

How it works in practice:

from sal import SALConfig, SALCallback

config = SALConfig.auto(model)  # auto-detects architecture
trainer = Trainer(model=model, callbacks=[SALCallback(config)])
trainer.train()

Key implementation details:

  • Forward pre-hooks on attention output projections, zeroing per-head slices
  • Heads accumulate (stay pruned, new ones added progressively)
  • Pruned set held through end of training so adaptation bakes into weights
  • Random selection — we validated that topology-guided selection doesn’t outperform random at tested scales

Results on DistilBERT/SST-2 @ 33% compression:

  • SAL-trained: 0.81 accuracy
  • Magnitude pruning: 0.53
  • Random post-hoc: 0.43

The package also includes:

  • FI (Fragility Index) : structural diagnostic that classifies each layer as IMMUNE/BUFFER/CRITICAL based on triangle-support in the head correlation graph
  • PlasticityScanner : 3-axis structural analysis (routing entropy, CKA similarity, mutual information) showing where a model has capacity to absorb compression
  • sal.compare() : benchmarks SAL vs magnitude vs random on your model

12 architectures supported via auto-detection: LLaMA, Mistral, GPT-2, BERT, RoBERTa, DistilBERT, ViT, Phi, Phi-3, Gemma, Gemma-2, Qwen2. Adding a new one is just a pattern entry in the registry.

Works with HuggingFace Trainer (callback) or standalone PyTorch (SALTrainer). Zero extra dependencies beyond torch, numpy, scipy.

pip install sal-torch
GitHub: GitHub - zeekmartin/sal-torch: Structurally Adaptive Learning - training-time sparsification for PyTorch · GitHub

BSL 1.1 free for research and evaluation.

Happy to answer questions about the mechanism or implementation choices.

1 Like