vitalroute 0.2.0

Task-aware training controller via layer vitality monitoring

VitalRoute

Task-aware training controller for feed-forward classifiers. It sits on top of the training optimizer (Adam, SGD, etc.) and decides when to apply training tactics based only on training-set shape (class counts and size) — not by hand-tuning flags for every dataset.

The library distills a research line on network vitality — stasis, weak coupling, saturation, and transferable structure — into probes, label-free parent selection, and class-aware sampling, without requiring any legacy codebase or naming scheme.

Vitality signals and tactics

A classic biological metaphor treats the network like a body that can be examined while it learns.

Signal	Meaning
Stasis	Hidden unit barely responds (dead ReLU, etc.)
Weak weights	Weight column has collapsed
Weak input	Incoming activations are tiny vs weights
Saturation	Unit stuck near a constant output

From those signals, VitalRoute provides:

1. Vitality sampler — For imbalanced data, oversample classes with high composite stress (all four signals, not only stasis).
2. Transfer pick — For scarce data, selects the best pretrained parent by lowest stasis on the new inputs (no labels needed), then warm-starts weights.
3. Hard-sample sampler — When class rebalancing is off, oversample individual examples with high per-sample stress (stasis + weak coupling + low confidence).
4. LR scale — Slow learning on layers with high stasis: lr_l = base_lr / (1 + α · stasis_l) (dampens high-stasis layers at elevated base LR).
5. Monitor — Logs layer health; resets stuck units when stasis is high (skipped for CNN-style models with a separate head).

An adaptive router turns (1)–(4) on or off from class counts and dataset size.

Install

pip install vitalroute==0.2.0
# PyTorch extras (probe, CNN, MLPerf hooks):
pip install "vitalroute[torch]==0.2.0"

Requires Python 3.10+ and NumPy. See PyPI.

Quick use

import numpy as np
from vitalroute import adaptive_controller, profile_task, route_plan
from vitalroute.backbone import MLP, LayerSpec, Adam

# Training arrays
X_train, y_train = ...
num_classes = 10

# Preview routing (no training)
prof = profile_task(y_train, num_classes)
plan = route_plan(prof, parent_pool_available=False)
print(plan.label)  # e.g. "imbalance", "transfer", "transfer+imbalance", "monitor"

# Training loop integration
ctrl = adaptive_controller(y_train, num_classes, parent_pool=None, verbose=True)
opt = ctrl.make_optimizer("adam", lr=1e-3)  # vitality-scaled when route includes lr_scale

sampler, _ = ctrl.bootstrap(model, X_train, y_train, num_classes=num_classes)

for epoch in range(epochs):
    ctrl.on_epoch_start(model, X_train, opt, epoch)  # refresh LR scales if enabled
    if sampler is not None:
        idx = sampler.sample_indices(epoch, model, X_train, y_train, len(y_train))
        X_ep, y_ep = X_train[idx], y_train[idx]
    else:
        X_ep, y_ep = X_train, y_train
    # ... standard mini-batches, loss, optimizer step ...
    ctrl.after_epoch(model, X_train, rng=np.random.default_rng(epoch))

Runnable example: examples/digits_imbalanced_demo.py.

Router rules (defaults)

Condition	Enabled
min_class / max_class < 0.25 and minority ≥ 15 samples	Vitality sampler (composite stress)
n ≤ 200 or min_class ≤ 12, and parent pool provided	Transfer pick
Scarce balanced data	Transfer + hard-sample sampler (no class sampler)
n ≥ 80, sampler off	LR scale
n ≥ 40, sampler off	Hard-sample sampler
Always (when training)	Monitor (+ conditional reset)

How it compares to inverse-frequency weighting

On a clean long-tail benchmark, VitalRoute ≈ inverse-frequency (inv_freq). They converge to the same answer because rare classes and broken-neuron classes heavily overlap — the network sees minority classes less, so their neurons die more.

Scenarios where VitalRoute differs from inv_freq:

Scenario	Distinction
Imbalanced but not uniformly scarce	A class with enough samples but high confusability (broken neurons) gets oversampled; inv_freq ignores it
Difficulty shifts mid-training	VitalRoute refreshes stress every N epochs; inv_freq is static
Label-free transfer selection	Selects the best pretrained parent by stasis on new inputs — no labels needed. inv_freq has no equivalent
Hard-sample curriculum	Per-sample stress (stasis + low confidence) for scarce balanced data; inv_freq only works at class level

For purely long-tail problems with clean class boundaries, inv_freq is simpler and nearly as good. When classes overlap, difficulty shifts, or label-free transfer selection is required, VitalRoute adds value.

Package layout

vitalroute/
  README.md
  PAPER.md            # research paper style writeup
  INTEGRATION.md      # NumPy, PyTorch, CNN, MLPerf integration
  pyproject.toml
  vitalroute/
    vitality.py         # layer stress probes + per-class/per-sample stress
    imbalance.py        # composite vitality class sampler (NumPy)
    hard_samples.py     # per-sample stress sampler (NumPy)
    lr_scale.py         # vitality-scaled Adam / SGD (NumPy)
    transfer.py         # label-free parent pick (NumPy)
    router.py           # task profile + adaptive controller (NumPy)
    torch_probe.py      # VitalityProbe — forward hooks for any nn.Module
    torch_probe_cnn.py  # CNNVitalityProbe — BatchNorm-aware conv probing
    torch_probes.py     # make_probe / detect_architecture registry
    torch_transfer.py   # PyTorch transfer pick + warm_start_from_parent
    torch_lr_scale.py   # per-layer param groups + apply_lr_scales
    torch_data.py       # stratified_probe_batch helper
    torch_samplers.py   # TorchVitalitySampler + TorchHardSampleSampler
    torch_controller.py # TorchTrainingController + torch_adaptive_controller
    mlperf_hooks.py     # VitalRouteMLPerfCallback for MLPerf-style loops
    backbone/           # optional reference MLP for demos
  examples/
    digits_imbalanced_demo.py     # NumPy backbone quick demo
    benchmark_baselines.py        # NumPy baseline comparison (digits)
    torch_probe_demo.py           # VitalityProbe on a PyTorch MLP
    torch_benchmark_fmnist.py     # PyTorch baseline comparison (Fashion-MNIST)
    cifar10_resnet_benchmark.py   # ResNet18 / CIFAR-10-LT
    cifar10_lt_benchmark.py       # CIFAR-10-LT + transfer-pick demo
    imagenet_lt_benchmark.py      # CIFAR-100-LT local / ImageNet-LT
    mlperf_resnet_integration.py  # MLPerf callback demo
    run_cnn_benchmarks.py         # run all CNN benchmarks (--quick / --full)
  tests/

Evidence summary

Measured on public-style benchmarks:

Setting	Typical gain
Imbalanced digits / Fashion minority classes	+2–4% minority accuracy vs uniform
vs inverse-frequency baseline (same imbalanced digits)	+0.7% minority, lower variance
Scarce digit subset with parent pool	up to +10% vs cold start
Scarce cat/dog (MLP / small CNN)	+2–3% with transfer pick

NumPy backbone benchmark (examples/benchmark_baselines.py), 3 seeds, 30 epochs, 5:1 imbalance on digits:

Method         Overall    Minority
uniform        93.7%±1.1%  87.9%±2.3%
inv_freq       94.4%±0.8%  90.1%±1.0%
vitalroute     95.1%±0.3%  90.8%±1.0%
stasis_only    95.0%±0.7%  90.7%±1.5%

Highest overall accuracy and lowest seed variance in this run.

PyTorch benchmark (examples/torch_benchmark_fmnist.py), 3 seeds, 20 epochs, 10:1 imbalance on Fashion-MNIST MLP:

Method         Overall    Minority
uniform        80.1%±0.4%  72.8%±1.2%
inv_freq       81.7%±0.5%  77.6%±0.9%
focal          80.0%±0.3%  72.6%±0.2%
vitalroute     81.7%±0.2%  76.5%±0.6%

Overall and minority accuracy match inv_freq; variance is lower than uniform and focal.

CNN benchmark (examples/cifar10_resnet_benchmark.py), ResNet18, 10:1 long-tail CIFAR-10 (classes 0–4: 1000 samples each; classes 5–9: 100 each):

Method	Role in comparison
uniform	Unweighted sampling baseline
inv_freq	Inverse-frequency WeightedRandomSampler
focal	Focal loss (γ=2), uniform sampling
vitalroute	Adaptive vitality class sampler + CNNVitalityProbe

Smoke run (2 epochs, CPU, 1 seed). Full numbers: python examples/run_cnn_benchmarks.py --full.

Method         Overall    Minority
uniform        26.6%       0.0%
inv_freq       38.7%      29.4%
focal          27.4%       0.0%
vitalroute     36.9%      33.5%

At epoch 2, minority accuracy is 33.5% (vitalroute) vs 29.4% (inv_freq). Uniform and focal report 0% on minority classes in this run.

CNN long-tail: mechanisms

Tactic	Behavior
Vitality class sampler	Oversamples classes with high composite stress on conv and head layers, not frequency alone
vs focal loss	Focal modulates loss per sample; VitalRoute modulates sampling from internal activation health
Adaptive router	Enables sampler, transfer pick, or LR scale from class counts and dataset size
CNNVitalityProbe	Reads Conv→BN→ReLU trunk and/or linear head (probe_zone=head, trunk, or all)
PyTorch transfer pick	Ranks parents by stasis on unlabeled inputs; warm-starts trunk via warm_start_from_parent
Per-layer LR	make_optimizer() assigns named param groups; stasis rate scales each group's learning rate
MLPerf callback	VitalRouteMLPerfCallback exposes the same epoch hooks as MLPerf reference training loops

Running CNN benchmarks

pip install -e ".[dev]"
python examples/run_cnn_benchmarks.py --quick   # 2 epochs per script
python examples/run_cnn_benchmarks.py --full    # 15 epochs, multiple trials

PyTorch integration

Probe only (read vitality signals)

VitalityProbe attaches to any torch.nn.Module via forward hooks — no changes to the model or optimizer required:

from vitalroute.torch_probe import VitalityProbe

probe = VitalityProbe(model)        # attach once; pairs Linear→ReLU automatically

for epoch in range(epochs):
    train_one_epoch(model, ...)
    probe.observe(X_train)          # one forward pass, no gradients
    print(probe.summary())          # per-layer stasis + composite stress
    print(f"mean stasis: {probe.mean_stasis():.3f}")

# Per-class and per-sample stress for custom sampling
class_scores  = probe.per_class_stress(X_train, y_train, num_classes=10)
sample_scores = probe.per_sample_stress(X_train, y_train)

probe.detach()                      # remove hooks

Full adaptive controller

torch_adaptive_controller selects tactics from the training label distribution. For CNNs, set architecture="cnn" and probe_zone to "head" (classifier only), "trunk" (conv blocks), or "all" (both).

from vitalroute.torch_controller import torch_adaptive_controller
from vitalroute.torch_data import stratified_probe_batch
from torch.utils.data import DataLoader

ctrl = torch_adaptive_controller(
    y_train, num_classes=10,
    architecture="cnn",
    probe_zone="all",
    parent_pool=None,   # optional: [("parent_a", model_a), ...]
    verbose=True,
)

X_probe, y_probe, _ = stratified_probe_batch(
    train_dataset, y_train, per_class=50, num_classes=10, device="cuda"
)
sampler = ctrl.setup(model, X_probe, y_probe, y_full=y_train, num_classes=10)
optimizer = ctrl.make_optimizer(model, torch.optim.SGD, lr=0.05, momentum=0.9)

loader = DataLoader(dataset, sampler=sampler, batch_size=64) if sampler else DataLoader(dataset, batch_size=64, shuffle=True)

for epoch in range(epochs):
    ctrl.on_epoch_start(model, X_probe, optimizer, epoch)
    for X_batch, y_batch in loader:
        ...  # standard loss + backward + step
    ctrl.after_epoch(model, X_probe, y_probe)

ctrl.detach()

MLPerf Training integration

VitalRouteMLPerfCallback wraps TorchTrainingController in an epoch callback interface aligned with MLPerf reference training loops. See INTEGRATION.md and examples/mlperf_resnet_integration.py.

Runnable examples: examples/torch_probe_demo.py, examples/torch_benchmark_fmnist.py, examples/cifar10_resnet_benchmark.py, examples/run_cnn_benchmarks.py.

License

MIT

Contributing

Contributions are welcome. See CONTRIBUTING.md for setup, PR requirements, and code guidelines. Please read CODE_OF_CONDUCT.md before participating. Report security issues via SECURITY.md, not public issues.

Related Work

The following related work is grouped by topic. Distinctions from VitalRoute are noted per entry.

Adaptive class resampling

• ART: Adaptive Resampling-based Training for Imbalanced Classification (2025) — periodically refreshes class sampling weights using class-wise F1 scores. VitalRoute uses internal neuron health signals instead of output metrics.

Dead neuron analysis and pruning

• When to Prune? A Policy towards Early Structural Pruning — uses dead-neuron rates to guide structured pruning during training. VitalRoute uses the same signal to drive sampling, not pruning.
• Dead neurons in Deep Learning (overview)

Dynamic network structure for imbalanced learning

• Adaptive Neuron Growth/Pruning for Imbalanced Classification (2025) — adds/removes neurons per class using gradient magnitude. Orthogonal to VitalRoute: modifies architecture rather than sampling.

Per-layer learning rate scaling

• LENA: Layer-wise Adaptive LR Scaling — scales per-layer LR by gradient variance. VitalRoute scales by stasis (dead unit fraction), a complementary signal.
• LLR: Heavy-Tail Guided Layerwise LR for LLMs (2025) — uses weight spectrum heavy-tailedness. Same goal, different diagnostic.
• AdaLip: Adaptive LR per Layer via Lipschitz Estimation — Lipschitz-constant-based per-layer LR.
• LARS / LAMB — weight/gradient ratio scaling; used in large-batch distributed training.

Label-free transfer model selection

• TURTLE: Unsupervised Transfer Learning (2024) — selects pretrained models without labels via representation-level generalization objectives. VitalRoute uses stasis rate on new data — simpler, different rationale.
• DISCO: Spectral Component Distribution for Transfer Assessment (2024) — SVD of feature distributions for transferability scoring.

Focal Loss (baseline used in benchmarks)

• Focal Loss for Dense Object Detection — Lin et al., 2017. Standard hard-example weighting via loss modulation.

Curriculum / hard-sample learning

• Self-Paced Learning — Bengio et al., 2009. Foundation for curriculum-style training.
• Online Hard Example Mining — Shrivastava et al., 2016. Per-sample difficulty weighting from loss values.