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:
191
ml/src/hyperopt/adapters/diffusion.rs
Normal file
191
ml/src/hyperopt/adapters/diffusion.rs
Normal 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");
|
||||
}
|
||||
}
|
||||
@@ -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};
|
||||
|
||||
Reference in New Issue
Block a user