339 lines
10 KiB
Rust
339 lines
10 KiB
Rust
//! Knowledge distillation module for model compression.
|
|
|
|
use forge_shared::{DistillLoss, DistillationConfig, ModelInfo};
|
|
|
|
/// Distiller for knowledge distillation.
|
|
#[derive(Debug)]
|
|
pub struct Distiller {
|
|
/// RNG state.
|
|
#[allow(dead_code)]
|
|
rng_state: u64,
|
|
}
|
|
|
|
impl Default for Distiller {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl Distiller {
|
|
/// Create a new distiller.
|
|
pub fn new() -> Self {
|
|
Self { rng_state: 42 }
|
|
}
|
|
|
|
/// Distill knowledge from teacher to student.
|
|
pub fn distill(&mut self, teacher: &ModelInfo, config: &DistillationConfig) -> ModelInfo {
|
|
// Create student model info
|
|
let student_layers = if config.layer_mapping.is_empty() {
|
|
teacher.num_layers / 2 // Default: half the layers
|
|
} else {
|
|
config.layer_mapping.len()
|
|
};
|
|
|
|
let student_hidden = teacher.hidden_dim;
|
|
let student_params = self.estimate_student_params(teacher, student_layers);
|
|
|
|
ModelInfo {
|
|
name: format!("{}-distilled", teacher.name),
|
|
architecture: teacher.architecture,
|
|
num_parameters: student_params,
|
|
num_layers: student_layers,
|
|
hidden_dim: student_hidden,
|
|
vocab_size: teacher.vocab_size,
|
|
image_size: teacher.image_size,
|
|
precision_bits: 16,
|
|
size_bytes: student_params * 2, // FP16
|
|
}
|
|
}
|
|
|
|
/// Estimate student model parameters.
|
|
fn estimate_student_params(&self, teacher: &ModelInfo, student_layers: usize) -> u64 {
|
|
let layer_ratio = student_layers as f64 / teacher.num_layers as f64;
|
|
|
|
// Params scale roughly with number of layers (plus embeddings)
|
|
let embedding_params =
|
|
teacher.vocab_size.unwrap_or(32000) as u64 * teacher.hidden_dim as u64;
|
|
let layer_params = (teacher.num_parameters - embedding_params * 2) as f64 * layer_ratio;
|
|
|
|
(embedding_params * 2 + layer_params as u64).max(1_000_000)
|
|
}
|
|
|
|
/// Compute distillation loss.
|
|
pub fn compute_loss(
|
|
&self,
|
|
teacher_logits: &[f64],
|
|
student_logits: &[f64],
|
|
config: &DistillationConfig,
|
|
) -> f64 {
|
|
match config.loss_type {
|
|
DistillLoss::KL => {
|
|
self.kl_divergence(teacher_logits, student_logits, config.temperature)
|
|
}
|
|
DistillLoss::MSE => self.mse_loss(teacher_logits, student_logits),
|
|
DistillLoss::Cosine => self.cosine_loss(teacher_logits, student_logits),
|
|
DistillLoss::AttentionTransfer => {
|
|
// Simplified: use MSE on attention patterns
|
|
self.mse_loss(teacher_logits, student_logits) * 0.5
|
|
}
|
|
DistillLoss::FeatureMatching => {
|
|
// Simplified: use MSE on features
|
|
self.mse_loss(teacher_logits, student_logits) * 0.3
|
|
}
|
|
DistillLoss::Contrastive => {
|
|
// Simplified contrastive loss
|
|
self.cosine_loss(teacher_logits, student_logits) * 0.5
|
|
}
|
|
DistillLoss::Combined => {
|
|
let kl = self.kl_divergence(teacher_logits, student_logits, config.temperature);
|
|
let mse = self.mse_loss(teacher_logits, student_logits);
|
|
config.alpha as f64 * kl + (1.0 - config.alpha as f64) * mse
|
|
}
|
|
}
|
|
}
|
|
|
|
/// KL divergence loss.
|
|
fn kl_divergence(&self, teacher: &[f64], student: &[f64], temperature: f32) -> f64 {
|
|
if teacher.is_empty() || student.is_empty() {
|
|
return 0.0;
|
|
}
|
|
|
|
let t = temperature as f64;
|
|
|
|
// Softmax with temperature
|
|
let teacher_soft = self.softmax_with_temp(teacher, t);
|
|
let student_soft = self.softmax_with_temp(student, t);
|
|
|
|
// KL divergence
|
|
let mut kl = 0.0;
|
|
for (p, q) in teacher_soft.iter().zip(student_soft.iter()) {
|
|
if *p > 1e-10 && *q > 1e-10 {
|
|
kl += p * (p.ln() - q.ln());
|
|
}
|
|
}
|
|
|
|
kl * t * t // Scale by T^2 as in original paper
|
|
}
|
|
|
|
/// Softmax with temperature.
|
|
fn softmax_with_temp(&self, logits: &[f64], temperature: f64) -> Vec<f64> {
|
|
if logits.is_empty() {
|
|
return vec![];
|
|
}
|
|
|
|
let scaled: Vec<f64> = logits.iter().map(|&x| x / temperature).collect();
|
|
let max = scaled.iter().copied().fold(f64::NEG_INFINITY, f64::max);
|
|
let exp_sum: f64 = scaled.iter().map(|&x| (x - max).exp()).sum();
|
|
|
|
scaled.iter().map(|&x| (x - max).exp() / exp_sum).collect()
|
|
}
|
|
|
|
/// MSE loss.
|
|
fn mse_loss(&self, teacher: &[f64], student: &[f64]) -> f64 {
|
|
if teacher.is_empty() || student.is_empty() {
|
|
return 0.0;
|
|
}
|
|
|
|
let n = teacher.len().min(student.len());
|
|
let sum: f64 = teacher[..n]
|
|
.iter()
|
|
.zip(student[..n].iter())
|
|
.map(|(t, s)| (t - s).powi(2))
|
|
.sum();
|
|
|
|
sum / n as f64
|
|
}
|
|
|
|
/// Cosine similarity loss.
|
|
fn cosine_loss(&self, teacher: &[f64], student: &[f64]) -> f64 {
|
|
if teacher.is_empty() || student.is_empty() {
|
|
return 1.0;
|
|
}
|
|
|
|
let n = teacher.len().min(student.len());
|
|
|
|
let dot: f64 = teacher[..n]
|
|
.iter()
|
|
.zip(student[..n].iter())
|
|
.map(|(t, s)| t * s)
|
|
.sum();
|
|
|
|
let norm_t: f64 = teacher[..n].iter().map(|t| t * t).sum::<f64>().sqrt();
|
|
let norm_s: f64 = student[..n].iter().map(|s| s * s).sum::<f64>().sqrt();
|
|
|
|
if norm_t > 1e-10 && norm_s > 1e-10 {
|
|
1.0 - dot / (norm_t * norm_s)
|
|
} else {
|
|
1.0
|
|
}
|
|
}
|
|
|
|
/// Estimate accuracy degradation from distillation.
|
|
pub fn estimate_degradation(&self, teacher: &ModelInfo, config: &DistillationConfig) -> f64 {
|
|
let layer_ratio = if config.layer_mapping.is_empty() {
|
|
0.5
|
|
} else {
|
|
config.layer_mapping.len() as f64 / teacher.num_layers as f64
|
|
};
|
|
|
|
// Base degradation from size reduction
|
|
let base_degradation = (1.0 - layer_ratio) * 0.1;
|
|
|
|
// Better distillation methods help
|
|
let method_factor = match config.loss_type {
|
|
DistillLoss::KL => 0.8,
|
|
DistillLoss::Combined => 0.7,
|
|
DistillLoss::AttentionTransfer => 0.6,
|
|
_ => 1.0,
|
|
};
|
|
|
|
// More training helps
|
|
let epoch_factor = 1.0 / (1.0 + config.epochs as f64 * 0.05);
|
|
|
|
// Intermediate matching helps
|
|
let intermediate_factor = if config.intermediate_matching {
|
|
0.8
|
|
} else {
|
|
1.0
|
|
};
|
|
|
|
base_degradation * method_factor * epoch_factor * intermediate_factor
|
|
}
|
|
|
|
/// Create layer mapping for progressive distillation.
|
|
pub fn create_layer_mapping(
|
|
&self,
|
|
teacher_layers: usize,
|
|
student_layers: usize,
|
|
) -> Vec<(usize, usize)> {
|
|
if student_layers >= teacher_layers {
|
|
return (0..student_layers).map(|i| (i, i)).collect();
|
|
}
|
|
|
|
let step = teacher_layers as f64 / student_layers as f64;
|
|
(0..student_layers)
|
|
.map(|i| {
|
|
let teacher_idx = (i as f64 * step) as usize;
|
|
(teacher_idx.min(teacher_layers - 1), i)
|
|
})
|
|
.collect()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_distiller_creation() {
|
|
let distiller = Distiller::new();
|
|
assert_eq!(distiller.rng_state, 42);
|
|
}
|
|
|
|
#[test]
|
|
fn test_distill() {
|
|
let mut distiller = Distiller::new();
|
|
let teacher = forge_shared::sample_model_info();
|
|
let config = DistillationConfig::default();
|
|
|
|
let student = distiller.distill(&teacher, &config);
|
|
|
|
assert!(student.num_layers < teacher.num_layers);
|
|
assert!(student.num_parameters < teacher.num_parameters);
|
|
}
|
|
|
|
#[test]
|
|
fn test_kl_divergence() {
|
|
let distiller = Distiller::new();
|
|
|
|
let teacher = vec![1.0, 2.0, 3.0];
|
|
let student = vec![1.0, 2.0, 3.0];
|
|
|
|
let loss = distiller.kl_divergence(&teacher, &student, 4.0);
|
|
assert!(loss.abs() < 0.01); // Same distributions should have low KL
|
|
}
|
|
|
|
#[test]
|
|
fn test_mse_loss() {
|
|
let distiller = Distiller::new();
|
|
|
|
let teacher = vec![1.0, 2.0, 3.0];
|
|
let student = vec![1.0, 2.0, 3.0];
|
|
|
|
let loss = distiller.mse_loss(&teacher, &student);
|
|
assert_eq!(loss, 0.0);
|
|
|
|
let student2 = vec![2.0, 3.0, 4.0];
|
|
let loss2 = distiller.mse_loss(&teacher, &student2);
|
|
assert_eq!(loss2, 1.0); // Each diff is 1, squared is 1, mean is 1
|
|
}
|
|
|
|
#[test]
|
|
fn test_cosine_loss() {
|
|
let distiller = Distiller::new();
|
|
|
|
// Same direction
|
|
let a = vec![1.0, 0.0, 0.0];
|
|
let b = vec![1.0, 0.0, 0.0];
|
|
let loss = distiller.cosine_loss(&a, &b);
|
|
assert!(loss.abs() < 0.01);
|
|
|
|
// Orthogonal
|
|
let c = vec![0.0, 1.0, 0.0];
|
|
let loss2 = distiller.cosine_loss(&a, &c);
|
|
assert!((loss2 - 1.0).abs() < 0.01);
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_loss_kl() {
|
|
let distiller = Distiller::new();
|
|
let config = DistillationConfig {
|
|
loss_type: DistillLoss::KL,
|
|
temperature: 4.0,
|
|
..Default::default()
|
|
};
|
|
|
|
let teacher = vec![1.0, 2.0, 3.0, 4.0];
|
|
let student = vec![1.5, 2.5, 2.5, 4.0];
|
|
|
|
let loss = distiller.compute_loss(&teacher, &student, &config);
|
|
assert!(loss.is_finite());
|
|
assert!(loss >= 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_estimate_degradation() {
|
|
let distiller = Distiller::new();
|
|
let teacher = forge_shared::sample_model_info();
|
|
|
|
let config_few_epochs = DistillationConfig {
|
|
epochs: 1,
|
|
..Default::default()
|
|
};
|
|
let config_many_epochs = DistillationConfig {
|
|
epochs: 100,
|
|
..Default::default()
|
|
};
|
|
|
|
let deg_few = distiller.estimate_degradation(&teacher, &config_few_epochs);
|
|
let deg_many = distiller.estimate_degradation(&teacher, &config_many_epochs);
|
|
|
|
// More epochs should reduce degradation
|
|
assert!(deg_many < deg_few);
|
|
}
|
|
|
|
#[test]
|
|
fn test_create_layer_mapping() {
|
|
let distiller = Distiller::new();
|
|
|
|
let mapping = distiller.create_layer_mapping(32, 6);
|
|
assert_eq!(mapping.len(), 6);
|
|
|
|
// First student layer should map to early teacher layer
|
|
assert_eq!(mapping[0].1, 0);
|
|
|
|
// Last student layer should map to late teacher layer
|
|
assert!(mapping[5].0 > 20);
|
|
}
|
|
}
|