From 37d617fb6a6b7c7422c36da223e413f1a42e3040 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 21 Feb 2026 13:17:31 +0100 Subject: [PATCH] feat(mamba): add directional MSE loss and fix ZOH->bilinear discretization Co-Authored-By: Claude Opus 4.6 --- ml/src/mamba/loss.rs | 112 +++++++++++++++++++++++++++++++++++++++++++ ml/src/mamba/mod.rs | 24 ++++++++-- 2 files changed, 132 insertions(+), 4 deletions(-) create mode 100644 ml/src/mamba/loss.rs diff --git a/ml/src/mamba/loss.rs b/ml/src/mamba/loss.rs new file mode 100644 index 000000000..5ef4a4b6e --- /dev/null +++ b/ml/src/mamba/loss.rs @@ -0,0 +1,112 @@ +//! Loss functions for Mamba2 SSM training +//! +//! Includes directional MSE that penalizes wrong-sign predictions +//! more heavily — critical for trading where direction matters more +//! than magnitude. + +use candle_core::Tensor; +use crate::{MLError, MLResult}; + +/// Standard MSE loss (baseline for comparison) +pub fn mse_loss(predictions: &Tensor, targets: &Tensor) -> MLResult { + let diff = predictions.sub(targets) + .map_err(|e| MLError::ModelError(format!("mse sub: {}", e)))?; + let sq = diff.sqr() + .map_err(|e| MLError::ModelError(format!("mse sqr: {}", e)))?; + sq.mean_all() + .map_err(|e| MLError::ModelError(format!("mse mean: {}", e))) +} + +/// Directional MSE loss. +/// +/// When prediction and target have opposite signs (wrong direction), +/// the squared error is multiplied by `direction_weight` (default: 2.0). +/// When signs agree, normal MSE is used. +/// +/// `direction_weight = 1.0` produces identical results to standard MSE. +/// `direction_weight = 2.0` penalizes wrong-direction errors 2x. +pub fn directional_mse_loss( + predictions: &Tensor, + targets: &Tensor, + direction_weight: f64, +) -> MLResult { + let diff = predictions.sub(targets) + .map_err(|e| MLError::ModelError(format!("dir_mse sub: {}", e)))?; + let sq = diff.sqr() + .map_err(|e| MLError::ModelError(format!("dir_mse sqr: {}", e)))?; + + // Detect wrong direction: sign(pred) != sign(target) + let zeros = Tensor::zeros_like(predictions) + .map_err(|e| MLError::ModelError(format!("dir_mse zeros: {}", e)))?; + let pred_positive = predictions.ge(&zeros) + .map_err(|e| MLError::ModelError(format!("dir_mse pred_ge: {}", e)))?; + let target_positive = targets.ge(&zeros) + .map_err(|e| MLError::ModelError(format!("dir_mse target_ge: {}", e)))?; + + // wrong_direction = (pred_positive XOR target_positive) as float + let wrong = pred_positive.ne(&target_positive) + .map_err(|e| MLError::ModelError(format!("dir_mse ne: {}", e)))? + .to_dtype(predictions.dtype()) + .map_err(|e| MLError::ModelError(format!("dir_mse to_dtype: {}", e)))?; + + // weight = 1.0 + wrong * (direction_weight - 1.0) + let extra = wrong.affine(direction_weight - 1.0, 1.0) + .map_err(|e| MLError::ModelError(format!("dir_mse affine: {}", e)))?; + + let weighted = sq.mul(&extra) + .map_err(|e| MLError::ModelError(format!("dir_mse mul: {}", e)))?; + + weighted.mean_all() + .map_err(|e| MLError::ModelError(format!("dir_mse mean: {}", e))) +} + +#[cfg(test)] +mod tests { + use super::*; + use candle_core::Device; + + #[test] + fn test_directional_loss_same_as_mse_when_correct_direction() { + let device = Device::Cpu; + let preds = Tensor::from_vec(vec![0.3_f64, 0.5, 0.8], 3, &device).unwrap(); + let targets = Tensor::from_vec(vec![0.2_f64, 0.6, 0.9], 3, &device).unwrap(); + + let dir_loss = directional_mse_loss(&preds, &targets, 2.0).unwrap(); + let mse = mse_loss(&preds, &targets).unwrap(); + + let dir_val: f64 = dir_loss.to_vec0().unwrap(); + let mse_val: f64 = mse.to_vec0().unwrap(); + assert!((dir_val - mse_val).abs() < 1e-6, + "same-direction: dir={} vs mse={}", dir_val, mse_val); + } + + #[test] + fn test_directional_loss_higher_for_wrong_direction() { + let device = Device::Cpu; + let preds = Tensor::from_vec(vec![0.3_f64], 1, &device).unwrap(); + let targets = Tensor::from_vec(vec![-0.2_f64], 1, &device).unwrap(); + + let dir_loss = directional_mse_loss(&preds, &targets, 2.0).unwrap(); + let mse = mse_loss(&preds, &targets).unwrap(); + + let dir_val: f64 = dir_loss.to_vec0().unwrap(); + let mse_val: f64 = mse.to_vec0().unwrap(); + assert!(dir_val > mse_val, + "wrong-direction loss {} should be > mse {}", dir_val, mse_val); + } + + #[test] + fn test_directional_loss_weight_1_equals_mse() { + let device = Device::Cpu; + let preds = Tensor::from_vec(vec![0.3_f64, -0.1], 2, &device).unwrap(); + let targets = Tensor::from_vec(vec![-0.2_f64, 0.5], 2, &device).unwrap(); + + let dir_loss = directional_mse_loss(&preds, &targets, 1.0).unwrap(); + let mse = mse_loss(&preds, &targets).unwrap(); + + let dir_val: f64 = dir_loss.to_vec0().unwrap(); + let mse_val: f64 = mse.to_vec0().unwrap(); + assert!((dir_val - mse_val).abs() < 1e-6, + "weight=1.0: dir={} should == mse={}", dir_val, mse_val); + } +} diff --git a/ml/src/mamba/mod.rs b/ml/src/mamba/mod.rs index 51c10baba..6d5d9cb3e 100644 --- a/ml/src/mamba/mod.rs +++ b/ml/src/mamba/mod.rs @@ -42,6 +42,7 @@ mod hardware_aware; mod scan_algorithms; pub mod selective_state; mod ssd_layer; +pub mod loss; pub mod trainable_adapter; // Public exports for types used in mod.rs and by external crates @@ -920,11 +921,15 @@ impl Mamba2SSM { // Create a 0-D scalar tensor with F64 dtype (matching mean_all output) let dt_tensor = Tensor::from_slice(&[dt_scalar], &[1], A_cont.device())?.reshape(&[])?; // Make it 0-D scalar - // A_discrete = exp(A_cont * dt) - // For simplicity, using first-order approximation: I + A_cont * dt - let A_scaled = A_cont.broadcast_mul(&dt_tensor)?; + // Bilinear (Tustin) approximation: A_disc = I + A*dt + (A*dt)^2 / 2 + // More accurate than ZOH (I + A*dt), matches ssd_layer.rs + let A_dt = A_cont.broadcast_mul(&dt_tensor)?; + let A_dt_sq = A_dt.matmul(&A_dt)?; + let half = Tensor::from_slice(&[0.5_f64], &[1], A_cont.device())?.reshape(&[])?; + let second_order = A_dt_sq.broadcast_mul(&half)?; + let identity = Tensor::eye(A_cont.dim(0)?, DType::F64, A_cont.device())?; - let A_discrete = (&identity + &A_scaled)?; + let A_discrete = ((&identity + &A_dt)? + &second_order)?; Ok(A_discrete) } @@ -3240,3 +3245,14 @@ fn test_mamba_parameter_count() -> anyhow::Result<()> { assert!(param_count > 0); Ok(()) } + +#[test] +fn test_bilinear_discretization_more_accurate_than_zoh() { + let zoh = 1.0 + (-1.0) * 0.1; + let bilinear = 1.0 + (-1.0) * 0.1 + ((-1.0) * 0.1_f64).powi(2) / 2.0; + let exact = (-0.1_f64).exp(); + + assert!((bilinear - exact).abs() < (zoh - exact).abs(), + "bilinear {} should be closer to exact {} than zoh {}", + bilinear, exact, zoh); +}