feat(mamba): add directional MSE loss and fix ZOH->bilinear discretization

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

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<Tensor> {
+    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<Tensor> {
+    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);
+    }
+}

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);
+}