refactor(ml): extract 8 supervised models into ml-supervised crate (task 8)
Move TFT, Mamba-2, Liquid, TGGN, TLOB, KAN, xLSTM, and Diffusion model implementations to ml-supervised. Bridge files (UnifiedTrainable adapters, Checkpointable impls) stay in ml. Delete AsyncDataLoader (replaced by StreamingDbnLoader + simple .chunks() batching). Remove empty ml-infra scaffold — the remaining ml modules are too tightly coupled for clean extraction, so ml stays as the orchestration facade. - ml-supervised: 234 tests, 0 failures - ml: 1687 tests, 0 failures - Workspace: 0 compilation errors Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
104
crates/ml-supervised/src/diffusion/config.rs
Normal file
104
crates/ml-supervised/src/diffusion/config.rs
Normal file
@@ -0,0 +1,104 @@
|
||||
//! Configuration for Diffusion model (DDPM/DDIM) for price path generation.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Noise schedule type for the diffusion process.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum NoiseSchedule {
|
||||
/// Linear beta schedule from beta_start to beta_end.
|
||||
Linear,
|
||||
/// Cosine schedule (Nichol & Dhariwal) -- smoother noise progression.
|
||||
Cosine,
|
||||
}
|
||||
|
||||
impl Default for NoiseSchedule {
|
||||
fn default() -> Self {
|
||||
Self::Cosine
|
||||
}
|
||||
}
|
||||
|
||||
/// Configuration for a Diffusion model.
|
||||
///
|
||||
/// OOM-safe defaults: small channels, short sequences, few res blocks.
|
||||
/// RTX 3050 Ti (4GB) safe at batch_size <= 32.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct DiffusionConfig {
|
||||
/// Number of diffusion timesteps (training noise levels).
|
||||
pub num_timesteps: usize,
|
||||
/// Number of DDIM sampling steps (inference, << num_timesteps).
|
||||
pub sampling_steps: usize,
|
||||
/// Sequence length of generated price paths.
|
||||
pub seq_len: usize,
|
||||
/// Feature dimension per timestep (1 = univariate price).
|
||||
pub feature_dim: usize,
|
||||
/// Hidden dimension for the denoiser network.
|
||||
pub hidden_dim: usize,
|
||||
/// Number of denoiser layers.
|
||||
pub num_layers: usize,
|
||||
/// Time embedding dimension.
|
||||
pub time_embed_dim: usize,
|
||||
/// Noise schedule type.
|
||||
pub schedule: NoiseSchedule,
|
||||
/// Learning rate.
|
||||
pub learning_rate: f64,
|
||||
/// Weight decay for regularization.
|
||||
pub weight_decay: f64,
|
||||
/// Gradient clipping max norm.
|
||||
pub grad_clip: f64,
|
||||
}
|
||||
|
||||
impl Default for DiffusionConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
num_timesteps: 1000,
|
||||
sampling_steps: 10,
|
||||
seq_len: 64,
|
||||
feature_dim: 1,
|
||||
hidden_dim: 128,
|
||||
num_layers: 3,
|
||||
time_embed_dim: 32,
|
||||
schedule: NoiseSchedule::Cosine,
|
||||
learning_rate: 1e-4,
|
||||
weight_decay: 1e-4,
|
||||
grad_clip: 1.0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl DiffusionConfig {
|
||||
/// Total data dimension (seq_len * feature_dim).
|
||||
pub fn data_dim(&self) -> usize {
|
||||
self.seq_len * self.feature_dim
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_default_config() {
|
||||
let config = DiffusionConfig::default();
|
||||
assert_eq!(config.num_timesteps, 1000);
|
||||
assert_eq!(config.sampling_steps, 10);
|
||||
assert_eq!(config.data_dim(), 64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_data_dim_multivariate() {
|
||||
let config = DiffusionConfig {
|
||||
seq_len: 32,
|
||||
feature_dim: 4,
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(config.data_dim(), 128);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_serde_roundtrip() {
|
||||
let config = DiffusionConfig::default();
|
||||
let json = serde_json::to_string(&config).unwrap_or_default();
|
||||
let restored: Result<DiffusionConfig, _> = serde_json::from_str(&json);
|
||||
assert!(restored.is_ok());
|
||||
}
|
||||
}
|
||||
299
crates/ml-supervised/src/diffusion/denoiser.rs
Normal file
299
crates/ml-supervised/src/diffusion/denoiser.rs
Normal file
@@ -0,0 +1,299 @@
|
||||
//! Denoiser network for the diffusion model.
|
||||
//!
|
||||
//! Uses a fully-connected architecture with sinusoidal time embedding.
|
||||
//! This is more memory-efficient than Conv1D U-Net while still effective
|
||||
//! for price sequence denoising at small sequence lengths (64-128).
|
||||
|
||||
use ml_core::MLError;
|
||||
use candle_core::{DType, Device, Tensor};
|
||||
use candle_nn::{linear, Linear, Module, VarBuilder};
|
||||
|
||||
/// Sinusoidal time embedding for diffusion timestep conditioning.
|
||||
///
|
||||
/// Maps scalar timestep t to a fixed-dimension vector using
|
||||
/// sin/cos positional encoding (same idea as Transformer PE).
|
||||
pub struct TimeEmbedding {
|
||||
proj: Linear,
|
||||
embed_dim: usize,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for TimeEmbedding {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("TimeEmbedding").finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl TimeEmbedding {
|
||||
pub fn new(embed_dim: usize, hidden_dim: usize, vb: VarBuilder<'_>) -> Result<Self, MLError> {
|
||||
let proj = linear(embed_dim, hidden_dim, vb.pp("time_proj"))
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
Ok(Self { proj, embed_dim })
|
||||
}
|
||||
|
||||
/// Compute sinusoidal embedding for timestep indices.
|
||||
/// Input: (batch,) u32 timestep indices → Output: (batch, hidden_dim)
|
||||
pub fn forward(&self, t: &Tensor, device: &Device) -> Result<Tensor, MLError> {
|
||||
let half_dim = self.embed_dim / 2;
|
||||
let t_f32 = t.to_dtype(DType::F32)
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
|
||||
// freq[i] = exp(-ln(10000) * i / half_dim)
|
||||
let mut freq_vals = Vec::with_capacity(half_dim);
|
||||
for i in 0..half_dim {
|
||||
let freq = (-(10000.0_f64.ln()) * i as f64 / half_dim.max(1) as f64).exp();
|
||||
freq_vals.push(freq as f32);
|
||||
}
|
||||
let freqs = Tensor::new(freq_vals, device)
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
|
||||
// (batch, 1) * (1, half_dim) → (batch, half_dim)
|
||||
let t_expanded = t_f32.unsqueeze(1)
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
let freqs_expanded = freqs.unsqueeze(0)
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
let angles = t_expanded.broadcast_mul(&freqs_expanded)
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
|
||||
let sin_emb = angles.sin()
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
let cos_emb = angles.cos()
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
|
||||
// Concat sin and cos: (batch, embed_dim)
|
||||
let emb = Tensor::cat(&[&sin_emb, &cos_emb], 1)
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
|
||||
// Cast to training dtype before projection through BF16 weights
|
||||
let emb = ml_core::mixed_precision::ensure_training_dtype(&emb)
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
|
||||
// Project to hidden_dim
|
||||
self.proj.forward(&emb)
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
/// A single denoiser block: linear → SiLU → linear + time conditioning + residual.
|
||||
struct DenoiserBlock {
|
||||
fc1: Linear,
|
||||
fc2: Linear,
|
||||
time_proj: Linear,
|
||||
has_residual: bool,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for DenoiserBlock {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("DenoiserBlock").finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl DenoiserBlock {
|
||||
fn new(
|
||||
input_dim: usize,
|
||||
hidden_dim: usize,
|
||||
time_dim: usize,
|
||||
vb: VarBuilder<'_>,
|
||||
) -> Result<Self, MLError> {
|
||||
let fc1 = linear(input_dim, hidden_dim, vb.pp("fc1"))
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
let fc2 = linear(hidden_dim, hidden_dim, vb.pp("fc2"))
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
let time_proj = linear(time_dim, hidden_dim, vb.pp("time"))
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
Ok(Self {
|
||||
fc1,
|
||||
fc2,
|
||||
time_proj,
|
||||
has_residual: input_dim == hidden_dim,
|
||||
})
|
||||
}
|
||||
|
||||
fn forward(&self, x: &Tensor, t_emb: &Tensor) -> Result<Tensor, MLError> {
|
||||
let map_err = |e: candle_core::Error| MLError::ModelError(e.to_string());
|
||||
|
||||
// fc1 → SiLU (x * sigmoid(x))
|
||||
let h = self.fc1.forward(x).map_err(map_err)?;
|
||||
let h_sig = candle_nn::ops::sigmoid(&h).map_err(map_err)?;
|
||||
let h = h.mul(&h_sig).map_err(map_err)?;
|
||||
|
||||
// Add time embedding
|
||||
let t_proj = self.time_proj.forward(t_emb).map_err(map_err)?;
|
||||
let h = h.add(&t_proj).map_err(map_err)?;
|
||||
|
||||
// fc2 → SiLU
|
||||
let h = self.fc2.forward(&h).map_err(map_err)?;
|
||||
let h_sig2 = candle_nn::ops::sigmoid(&h).map_err(map_err)?;
|
||||
let h = h.mul(&h_sig2).map_err(map_err)?;
|
||||
|
||||
// Residual connection
|
||||
if self.has_residual {
|
||||
h.add(x).map_err(map_err)
|
||||
} else {
|
||||
Ok(h)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Fully-connected denoiser network for diffusion.
|
||||
///
|
||||
/// Architecture: input projection → N denoiser blocks → output projection.
|
||||
/// Each block is time-conditioned via additive time embedding.
|
||||
///
|
||||
/// Memory usage at batch_size=32, data_dim=64, hidden=128:
|
||||
/// ~32 * 128 * num_layers * 4 bytes per layer ≈ 48KB -- very safe for 4GB GPU.
|
||||
pub struct Denoiser {
|
||||
input_proj: Linear,
|
||||
blocks: Vec<DenoiserBlock>,
|
||||
output_proj: Linear,
|
||||
time_embed: TimeEmbedding,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for Denoiser {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("Denoiser").finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl Denoiser {
|
||||
pub fn new(
|
||||
data_dim: usize,
|
||||
hidden_dim: usize,
|
||||
num_layers: usize,
|
||||
time_embed_dim: usize,
|
||||
vb: VarBuilder<'_>,
|
||||
device: &Device,
|
||||
) -> Result<Self, MLError> {
|
||||
if data_dim == 0 || hidden_dim == 0 {
|
||||
return Err(MLError::ConfigError(format!(
|
||||
"Diffusion denoiser requires data_dim > 0 and hidden_dim > 0 (got {}x{})",
|
||||
data_dim, hidden_dim
|
||||
)));
|
||||
}
|
||||
|
||||
let time_embed = TimeEmbedding::new(time_embed_dim, hidden_dim, vb.pp("time_embed"))?;
|
||||
|
||||
let input_proj = linear(data_dim, hidden_dim, vb.pp("input_proj"))
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
|
||||
let mut blocks = Vec::with_capacity(num_layers);
|
||||
for i in 0..num_layers {
|
||||
let block = DenoiserBlock::new(
|
||||
hidden_dim,
|
||||
hidden_dim,
|
||||
hidden_dim,
|
||||
vb.pp(format!("block_{i}")),
|
||||
)?;
|
||||
blocks.push(block);
|
||||
}
|
||||
|
||||
let output_proj = linear(hidden_dim, data_dim, vb.pp("output_proj"))
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
|
||||
Ok(Self {
|
||||
input_proj,
|
||||
blocks,
|
||||
output_proj,
|
||||
time_embed,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Predict noise epsilon given noisy input x_t and timestep t.
|
||||
///
|
||||
/// Input x: (batch, data_dim), t: (batch,) → Output: (batch, data_dim)
|
||||
pub fn forward(&self, x: &Tensor, t: &Tensor) -> Result<Tensor, MLError> {
|
||||
let x = ml_core::mixed_precision::ensure_training_dtype(x)
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
let map_err = |e: candle_core::Error| MLError::ModelError(e.to_string());
|
||||
|
||||
// Time embedding
|
||||
let t_emb = self.time_embed.forward(t, &self.device)?;
|
||||
|
||||
// Input projection → SiLU
|
||||
let mut h = self.input_proj.forward(&x).map_err(map_err)?;
|
||||
let h_sig = candle_nn::ops::sigmoid(&h).map_err(map_err)?;
|
||||
h = h.mul(&h_sig).map_err(map_err)?;
|
||||
|
||||
// Denoiser blocks
|
||||
for block in &self.blocks {
|
||||
h = block.forward(&h, &t_emb)?;
|
||||
}
|
||||
|
||||
// Output projection → predicted noise
|
||||
let output = self.output_proj.forward(&h).map_err(map_err)?;
|
||||
// Cast output back to F32 for API compatibility
|
||||
output.to_dtype(candle_core::DType::F32).map_err(map_err)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use candle_nn::VarMap;
|
||||
use ml_core::mixed_precision::training_dtype;
|
||||
|
||||
#[test]
|
||||
fn test_time_embedding_shape() {
|
||||
let dev = Device::Cpu;
|
||||
let var_map = VarMap::new();
|
||||
let vb = VarBuilder::from_varmap(&var_map, training_dtype(&dev), &dev);
|
||||
let te = TimeEmbedding::new(32, 64, vb).unwrap();
|
||||
let t = Tensor::new(&[0_u32, 100, 500, 999], &dev).unwrap();
|
||||
let emb = te.forward(&t, &dev).unwrap();
|
||||
assert_eq!(emb.dims(), &[4, 64]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_time_embedding_different_timesteps_differ() {
|
||||
let dev = Device::Cpu;
|
||||
let var_map = VarMap::new();
|
||||
let vb = VarBuilder::from_varmap(&var_map, training_dtype(&dev), &dev);
|
||||
let te = TimeEmbedding::new(32, 64, vb).unwrap();
|
||||
let t1 = Tensor::new(&[0_u32], &dev).unwrap();
|
||||
let t2 = Tensor::new(&[500_u32], &dev).unwrap();
|
||||
let e1 = te.forward(&t1, &dev).unwrap();
|
||||
let e2 = te.forward(&t2, &dev).unwrap();
|
||||
let diff: f32 = e1.sub(&e2).unwrap().abs().unwrap().sum_all().unwrap().to_scalar().unwrap();
|
||||
assert!(diff > 0.0, "Different timesteps should produce different embeddings");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_denoiser_output_shape() {
|
||||
let dev = Device::Cpu;
|
||||
let var_map = VarMap::new();
|
||||
let vb = VarBuilder::from_varmap(&var_map, training_dtype(&dev), &dev);
|
||||
let denoiser = Denoiser::new(64, 128, 3, 32, vb, &dev).unwrap();
|
||||
let x = Tensor::randn(0_f32, 1.0, &[4, 64], &dev).unwrap();
|
||||
let t = Tensor::new(&[100_u32, 200, 300, 400], &dev).unwrap();
|
||||
let out = denoiser.forward(&x, &t).unwrap();
|
||||
assert_eq!(out.dims(), &[4, 64], "Output should match input shape");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_denoiser_produces_gradients() {
|
||||
let dev = Device::Cpu;
|
||||
let var_map = VarMap::new();
|
||||
let vb = VarBuilder::from_varmap(&var_map, training_dtype(&dev), &dev);
|
||||
let denoiser = Denoiser::new(64, 128, 2, 32, vb, &dev).unwrap();
|
||||
let x = Tensor::randn(0_f32, 1.0, &[4, 64], &dev).unwrap();
|
||||
let t = Tensor::new(&[50_u32, 100, 200, 300], &dev).unwrap();
|
||||
let out = denoiser.forward(&x, &t).unwrap();
|
||||
let loss = out.sqr().unwrap().mean_all().unwrap();
|
||||
let grads = loss.backward().unwrap();
|
||||
let has_grads = var_map.all_vars().iter().any(|v| grads.get(v.as_tensor()).is_some());
|
||||
assert!(has_grads, "Should produce gradients for training");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_denoiser_single_layer() {
|
||||
let dev = Device::Cpu;
|
||||
let var_map = VarMap::new();
|
||||
let vb = VarBuilder::from_varmap(&var_map, training_dtype(&dev), &dev);
|
||||
let denoiser = Denoiser::new(32, 64, 1, 16, vb, &dev).unwrap();
|
||||
let x = Tensor::randn(0_f32, 1.0, &[2, 32], &dev).unwrap();
|
||||
let t = Tensor::new(&[0_u32, 999], &dev).unwrap();
|
||||
let out = denoiser.forward(&x, &t).unwrap();
|
||||
assert_eq!(out.dims(), &[2, 32]);
|
||||
}
|
||||
}
|
||||
17
crates/ml-supervised/src/diffusion/mod.rs
Normal file
17
crates/ml-supervised/src/diffusion/mod.rs
Normal file
@@ -0,0 +1,17 @@
|
||||
//! Diffusion model (DDPM/DDIM) for price path generation.
|
||||
//!
|
||||
//! Implements a denoising diffusion probabilistic model with:
|
||||
//! - Cosine/linear noise schedules
|
||||
//! - Fully-connected denoiser with sinusoidal time embedding
|
||||
//! - DDIM sampling for fast inference
|
||||
//! - UnifiedTrainable adapter for the training pipeline
|
||||
|
||||
pub mod config;
|
||||
pub mod denoiser;
|
||||
pub mod noise;
|
||||
pub mod sampler;
|
||||
|
||||
pub use config::{DiffusionConfig, NoiseSchedule};
|
||||
pub use denoiser::Denoiser;
|
||||
pub use noise::NoiseScheduler;
|
||||
pub use sampler::DDIMSampler;
|
||||
220
crates/ml-supervised/src/diffusion/noise.rs
Normal file
220
crates/ml-supervised/src/diffusion/noise.rs
Normal file
@@ -0,0 +1,220 @@
|
||||
//! Noise scheduler for the diffusion process.
|
||||
//!
|
||||
//! Precomputes alpha_bar_t for all timesteps and provides
|
||||
//! forward process (add noise) operations.
|
||||
|
||||
use ml_core::MLError;
|
||||
use candle_core::{Device, Tensor};
|
||||
|
||||
use super::config::NoiseSchedule;
|
||||
|
||||
/// Precomputed noise schedule for the diffusion process.
|
||||
///
|
||||
/// Stores alpha_bar_t (cumulative product of (1 - beta_t)) for
|
||||
/// all T timesteps, enabling efficient forward-process noise addition.
|
||||
pub struct NoiseScheduler {
|
||||
/// Cumulative alpha products: alpha_bar_t for each timestep.
|
||||
alpha_bars: Vec<f32>,
|
||||
num_timesteps: usize,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for NoiseScheduler {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("NoiseScheduler").finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl NoiseScheduler {
|
||||
/// Create a new noise scheduler with precomputed schedule.
|
||||
pub fn new(
|
||||
num_timesteps: usize,
|
||||
schedule: &NoiseSchedule,
|
||||
device: &Device,
|
||||
) -> Result<Self, MLError> {
|
||||
if num_timesteps == 0 {
|
||||
return Err(MLError::ConfigError("num_timesteps must be > 0".to_owned()));
|
||||
}
|
||||
|
||||
let alpha_bars = match schedule {
|
||||
NoiseSchedule::Linear => Self::linear_schedule(num_timesteps),
|
||||
NoiseSchedule::Cosine => Self::cosine_schedule(num_timesteps),
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
alpha_bars,
|
||||
num_timesteps,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Linear beta schedule: beta linearly from 1e-4 to 0.02.
|
||||
fn linear_schedule(t_max: usize) -> Vec<f32> {
|
||||
let beta_start = 1e-4_f64;
|
||||
let beta_end = 0.02_f64;
|
||||
let mut alpha_bars = Vec::with_capacity(t_max);
|
||||
let mut cumulative = 1.0_f64;
|
||||
for i in 0..t_max {
|
||||
let beta = beta_start + (beta_end - beta_start) * (i as f64 / (t_max - 1).max(1) as f64);
|
||||
cumulative *= 1.0 - beta;
|
||||
alpha_bars.push(cumulative as f32);
|
||||
}
|
||||
alpha_bars
|
||||
}
|
||||
|
||||
/// Cosine schedule (Nichol & Dhariwal 2021).
|
||||
/// alpha_bar_t = cos((t/T + s) / (1+s) * pi/2)^2
|
||||
fn cosine_schedule(t_max: usize) -> Vec<f32> {
|
||||
let s = 0.008_f64; // Small offset to prevent beta_0 = 0
|
||||
let mut alpha_bars = Vec::with_capacity(t_max);
|
||||
for i in 0..t_max {
|
||||
let t_frac = i as f64 / t_max as f64;
|
||||
let angle = (t_frac + s) / (1.0 + s) * std::f64::consts::FRAC_PI_2;
|
||||
let alpha_bar = angle.cos().powi(2);
|
||||
alpha_bars.push(alpha_bar as f32);
|
||||
}
|
||||
// Normalize so alpha_bar[0] ≈ 1.0
|
||||
let first = alpha_bars.first().copied().unwrap_or(1.0);
|
||||
if first > 0.0 {
|
||||
for ab in &mut alpha_bars {
|
||||
*ab /= first;
|
||||
}
|
||||
}
|
||||
alpha_bars
|
||||
}
|
||||
|
||||
/// Get alpha_bar for a specific timestep.
|
||||
pub fn get_alpha_bar(&self, t: usize) -> Result<f32, MLError> {
|
||||
self.alpha_bars
|
||||
.get(t)
|
||||
.copied()
|
||||
.ok_or_else(|| MLError::ConfigError(format!("Timestep {} out of range (max {})", t, self.num_timesteps)))
|
||||
}
|
||||
|
||||
/// Forward process: add noise to clean data x0 at timestep t.
|
||||
///
|
||||
/// q(x_t | x_0) = N(sqrt(alpha_bar_t) * x_0, (1 - alpha_bar_t) * I)
|
||||
///
|
||||
/// Returns (noisy_x, noise) where noise is the sampled epsilon.
|
||||
pub fn add_noise(
|
||||
&self,
|
||||
x0: &Tensor,
|
||||
t: usize,
|
||||
) -> Result<(Tensor, Tensor), MLError> {
|
||||
let alpha_bar = self.get_alpha_bar(t)?;
|
||||
let sqrt_alpha_bar = alpha_bar.sqrt();
|
||||
let sqrt_one_minus_alpha_bar = (1.0_f32 - alpha_bar).max(0.0_f32).sqrt();
|
||||
|
||||
// Sample noise epsilon ~ N(0, I)
|
||||
let noise = Tensor::randn(0_f32, 1.0, x0.dims(), &self.device)
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
|
||||
// x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon
|
||||
let signal = x0.affine(sqrt_alpha_bar as f64, 0.0)
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
let noise_scaled = noise.affine(sqrt_one_minus_alpha_bar as f64, 0.0)
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
let noisy = signal.add(&noise_scaled)
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
|
||||
Ok((noisy, noise))
|
||||
}
|
||||
|
||||
/// Get alpha_bar as a tensor (scalar) for a given timestep.
|
||||
pub fn alpha_bar_tensor(&self, t: usize) -> Result<Tensor, MLError> {
|
||||
let ab = self.get_alpha_bar(t)?;
|
||||
Tensor::new(&[ab], &self.device)
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))
|
||||
}
|
||||
|
||||
/// Number of timesteps.
|
||||
pub fn num_timesteps(&self) -> usize {
|
||||
self.num_timesteps
|
||||
}
|
||||
|
||||
/// Get alpha_bars slice for DDIM sampling.
|
||||
pub fn alpha_bars(&self) -> &[f32] {
|
||||
&self.alpha_bars
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use candle_core::DType;
|
||||
|
||||
#[test]
|
||||
fn test_linear_schedule_decreasing() {
|
||||
let sched = NoiseScheduler::new(1000, &NoiseSchedule::Linear, &Device::Cpu).unwrap();
|
||||
let a0 = sched.get_alpha_bar(0).unwrap();
|
||||
let a500 = sched.get_alpha_bar(500).unwrap();
|
||||
let a999 = sched.get_alpha_bar(999).unwrap();
|
||||
assert!(a0 > a500, "a0={a0} should be > a500={a500}");
|
||||
assert!(a500 > a999, "a500={a500} should be > a999={a999}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cosine_schedule_decreasing() {
|
||||
let sched = NoiseScheduler::new(1000, &NoiseSchedule::Cosine, &Device::Cpu).unwrap();
|
||||
let a0 = sched.get_alpha_bar(0).unwrap();
|
||||
let a500 = sched.get_alpha_bar(500).unwrap();
|
||||
let a999 = sched.get_alpha_bar(999).unwrap();
|
||||
assert!(a0 > a500, "a0={a0} should be > a500={a500}");
|
||||
assert!(a500 > a999, "a500={a500} should be > a999={a999}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_alpha_bar_first_near_one() {
|
||||
let sched = NoiseScheduler::new(1000, &NoiseSchedule::Cosine, &Device::Cpu).unwrap();
|
||||
let a0 = sched.get_alpha_bar(0).unwrap();
|
||||
assert!((a0 - 1.0).abs() < 0.05, "First alpha_bar should be near 1.0, got {a0}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_alpha_bar_last_near_zero() {
|
||||
let sched = NoiseScheduler::new(1000, &NoiseSchedule::Cosine, &Device::Cpu).unwrap();
|
||||
let a_last = sched.get_alpha_bar(999).unwrap();
|
||||
assert!(a_last < 0.1, "Last alpha_bar should be near 0, got {a_last}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_add_noise_preserves_shape() {
|
||||
let sched = NoiseScheduler::new(1000, &NoiseSchedule::Cosine, &Device::Cpu).unwrap();
|
||||
let x = Tensor::ones(&[4, 64], DType::F32, &Device::Cpu).unwrap();
|
||||
let (noisy, noise) = sched.add_noise(&x, 100).unwrap();
|
||||
assert_eq!(noisy.dims(), &[4, 64]);
|
||||
assert_eq!(noise.dims(), &[4, 64]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_add_noise_at_t0_preserves_signal() {
|
||||
let sched = NoiseScheduler::new(1000, &NoiseSchedule::Cosine, &Device::Cpu).unwrap();
|
||||
let x = Tensor::ones(&[4, 64], DType::F32, &Device::Cpu).unwrap();
|
||||
let (noisy, _noise) = sched.add_noise(&x, 0).unwrap();
|
||||
let diff: f32 = noisy.sub(&x).unwrap().abs().unwrap().mean_all().unwrap()
|
||||
.to_scalar().unwrap();
|
||||
assert!(diff < 0.5, "t=0 should add minimal noise, got diff={diff}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_add_noise_at_high_t_destroys_signal() {
|
||||
let sched = NoiseScheduler::new(1000, &NoiseSchedule::Cosine, &Device::Cpu).unwrap();
|
||||
let x = Tensor::ones(&[4, 64], DType::F32, &Device::Cpu).unwrap();
|
||||
let (noisy, _noise) = sched.add_noise(&x, 999).unwrap();
|
||||
// At high t, noisy should be mostly noise (mean near 0)
|
||||
let mean: f32 = noisy.mean_all().unwrap().to_scalar().unwrap();
|
||||
assert!(mean.abs() < 1.5, "At t=999, signal should be destroyed, mean={mean}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_out_of_range_errors() {
|
||||
let sched = NoiseScheduler::new(100, &NoiseSchedule::Linear, &Device::Cpu).unwrap();
|
||||
assert!(sched.get_alpha_bar(100).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_zero_timesteps_errors() {
|
||||
let result = NoiseScheduler::new(0, &NoiseSchedule::Linear, &Device::Cpu);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
}
|
||||
237
crates/ml-supervised/src/diffusion/sampler.rs
Normal file
237
crates/ml-supervised/src/diffusion/sampler.rs
Normal file
@@ -0,0 +1,237 @@
|
||||
//! DDIM (Denoising Diffusion Implicit Models) sampler.
|
||||
//!
|
||||
//! Provides deterministic, fast sampling from a trained diffusion model
|
||||
//! using a small number of steps (e.g., 10) instead of the full T=1000.
|
||||
|
||||
use ml_core::MLError;
|
||||
use candle_core::{Device, Tensor};
|
||||
|
||||
use super::denoiser::Denoiser;
|
||||
use super::noise::NoiseScheduler;
|
||||
|
||||
/// DDIM sampler for fast, deterministic inference.
|
||||
///
|
||||
/// Given a trained denoiser and noise scheduler, generates samples
|
||||
/// by iteratively denoising from pure noise using uniformly spaced
|
||||
/// timestep subsequence.
|
||||
#[derive(Debug)]
|
||||
pub struct DDIMSampler {
|
||||
/// Number of DDIM steps (much less than training timesteps).
|
||||
num_steps: usize,
|
||||
/// Total training timesteps.
|
||||
num_timesteps: usize,
|
||||
/// DDIM eta parameter: 0.0 = deterministic, 1.0 = DDPM.
|
||||
eta: f32,
|
||||
}
|
||||
|
||||
impl DDIMSampler {
|
||||
/// Create a new DDIM sampler.
|
||||
///
|
||||
/// `num_steps`: number of denoising steps (e.g., 10-50).
|
||||
/// `num_timesteps`: total training timesteps (e.g., 1000).
|
||||
/// `eta`: stochasticity parameter (0.0 for deterministic).
|
||||
pub fn new(num_steps: usize, num_timesteps: usize, eta: f32) -> Self {
|
||||
Self {
|
||||
num_steps: num_steps.max(1),
|
||||
num_timesteps,
|
||||
eta,
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute the uniformly spaced timestep subsequence for DDIM.
|
||||
fn timestep_sequence(&self) -> Vec<usize> {
|
||||
let mut seq = Vec::with_capacity(self.num_steps);
|
||||
for i in 0..self.num_steps {
|
||||
// Reverse order: from T-1 down to 0
|
||||
let t = ((self.num_timesteps - 1) as f64 * (1.0 - i as f64 / self.num_steps.max(1) as f64)) as usize;
|
||||
seq.push(t.min(self.num_timesteps - 1));
|
||||
}
|
||||
seq
|
||||
}
|
||||
|
||||
/// Generate samples from pure noise using DDIM.
|
||||
///
|
||||
/// `denoiser`: the trained noise prediction network.
|
||||
/// `scheduler`: noise schedule with alpha_bar values.
|
||||
/// `num_samples`: number of samples to generate.
|
||||
/// `data_dim`: dimension of each sample.
|
||||
/// `device`: computation device.
|
||||
///
|
||||
/// Returns: (num_samples, data_dim) tensor.
|
||||
pub fn sample(
|
||||
&self,
|
||||
denoiser: &Denoiser,
|
||||
scheduler: &NoiseScheduler,
|
||||
num_samples: usize,
|
||||
data_dim: usize,
|
||||
device: &Device,
|
||||
) -> Result<Tensor, MLError> {
|
||||
let map_err = |e: candle_core::Error| MLError::ModelError(e.to_string());
|
||||
|
||||
// Start from pure noise
|
||||
let mut x_t = Tensor::randn(0_f32, 1.0, &[num_samples, data_dim], device)
|
||||
.map_err(map_err)?;
|
||||
|
||||
let timesteps = self.timestep_sequence();
|
||||
|
||||
for (i, &t) in timesteps.iter().enumerate() {
|
||||
// Create timestep tensor for batch
|
||||
let t_tensor = Tensor::new(
|
||||
vec![t as u32; num_samples],
|
||||
device,
|
||||
).map_err(map_err)?;
|
||||
|
||||
// Predict noise
|
||||
let eps_pred = denoiser.forward(&x_t, &t_tensor)?;
|
||||
|
||||
let alpha_bar_t: f32 = scheduler.get_alpha_bar(t)?;
|
||||
|
||||
// Get alpha_bar for next timestep (or 1.0 at the end)
|
||||
let alpha_bar_prev: f32 = if i + 1 < timesteps.len() {
|
||||
let t_prev = timesteps.get(i + 1).copied().unwrap_or(0);
|
||||
scheduler.get_alpha_bar(t_prev)?
|
||||
} else {
|
||||
1.0_f32 // Final step: fully denoised
|
||||
};
|
||||
|
||||
// DDIM update rule:
|
||||
// predicted x_0 = (x_t - sqrt(1 - alpha_bar_t) * eps) / sqrt(alpha_bar_t)
|
||||
let sqrt_ab = alpha_bar_t.sqrt();
|
||||
let sqrt_one_minus_ab = (1.0_f32 - alpha_bar_t).max(0.0_f32).sqrt();
|
||||
|
||||
// Predicted clean sample:
|
||||
// x0 = (x_t - sqrt(1-ab)*eps) / sqrt(ab)
|
||||
let scale = 1.0_f64 / sqrt_ab.max(1e-8) as f64;
|
||||
let noise_sub = eps_pred.affine(sqrt_one_minus_ab as f64, 0.0).map_err(map_err)?;
|
||||
let x0_pred = x_t.sub(&noise_sub).map_err(map_err)?
|
||||
.affine(scale, 0.0).map_err(map_err)?;
|
||||
|
||||
// Direction pointing to x_t
|
||||
let sqrt_one_minus_ab_prev = (1.0_f32 - alpha_bar_prev).max(0.0_f32).sqrt();
|
||||
|
||||
// Compute sigma for stochasticity
|
||||
let sigma: f32 = if self.eta > 0.0_f32 {
|
||||
let ratio = (1.0_f32 - alpha_bar_prev) / (1.0_f32 - alpha_bar_t).max(1e-8_f32);
|
||||
let beta_t = 1.0_f32 - alpha_bar_t / alpha_bar_prev.max(1e-8_f32);
|
||||
self.eta * (ratio * beta_t).max(0.0_f32).sqrt()
|
||||
} else {
|
||||
0.0_f32
|
||||
};
|
||||
|
||||
let dir_coeff = (sqrt_one_minus_ab_prev.powi(2) - sigma.powi(2)).max(0.0_f32).sqrt();
|
||||
|
||||
// x_{t-1} = sqrt(alpha_bar_{t-1}) * x0_pred + dir_coeff * eps_pred + sigma * noise
|
||||
let sqrt_ab_prev = alpha_bar_prev.sqrt();
|
||||
let term1 = x0_pred.affine(sqrt_ab_prev as f64, 0.0).map_err(map_err)?;
|
||||
let term2 = eps_pred.affine(dir_coeff as f64, 0.0).map_err(map_err)?;
|
||||
x_t = term1.add(&term2).map_err(map_err)?;
|
||||
|
||||
// Add stochastic noise if eta > 0
|
||||
if sigma > 1e-8_f32 {
|
||||
let noise = Tensor::randn(0_f32, 1.0, &[num_samples, data_dim], device)
|
||||
.map_err(map_err)?;
|
||||
let noise_term = noise.affine(sigma as f64, 0.0).map_err(map_err)?;
|
||||
x_t = x_t.add(&noise_term).map_err(map_err)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(x_t)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use super::super::config::{DiffusionConfig, NoiseSchedule};
|
||||
use candle_nn::{VarBuilder, VarMap};
|
||||
use ml_core::mixed_precision::training_dtype;
|
||||
|
||||
fn make_test_components() -> (Denoiser, NoiseScheduler, DDIMSampler) {
|
||||
let dev = Device::Cpu;
|
||||
let config = DiffusionConfig {
|
||||
num_timesteps: 100, // Small for fast tests
|
||||
sampling_steps: 5,
|
||||
seq_len: 16,
|
||||
feature_dim: 1,
|
||||
hidden_dim: 32,
|
||||
num_layers: 1,
|
||||
time_embed_dim: 16,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let var_map = VarMap::new();
|
||||
let vb = VarBuilder::from_varmap(&var_map, training_dtype(&dev), &dev);
|
||||
let denoiser = Denoiser::new(
|
||||
config.data_dim(),
|
||||
config.hidden_dim,
|
||||
config.num_layers,
|
||||
config.time_embed_dim,
|
||||
vb,
|
||||
&dev,
|
||||
).unwrap();
|
||||
|
||||
let scheduler = NoiseScheduler::new(
|
||||
config.num_timesteps,
|
||||
&config.schedule,
|
||||
&dev,
|
||||
).unwrap();
|
||||
|
||||
let sampler = DDIMSampler::new(
|
||||
config.sampling_steps,
|
||||
config.num_timesteps,
|
||||
0.0, // Deterministic
|
||||
);
|
||||
|
||||
(denoiser, scheduler, sampler)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_timestep_sequence_decreasing() {
|
||||
let sampler = DDIMSampler::new(5, 100, 0.0);
|
||||
let seq = sampler.timestep_sequence();
|
||||
assert_eq!(seq.len(), 5);
|
||||
// Should be decreasing
|
||||
for i in 1..seq.len() {
|
||||
assert!(seq.get(i).copied().unwrap_or(0) <= seq.get(i - 1).copied().unwrap_or(0),
|
||||
"Timesteps should be decreasing: {:?}", seq);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ddim_sample_shape() {
|
||||
let (denoiser, scheduler, sampler) = make_test_components();
|
||||
let samples = sampler.sample(&denoiser, &scheduler, 8, 16, &Device::Cpu).unwrap();
|
||||
assert_eq!(samples.dims(), &[8, 16]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ddim_sample_finite() {
|
||||
let (denoiser, scheduler, sampler) = make_test_components();
|
||||
let samples = sampler.sample(&denoiser, &scheduler, 4, 16, &Device::Cpu).unwrap();
|
||||
let vals: Vec<f32> = samples.flatten_all().unwrap().to_vec1().unwrap();
|
||||
for v in &vals {
|
||||
assert!(v.is_finite(), "Sample contains non-finite value: {v}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ddim_deterministic_eta0() {
|
||||
// With eta=0, two runs from same initial noise should give same result
|
||||
// (But we can't control the initial noise easily, so just verify it runs)
|
||||
let (denoiser, scheduler, sampler) = make_test_components();
|
||||
let s1 = sampler.sample(&denoiser, &scheduler, 2, 16, &Device::Cpu);
|
||||
assert!(s1.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ddim_single_step() {
|
||||
let dev = Device::Cpu;
|
||||
let var_map = VarMap::new();
|
||||
let vb = VarBuilder::from_varmap(&var_map, training_dtype(&dev), &dev);
|
||||
let denoiser = Denoiser::new(8, 16, 1, 8, vb, &dev).unwrap();
|
||||
let scheduler = NoiseScheduler::new(10, &NoiseSchedule::Linear, &dev).unwrap();
|
||||
let sampler = DDIMSampler::new(1, 10, 0.0);
|
||||
let result = sampler.sample(&denoiser, &scheduler, 2, 8, &dev);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user