Pruning Schedules

Pruning schedules control when and how sparsity is introduced during training. The right schedule can significantly impact final model quality.

Overview

Entrenar supports three pruning schedules:

OneShot Schedule

Applies target sparsity in a single step. Simple and effective for post-training compression.

#![allow(unused)]
fn main() {
use entrenar::prune::PruningSchedule;

let schedule = PruningSchedule::OneShot { step: 1000 };

// Sparsity transitions instantly at step 1000
assert_eq!(schedule.sparsity_at_step(999), 0.0);
assert_eq!(schedule.sparsity_at_step(1000), 1.0);  // Returns multiplier
assert_eq!(schedule.sparsity_at_step(2000), 1.0);
}

Use Cases

• LLM pruning (SparseGPT, Wanda)
• Post-training compression
• When fine-tuning budget is limited

Pros and Cons

Pros:

• Simplest to implement
• Works well with calibration-based methods
• No hyperparameter tuning for schedule

Cons:

• Can cause accuracy drop without fine-tuning
• Not ideal for very high sparsity (>70%)

Gradual Schedule

Linearly interpolates from initial to final sparsity over a range of steps.

#![allow(unused)]
fn main() {
let schedule = PruningSchedule::Gradual {
    start_step: 1000,
    end_step: 5000,
    initial_sparsity: 0.0,
    final_sparsity: 0.5,
    frequency: 500,
};

// Before start: no pruning
assert_eq!(schedule.sparsity_at_step(500), 0.0);

// During pruning: linear interpolation
assert_eq!(schedule.sparsity_at_step(1000), 0.0);
assert_eq!(schedule.sparsity_at_step(3000), 0.25);
assert_eq!(schedule.sparsity_at_step(5000), 0.5);

// After end: stay at final sparsity
assert_eq!(schedule.sparsity_at_step(6000), 0.5);
}

Parameters

Frequency Effect

The frequency parameter controls how often the mask is updated:

#![allow(unused)]
fn main() {
// Update mask every 500 steps
let frequent = PruningSchedule::Gradual {
    start_step: 0,
    end_step: 10000,
    initial_sparsity: 0.0,
    final_sparsity: 0.5,
    frequency: 500,  // 20 updates total
};

// Update mask every 2000 steps
let sparse_updates = PruningSchedule::Gradual {
    start_step: 0,
    end_step: 10000,
    initial_sparsity: 0.0,
    final_sparsity: 0.5,
    frequency: 2000,  // 5 updates total
};
}

More frequent updates allow finer control but add overhead.

Cubic Schedule (Zhu & Gupta 2017)

Uses a cubic polynomial that prunes aggressively early and slows toward the end.

#![allow(unused)]
fn main() {
let schedule = PruningSchedule::Cubic {
    start_step: 0,
    end_step: 10000,
    final_sparsity: 0.7,
};

// Formula: s_t = s_f * (1 - (1 - t/T)^3)
}

Mathematical Formula

The cubic schedule follows:

s_t = s_f * (1 - (1 - t/T)^3)

Where:

• s_t = sparsity at step t
• s_f = final target sparsity (e.g., 0.7)
• t = current step within pruning window
• T = total pruning steps (end_step - start_step)

Sparsity Progression

For final_sparsity = 0.7 over 10000 steps:

#![allow(unused)]
fn main() {
let schedule = PruningSchedule::Cubic {
    start_step: 0,
    end_step: 10000,
    final_sparsity: 0.7,
};

// Verify progression
assert!((schedule.sparsity_at_step(0) - 0.0).abs() < 0.01);
assert!((schedule.sparsity_at_step(2500) - 0.488).abs() < 0.01);
assert!((schedule.sparsity_at_step(5000) - 0.613).abs() < 0.01);
assert!((schedule.sparsity_at_step(7500) - 0.689).abs() < 0.01);
assert!((schedule.sparsity_at_step(10000) - 0.7).abs() < 0.01);
}

Why Cubic?

The cubic curve has desirable properties:

1. Aggressive early pruning - Model is most plastic early in training
2. Gradual convergence - Allows fine-tuning of remaining weights
3. Smooth transitions - No sudden sparsity jumps

Reference

Zhu, M., & Gupta, S. (2017). "To Prune, or Not to Prune: Exploring the Efficacy of Pruning for Model Compression." arXiv:1710.01878

Choosing a Schedule

Decision Tree

Is this post-training compression?
├── Yes → OneShot
└── No → Target sparsity > 50%?
    ├── Yes → Cubic
    └── No → Gradual

Recommendations by Scenario

Configuration Validation

All schedules validate their parameters:

#![allow(unused)]
fn main() {
use entrenar::prune::PruningConfig;

// Invalid: end_step before start_step
let bad_config = PruningConfig::new()
    .with_schedule(PruningSchedule::Gradual {
        start_step: 5000,
        end_step: 1000,  // Invalid!
        initial_sparsity: 0.0,
        final_sparsity: 0.5,
        frequency: 100,
    });

match bad_config.validate() {
    Ok(()) => unreachable!(),
    Err(e) => println!("Validation error: {}", e),
}
}

Combining with Fine-Tuning

For best results, allocate training steps for recovery:

#![allow(unused)]
fn main() {
let total_steps = 100000;

// Prune during first 60% of training
let schedule = PruningSchedule::Cubic {
    start_step: 0,
    end_step: 60000,  // 60% of total
    final_sparsity: 0.7,
};

// Remaining 40% for fine-tuning at final sparsity
}