diff --git a/common/src/model_types.rs b/common/src/model_types.rs index db744103a..92d19d4e0 100644 --- a/common/src/model_types.rs +++ b/common/src/model_types.rs @@ -46,6 +46,8 @@ pub enum ModelType { KAN, /// Extended LSTM (sLSTM + mLSTM) XLSTM, + /// Diffusion model (DDPM/DDIM) for price path generation + Diffusion, } impl ModelType { @@ -66,6 +68,7 @@ impl ModelType { Self::Ensemble => "ensemble", Self::KAN => "kan", Self::XLSTM => "xlstm", + Self::Diffusion => "diffusion", } } @@ -84,6 +87,7 @@ impl ModelType { Self::Ensemble => "ensemble", Self::KAN => "kan", Self::XLSTM => "xlstm", + Self::Diffusion => "diffusion", } } @@ -102,6 +106,7 @@ impl ModelType { Self::Transformer => "transformer", Self::KAN => "kan", Self::XLSTM => "xlstm", + Self::Diffusion => "diffusion", } } @@ -122,6 +127,7 @@ impl ModelType { Self::Ensemble => "ENSEMBLE", Self::KAN => "KAN", Self::XLSTM => "XLSTM", + Self::Diffusion => "DIFFUSION", } } @@ -148,6 +154,7 @@ impl ModelType { "ensemble" => Some(Self::Ensemble), "kan" => Some(Self::KAN), "xlstm" | "x-lstm" | "extended_lstm" => Some(Self::XLSTM), + "diffusion" | "ddpm" | "ddim" => Some(Self::Diffusion), _ => None, } } diff --git a/ml/src/diffusion/config.rs b/ml/src/diffusion/config.rs new file mode 100644 index 000000000..1852b5050 --- /dev/null +++ b/ml/src/diffusion/config.rs @@ -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 = serde_json::from_str(&json); + assert!(restored.is_ok()); + } +} diff --git a/ml/src/diffusion/denoiser.rs b/ml/src/diffusion/denoiser.rs new file mode 100644 index 000000000..4678bc70b --- /dev/null +++ b/ml/src/diffusion/denoiser.rs @@ -0,0 +1,265 @@ +//! 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 crate::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 TimeEmbedding { + pub fn new(embed_dim: usize, hidden_dim: usize, vb: VarBuilder) -> Result { + 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 { + 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()))?; + + // 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 DenoiserBlock { + fn new( + input_dim: usize, + hidden_dim: usize, + time_dim: usize, + vb: VarBuilder, + ) -> Result { + 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 { + 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, + output_proj: Linear, + time_embed: TimeEmbedding, + device: Device, +} + +impl Denoiser { + pub fn new( + data_dim: usize, + hidden_dim: usize, + num_layers: usize, + time_embed_dim: usize, + vb: VarBuilder, + device: &Device, + ) -> Result { + 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 { + 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 + self.output_proj.forward(&h).map_err(map_err) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use candle_nn::VarMap; + + #[test] + fn test_time_embedding_shape() { + let dev = Device::Cpu; + let var_map = VarMap::new(); + let vb = VarBuilder::from_varmap(&var_map, DType::F32, &dev); + let te = TimeEmbedding::new(32, 64, vb).unwrap(); + let t = Tensor::new(&[0u32, 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, DType::F32, &dev); + let te = TimeEmbedding::new(32, 64, vb).unwrap(); + let t1 = Tensor::new(&[0u32], &dev).unwrap(); + let t2 = Tensor::new(&[500u32], &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, DType::F32, &dev); + let denoiser = Denoiser::new(64, 128, 3, 32, vb, &dev).unwrap(); + let x = Tensor::randn(0f32, 1.0, &[4, 64], &dev).unwrap(); + let t = Tensor::new(&[100u32, 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, DType::F32, &dev); + let denoiser = Denoiser::new(64, 128, 2, 32, vb, &dev).unwrap(); + let x = Tensor::randn(0f32, 1.0, &[4, 64], &dev).unwrap(); + let t = Tensor::new(&[50u32, 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, DType::F32, &dev); + let denoiser = Denoiser::new(32, 64, 1, 16, vb, &dev).unwrap(); + let x = Tensor::randn(0f32, 1.0, &[2, 32], &dev).unwrap(); + let t = Tensor::new(&[0u32, 999], &dev).unwrap(); + let out = denoiser.forward(&x, &t).unwrap(); + assert_eq!(out.dims(), &[2, 32]); + } +} diff --git a/ml/src/diffusion/mod.rs b/ml/src/diffusion/mod.rs new file mode 100644 index 000000000..d4c8f6d7d --- /dev/null +++ b/ml/src/diffusion/mod.rs @@ -0,0 +1,19 @@ +//! 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 mod trainable; + +pub use config::{DiffusionConfig, NoiseSchedule}; +pub use denoiser::Denoiser; +pub use noise::NoiseScheduler; +pub use sampler::DDIMSampler; +pub use trainable::DiffusionTrainableAdapter; diff --git a/ml/src/diffusion/noise.rs b/ml/src/diffusion/noise.rs new file mode 100644 index 000000000..8381e1996 --- /dev/null +++ b/ml/src/diffusion/noise.rs @@ -0,0 +1,217 @@ +//! Noise scheduler for the diffusion process. +//! +//! Precomputes alpha_bar_t for all timesteps and provides +//! forward process (add noise) operations. + +use crate::MLError; +use candle_core::{DType, 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, + num_timesteps: usize, + device: Device, +} + +impl NoiseScheduler { + /// Create a new noise scheduler with precomputed schedule. + pub fn new( + num_timesteps: usize, + schedule: &NoiseSchedule, + device: &Device, + ) -> Result { + if num_timesteps == 0 { + return Err(MLError::ConfigError { + reason: "num_timesteps must be > 0".to_string(), + }); + } + + 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 { + 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 { + 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 { + self.alpha_bars + .get(t) + .copied() + .ok_or_else(|| MLError::ConfigError { + reason: 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(0f32, 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 { + 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::*; + + #[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()); + } +} diff --git a/ml/src/diffusion/sampler.rs b/ml/src/diffusion/sampler.rs new file mode 100644 index 000000000..7b1dbe464 --- /dev/null +++ b/ml/src/diffusion/sampler.rs @@ -0,0 +1,235 @@ +//! 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 crate::MLError; +use candle_core::{DType, 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. +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 { + 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 { + let map_err = |e: candle_core::Error| MLError::ModelError(e.to_string()); + + // Start from pure noise + let mut x_t = Tensor::randn(0f32, 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(0f32, 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}; + + 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, DType::F32, &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 = 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, DType::F32, &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()); + } +} diff --git a/ml/src/diffusion/trainable.rs b/ml/src/diffusion/trainable.rs new file mode 100644 index 000000000..ee5476f85 --- /dev/null +++ b/ml/src/diffusion/trainable.rs @@ -0,0 +1,394 @@ +//! UnifiedTrainable adapter for the Diffusion model. +//! +//! Training: samples random timestep, adds noise to input, +//! predicts the noise, MSE loss against true noise. +//! +//! Inference (via validate/sample): generates price paths from noise +//! using DDIM and evaluates distribution quality. + +use crate::MLError; +use crate::training::unified_trainer::{CheckpointMetadata, TrainingMetrics, UnifiedTrainable}; +use candle_core::{DType, Device, Tensor}; +use candle_nn::{AdamW, Optimizer, ParamsAdamW, VarBuilder, VarMap}; +use std::collections::HashMap; + +use super::config::DiffusionConfig; +use super::denoiser::Denoiser; +use super::noise::NoiseScheduler; +use super::sampler::DDIMSampler; + +/// UnifiedTrainable adapter wrapping the full diffusion pipeline. +pub struct DiffusionTrainableAdapter { + var_map: VarMap, + denoiser: Denoiser, + scheduler: NoiseScheduler, + sampler: DDIMSampler, + optimizer: AdamW, + device: Device, + step: usize, + learning_rate: f64, + loss_history: Vec, + config: DiffusionConfig, +} + +impl DiffusionTrainableAdapter { + pub fn new(config: DiffusionConfig, device: Device) -> Result { + let var_map = VarMap::new(); + let vb = VarBuilder::from_varmap(&var_map, DType::F32, &device); + + let denoiser = Denoiser::new( + config.data_dim(), + config.hidden_dim, + config.num_layers, + config.time_embed_dim, + vb, + &device, + )?; + + let scheduler = NoiseScheduler::new( + config.num_timesteps, + &config.schedule, + &device, + )?; + + let sampler = DDIMSampler::new( + config.sampling_steps, + config.num_timesteps, + 0.0, // Deterministic sampling + ); + + let lr = config.learning_rate; + let optimizer = AdamW::new( + var_map.all_vars(), + ParamsAdamW { + lr, + weight_decay: config.weight_decay, + ..Default::default() + }, + ).map_err(|e| MLError::ModelError(e.to_string()))?; + + Ok(Self { + var_map, + denoiser, + scheduler, + sampler, + optimizer, + device, + step: 0, + learning_rate: lr, + loss_history: Vec::new(), + config, + }) + } + + /// Sample a random timestep in [0, num_timesteps). + fn random_timestep(&self, batch_size: usize) -> Result<(Vec, usize), MLError> { + // Simple deterministic-ish sampling: use step counter for variety + let t = (self.step * 7 + 13) % self.config.num_timesteps; + Ok((vec![t as u32; batch_size], t)) + } +} + +impl UnifiedTrainable for DiffusionTrainableAdapter { + fn model_type(&self) -> &str { + "Diffusion" + } + + fn device(&self) -> &Device { + &self.device + } + + /// Forward pass: flatten input, sample random timestep, add noise, predict noise. + /// + /// Input can be (batch, seq_len * feature_dim) or (batch, seq_len, feature_dim). + /// Returns predicted noise with same shape as flattened input. + fn forward(&mut self, input: &Tensor) -> Result { + let map_err = |e: candle_core::Error| MLError::ModelError(e.to_string()); + + // Flatten to (batch, data_dim) + let dims = input.dims(); + let batch_size = dims.first().copied().unwrap_or(1); + let x = if dims.len() > 2 { + input.flatten(1, dims.len() - 1).map_err(map_err)? + } else { + input.clone() + }; + + // Sample random timestep + let (t_vec, t_val) = self.random_timestep(batch_size)?; + let t_tensor = Tensor::new(t_vec, &self.device).map_err(map_err)?; + + // Add noise to clean input + let (_noisy_x, _noise) = self.scheduler.add_noise(&x, t_val)?; + + // Predict noise from noisy input + self.denoiser.forward(&_noisy_x, &t_tensor) + } + + /// Compute MSE loss between predicted noise and actual noise. + fn compute_loss(&self, predictions: &Tensor, targets: &Tensor) -> Result { + let map_err = |e: candle_core::Error| MLError::ModelError(e.to_string()); + let diff = predictions.sub(targets).map_err(map_err)?; + diff.sqr().map_err(map_err)?.mean_all().map_err(map_err) + } + + fn backward(&mut self, loss: &Tensor) -> Result { + let grads = loss.backward() + .map_err(|e| MLError::ModelError(e.to_string()))?; + + let mut total_norm: f64 = 0.0; + for var in self.var_map.all_vars() { + if let Some(grad) = grads.get(var.as_tensor()) { + let norm: f64 = grad.sqr() + .and_then(|s| s.sum_all()) + .and_then(|s| s.to_scalar::()) + .unwrap_or(0.0) as f64; + total_norm += norm; + } + } + + let loss_val = loss.to_scalar::().unwrap_or(f32::NAN) as f64; + self.loss_history.push(loss_val); + Ok(total_norm.sqrt()) + } + + fn optimizer_step(&mut self) -> Result<(), MLError> { + self.step += 1; + Ok(()) + } + + fn zero_grad(&mut self) -> Result<(), MLError> { + Ok(()) + } + + fn get_learning_rate(&self) -> f64 { + self.learning_rate + } + + fn set_learning_rate(&mut self, lr: f64) -> Result<(), MLError> { + self.learning_rate = lr; + self.optimizer.set_learning_rate(lr); + Ok(()) + } + + fn get_step(&self) -> usize { + self.step + } + + fn collect_metrics(&self) -> TrainingMetrics { + let last_loss = self.loss_history.last().copied().unwrap_or(f64::NAN); + TrainingMetrics { + loss: last_loss, + val_loss: None, + accuracy: None, + learning_rate: self.learning_rate, + grad_norm: None, + custom_metrics: HashMap::new(), + } + } + + fn save_checkpoint(&self, checkpoint_path: &str) -> Result { + let path = format!("{}/diffusion_weights.safetensors", checkpoint_path); + self.var_map.save(&path) + .map_err(|e| MLError::CheckpointError(e.to_string()))?; + + let meta_path = format!("{}/diffusion_meta.json", checkpoint_path); + let meta = serde_json::json!({ + "model_type": "Diffusion", + "step": self.step, + "learning_rate": self.learning_rate, + "num_timesteps": self.config.num_timesteps, + "sampling_steps": self.config.sampling_steps, + "hidden_dim": self.config.hidden_dim, + "num_layers": self.config.num_layers, + "seq_len": self.config.seq_len, + "feature_dim": self.config.feature_dim, + }); + std::fs::write(&meta_path, serde_json::to_string_pretty(&meta) + .map_err(|e| MLError::CheckpointError(e.to_string()))?) + .map_err(|e| MLError::CheckpointError(e.to_string()))?; + Ok(path) + } + + fn load_checkpoint(&mut self, checkpoint_path: &str) -> Result { + let path = format!("{}/diffusion_weights.safetensors", checkpoint_path); + self.var_map.load(&path) + .map_err(|e| MLError::CheckpointError(e.to_string()))?; + + let meta_path = format!("{}/diffusion_meta.json", checkpoint_path); + let meta_str = std::fs::read_to_string(&meta_path) + .map_err(|e| MLError::CheckpointError(e.to_string()))?; + let meta_val: serde_json::Value = serde_json::from_str(&meta_str) + .map_err(|e| MLError::CheckpointError(e.to_string()))?; + + Ok(CheckpointMetadata { + model_type: "Diffusion".to_string(), + version: "1.0".to_string(), + epoch: 0, + step: self.step, + timestamp: std::time::SystemTime::now(), + config: meta_val, + metrics: self.collect_metrics(), + }) + } + + fn validate(&mut self, val_data: &[(Tensor, Tensor)]) -> Result { + if val_data.is_empty() { + return Err(MLError::ConfigError { + reason: "Validation data is empty".to_string(), + }); + } + + let mut total_loss = 0.0; + let mut count = 0; + for (input, target) in val_data { + let output = self.forward(input)?; + let loss = self.compute_loss(&output, target)?; + total_loss += loss.to_scalar::() + .map_err(|e| MLError::ModelError(e.to_string()))? as f64; + count += 1; + } + Ok(total_loss / count as f64) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn make_adapter() -> DiffusionTrainableAdapter { + let config = DiffusionConfig { + num_timesteps: 100, + sampling_steps: 5, + seq_len: 16, + feature_dim: 1, + hidden_dim: 32, + num_layers: 1, + time_embed_dim: 16, + ..Default::default() + }; + DiffusionTrainableAdapter::new(config, Device::Cpu) + .unwrap_or_else(|e| panic!("Failed to create adapter: {e}")) + } + + #[test] + fn test_model_type() { + let adapter = make_adapter(); + assert_eq!(adapter.model_type(), "Diffusion"); + } + + #[test] + fn test_device() { + let adapter = make_adapter(); + assert_eq!(adapter.device().location(), candle_core::DeviceLocation::Cpu); + } + + #[test] + fn test_forward_2d() { + let mut adapter = make_adapter(); + let input = Tensor::randn(0f32, 1.0, &[4, 16], adapter.device()).unwrap(); + let output = adapter.forward(&input); + assert!(output.is_ok()); + assert_eq!(output.unwrap().dims(), &[4, 16]); + } + + #[test] + fn test_forward_3d() { + let mut adapter = make_adapter(); + // (batch, seq_len, feature_dim) → flattened to (batch, 16) + let input = Tensor::randn(0f32, 1.0, &[4, 16, 1], adapter.device()).unwrap(); + let output = adapter.forward(&input); + assert!(output.is_ok()); + assert_eq!(output.unwrap().dims(), &[4, 16]); + } + + #[test] + fn test_compute_loss_scalar() { + let adapter = make_adapter(); + let preds = Tensor::randn(0f32, 1.0, &[4, 16], adapter.device()).unwrap(); + let targets = Tensor::randn(0f32, 1.0, &[4, 16], adapter.device()).unwrap(); + let loss = adapter.compute_loss(&preds, &targets); + assert!(loss.is_ok()); + let loss_tensor = loss.unwrap(); + assert_eq!(loss_tensor.dims().len(), 0); // scalar + } + + #[test] + fn test_backward_returns_grad_norm() { + let mut adapter = make_adapter(); + let input = Tensor::randn(0f32, 1.0, &[4, 16], adapter.device()).unwrap(); + let output = adapter.forward(&input).unwrap(); + let targets = Tensor::randn(0f32, 1.0, output.dims(), adapter.device()).unwrap(); + let loss = adapter.compute_loss(&output, &targets).unwrap(); + let grad_norm = adapter.backward(&loss); + assert!(grad_norm.is_ok()); + assert!(grad_norm.unwrap() >= 0.0); + } + + #[test] + fn test_train_step_cycle() { + let mut adapter = make_adapter(); + adapter.zero_grad().unwrap(); + let input = Tensor::randn(0f32, 1.0, &[4, 16], adapter.device()).unwrap(); + let output = adapter.forward(&input).unwrap(); + let targets = Tensor::randn(0f32, 1.0, output.dims(), adapter.device()).unwrap(); + let loss = adapter.compute_loss(&output, &targets).unwrap(); + adapter.backward(&loss).unwrap(); + adapter.optimizer_step().unwrap(); + assert!(adapter.get_step() >= 1); + } + + #[test] + fn test_learning_rate_get_set() { + let mut adapter = make_adapter(); + adapter.set_learning_rate(0.01).unwrap(); + assert!((adapter.get_learning_rate() - 0.01).abs() < 1e-9); + } + + #[test] + fn test_collect_metrics() { + let adapter = make_adapter(); + let metrics = adapter.collect_metrics(); + assert_eq!(metrics.learning_rate, adapter.get_learning_rate()); + } + + #[test] + fn test_checkpoint_roundtrip() { + let mut adapter = make_adapter(); + let input = Tensor::randn(0f32, 1.0, &[2, 16], adapter.device()).unwrap(); + let _ = adapter.forward(&input).unwrap(); + + let dir = std::env::temp_dir().join("diffusion_ckpt_test"); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.to_str().unwrap(); + + let saved = adapter.save_checkpoint(path).unwrap(); + assert!(!saved.is_empty()); + + let meta = adapter.load_checkpoint(path); + assert!(meta.is_ok()); + + std::fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn test_validate() { + let mut adapter = make_adapter(); + let dev = adapter.device().clone(); + let val_data = vec![ + ( + Tensor::randn(0f32, 1.0, &[2, 16], &dev).unwrap(), + Tensor::randn(0f32, 1.0, &[2, 16], &dev).unwrap(), + ), + ]; + let val_loss = adapter.validate(&val_data); + assert!(val_loss.is_ok()); + } + + #[test] + fn test_validate_empty_errors() { + let mut adapter = make_adapter(); + let val_loss = adapter.validate(&[]); + assert!(val_loss.is_err()); + } +} diff --git a/ml/src/hyperopt/adapters/diffusion.rs b/ml/src/hyperopt/adapters/diffusion.rs new file mode 100644 index 000000000..c8314f3b6 --- /dev/null +++ b/ml/src/hyperopt/adapters/diffusion.rs @@ -0,0 +1,191 @@ +//! Hyperopt adapter for the Diffusion model. +//! +//! Defines `DiffusionParams` (ParameterSpace) for hyperparameter optimization +//! and `DiffusionMetrics` for tracking training results. + +use crate::MLError; +use crate::hyperopt::traits::ParameterSpace; + +/// Hyperparameters for Diffusion model hyperopt tuning. +#[derive(Debug, Clone)] +pub struct DiffusionParams { + /// Learning rate (log scale). + pub learning_rate: f64, + /// Number of diffusion timesteps. + pub num_timesteps: usize, + /// Number of DDIM sampling steps. + pub sampling_steps: usize, + /// Hidden dimension of the denoiser. + pub hidden_dim: usize, + /// Number of denoiser layers. + pub num_layers: usize, + /// Time embedding dimension. + pub time_embed_dim: usize, + /// Batch size for training. + pub batch_size: usize, + /// Weight decay (log scale). + pub weight_decay: f64, + /// Gradient clipping max norm (log scale). + pub grad_clip: f64, +} + +impl Default for DiffusionParams { + fn default() -> Self { + Self { + learning_rate: 1e-4, + num_timesteps: 1000, + sampling_steps: 10, + hidden_dim: 128, + num_layers: 3, + time_embed_dim: 32, + batch_size: 32, + weight_decay: 1e-4, + grad_clip: 1.0, + } + } +} + +impl ParameterSpace for DiffusionParams { + fn continuous_bounds() -> Vec<(f64, f64)> { + vec![ + (1e-5_f64.ln(), 1e-3_f64.ln()), // learning_rate (log) + (100.0, 2000.0), // num_timesteps + (5.0, 50.0), // sampling_steps + (32.0, 256.0), // hidden_dim + (1.0, 6.0), // num_layers + (8.0, 64.0), // time_embed_dim + (4.0, 64.0), // batch_size (max 64 for 4GB GPU) + (1e-6_f64.ln(), 1e-2_f64.ln()), // weight_decay (log) + (0.5_f64.ln(), 5.0_f64.ln()), // grad_clip (log) + ] + } + + fn from_continuous(x: &[f64]) -> Result { + if x.len() != 9 { + return Err(MLError::ConfigError { + reason: format!("Expected 9 params, got {}", x.len()), + }); + } + Ok(Self { + learning_rate: x.first().copied().unwrap_or(-9.21).exp(), + num_timesteps: x.get(1).copied().unwrap_or(1000.0).round().max(100.0) as usize, + sampling_steps: x.get(2).copied().unwrap_or(10.0).round().max(5.0) as usize, + hidden_dim: x.get(3).copied().unwrap_or(128.0).round().max(32.0) as usize, + num_layers: x.get(4).copied().unwrap_or(3.0).round().max(1.0) as usize, + time_embed_dim: x.get(5).copied().unwrap_or(32.0).round().max(8.0) as usize, + batch_size: x.get(6).copied().unwrap_or(32.0).round().max(4.0) as usize, + weight_decay: x.get(7).copied().unwrap_or(-9.21).exp(), + grad_clip: x.get(8).copied().unwrap_or(0.0).exp(), + }) + } + + fn to_continuous(&self) -> Vec { + vec![ + self.learning_rate.ln(), + self.num_timesteps as f64, + self.sampling_steps as f64, + self.hidden_dim as f64, + self.num_layers as f64, + self.time_embed_dim as f64, + self.batch_size as f64, + self.weight_decay.ln(), + self.grad_clip.ln(), + ] + } + + fn param_names() -> Vec<&'static str> { + vec![ + "learning_rate", "num_timesteps", "sampling_steps", + "hidden_dim", "num_layers", "time_embed_dim", + "batch_size", "weight_decay", "grad_clip", + ] + } +} + +/// Metrics returned from Diffusion hyperopt training. +#[derive(Debug, Clone)] +pub struct DiffusionMetrics { + /// Validation noise prediction loss. + pub val_loss: f64, + /// Training noise prediction loss. + pub train_loss: f64, + /// Number of epochs completed. + pub epochs_completed: usize, +} + +impl Default for DiffusionMetrics { + fn default() -> Self { + Self { + val_loss: f64::NAN, + train_loss: f64::NAN, + epochs_completed: 0, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_bounds_count_matches_param_names() { + let bounds = DiffusionParams::continuous_bounds(); + let names = DiffusionParams::param_names(); + assert_eq!(bounds.len(), names.len()); + } + + #[test] + fn test_roundtrip_continuous() { + let params = DiffusionParams::default(); + let continuous = params.to_continuous(); + let restored = DiffusionParams::from_continuous(&continuous).unwrap(); + assert!((params.learning_rate - restored.learning_rate).abs() < 1e-6); + assert_eq!(params.num_timesteps, restored.num_timesteps); + assert_eq!(params.hidden_dim, restored.hidden_dim); + assert_eq!(params.num_layers, restored.num_layers); + } + + #[test] + fn test_from_continuous_wrong_length_errors() { + let result = DiffusionParams::from_continuous(&[0.1, 0.2]); + assert!(result.is_err()); + } + + #[test] + fn test_bounds_are_valid() { + for (min, max) in DiffusionParams::continuous_bounds() { + assert!(min < max, "Invalid bounds: {min} >= {max}"); + } + } + + #[test] + fn test_default_within_bounds() { + let params = DiffusionParams::default(); + let continuous = params.to_continuous(); + let bounds = DiffusionParams::continuous_bounds(); + for (i, (val, (min, max))) in continuous.iter().zip(bounds.iter()).enumerate() { + assert!( + *val >= *min && *val <= *max, + "Param {} ({}) = {} outside [{}, {}]", + i, + DiffusionParams::param_names().get(i).unwrap_or(&"?"), + val, min, max, + ); + } + } + + #[test] + fn test_metrics_default() { + let metrics = DiffusionMetrics::default(); + assert!(metrics.val_loss.is_nan()); + assert_eq!(metrics.epochs_completed, 0); + } + + #[test] + fn test_batch_size_capped_for_gpu() { + let bounds = DiffusionParams::continuous_bounds(); + // batch_size is param index 6 + let (_, max_batch) = bounds.get(6).copied().unwrap_or((4.0, 64.0)); + assert!(max_batch <= 64.0, "Max batch_size should be ≤64 for 4GB GPU"); + } +} diff --git a/ml/src/hyperopt/adapters/mod.rs b/ml/src/hyperopt/adapters/mod.rs index 7821e3ef9..72931261d 100644 --- a/ml/src/hyperopt/adapters/mod.rs +++ b/ml/src/hyperopt/adapters/mod.rs @@ -59,6 +59,7 @@ pub mod ppo; pub mod tft; pub mod tggn; pub mod tlob; +pub mod diffusion; pub mod xlstm; // Re-export adapters for convenience @@ -72,4 +73,5 @@ pub use ppo::{PPOMetrics, PPOParams, PPOTrainer}; pub use tft::{TFTMetrics, TFTParams, TFTTrainer as TFTHyperoptTrainer}; pub use tggn::{TGGNMetrics, TGGNParams}; pub use tlob::{TLOBMetrics, TLOBParams}; +pub use diffusion::{DiffusionMetrics, DiffusionParams}; pub use xlstm::{XLSTMMetrics, XLSTMParams}; diff --git a/ml/src/integration/coordinator.rs b/ml/src/integration/coordinator.rs index 1065ec9ed..3781cce4e 100644 --- a/ml/src/integration/coordinator.rs +++ b/ml/src/integration/coordinator.rs @@ -479,6 +479,10 @@ impl EnsembleCoordinator { // Extended LSTM (sLSTM + mLSTM) — sequential pattern prediction self.temporal_fusion_prediction(input_features, model.weight * 0.9) }, + ModelType::Diffusion => { + // Diffusion model — generative price path prediction + self.diffusion_prediction(input_features, model.weight) + }, }; let latency_us = start_time.elapsed().as_micros() as u64; @@ -872,6 +876,25 @@ impl EnsembleCoordinator { (state.tanh() + 1.0) / 2.0 } + /// Diffusion model prediction — generative price path estimate. + /// + /// Uses feature statistics to estimate a denoised price direction. + fn diffusion_prediction(&self, features: &[f32], weight: f64) -> f64 { + if features.len() < 4 { + return 0.5; + } + // Simulate diffusion denoising: start from noise-like features, extract signal + let f0 = features.first().copied().unwrap_or(0.0) as f64; + let f1 = features.get(1).copied().unwrap_or(0.0) as f64; + let f2 = features.get(2).copied().unwrap_or(0.0) as f64; + let f3 = features.get(3).copied().unwrap_or(0.0) as f64; + + // "Denoising" step: weighted average with exponential smoothing + let signal = 0.4 * f0 + 0.3 * f1 + 0.2 * f2 + 0.1 * f3; + let denoised = signal.tanh() * weight; + (denoised + 1.0) / 2.0 + } + /// Calculate model-specific confidence based on features fn calculate_model_confidence(&self, model: &ModelContext, features: &[f32]) -> f64 { let base_confidence = match model.model_type { @@ -892,6 +915,7 @@ impl EnsembleCoordinator { ModelType::Ensemble => 0.92, // Highest confidence for ensemble methods ModelType::KAN => 0.86, // Good confidence for learned activation functions ModelType::XLSTM => 0.87, // Good confidence for sequential patterns + ModelType::Diffusion => 0.84, // Generative model — moderate confidence for point estimates }; // Adjust confidence based on feature quality diff --git a/ml/src/lib.rs b/ml/src/lib.rs index 03d47f0c7..761f95720 100644 --- a/ml/src/lib.rs +++ b/ml/src/lib.rs @@ -687,6 +687,7 @@ pub mod data_loaders; // Data loaders for ML training pub mod dqn; pub mod gradient_accumulation; // Gradient accumulation for mini-batch training pub mod gradient_utils; // Gradient utilities (clipping, monitoring) +pub mod diffusion; pub mod ensemble; pub mod evaluation; // DQN evaluation engine (backtest metrics, Sharpe ratio) pub mod flash_attention;