feat(ml): add Diffusion model (DDPM/DDIM) for price path generation

- NoiseScheduler: precomputed cosine/linear alpha_bar schedules
- Denoiser: FC network with sinusoidal time embedding + SiLU + residual
- DDIMSampler: deterministic fast sampling (10 steps from 1000 timesteps)
- DiffusionTrainableAdapter: UnifiedTrainable for unified training pipeline
- Hyperopt adapter with ParameterSpace (9 params, batch ≤64 for 4GB GPU)
- ModelType::Diffusion registered in common + coordinator
- 41 tests passing (config=3, noise=7, denoiser=4, sampler=5, trainable=12, hyperopt=7)
- OOM-safe: FC denoiser instead of U-Net, small hidden dims, conservative defaults

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-23 09:32:44 +01:00
parent 83054548b8
commit b88fd62af2
11 changed files with 1459 additions and 0 deletions

View File

@@ -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<Self, MLError> {
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<f64> {
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");
}
}

View File

@@ -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};