From d4a02555954a7a2ac88b8e9dd077f5d5a91fd80e Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 23 Feb 2026 01:03:53 +0100 Subject: [PATCH] feat(ml): AttentionMask causal masking, document multi-asset DQN and conversion layer Co-Authored-By: Claude Opus 4.6 --- ml/src/observability/metrics.rs | 7 ++- ml/src/trainers/dqn/trainer.rs | 11 +++- ml/src/transformers/attention.rs | 63 ++++++++++++++++--- ml/src/transformers/mod.rs | 8 +-- .../ml_training_service/src/orchestrator.rs | 10 ++- 5 files changed, 81 insertions(+), 18 deletions(-) diff --git a/ml/src/observability/metrics.rs b/ml/src/observability/metrics.rs index fd198d973..77405658e 100644 --- a/ml/src/observability/metrics.rs +++ b/ml/src/observability/metrics.rs @@ -17,8 +17,11 @@ use tokio::sync::RwLock; use crate::{MLError, MLResult, ModelPrediction, ModelType}; -// Helper function for asset class bucketing (duplicated here to avoid circular dependency) -// TODO: Move to common crate utility module +/// Helper function for asset class bucketing. +/// +/// Duplicated from common crate to avoid a circular dependency (ml -> common -> ml). +/// Tracked for dedup: when `common` exposes a public `bucket_symbol()`, replace this +/// local copy with a re-import. fn bucket_symbol(symbol: &str) -> &'static str { let upper = symbol.to_uppercase(); let upper_str = upper.as_str(); diff --git a/ml/src/trainers/dqn/trainer.rs b/ml/src/trainers/dqn/trainer.rs index 56b32fcc4..f0f1370b4 100644 --- a/ml/src/trainers/dqn/trainer.rs +++ b/ml/src/trainers/dqn/trainer.rs @@ -445,8 +445,15 @@ impl DQNTrainer { None }; - // Multi-asset portfolio tracking (disabled by default - single-asset mode) - let multi_asset_portfolio: Option> = None; // TODO: Add enable_multi_asset flag when needed + // Multi-asset portfolio tracking (disabled — single-asset is the current production mode) + // + // When expanding to multi-asset trading: + // 1. Add `enable_multi_asset: bool` to DQNHyperparams (default false). + // 2. Initialize MultiAssetPortfolioTracker here when the flag is set. + // 3. Wire portfolio state into the DQN observation: expand state_dim to + // include per-asset position, PnL, and correlation features so the + // agent can learn cross-asset hedging and allocation. + let multi_asset_portfolio: Option> = None; // Stress testing for robustness validation let stress_tester: Option> = if hyperparams.enable_stress_testing { diff --git a/ml/src/transformers/attention.rs b/ml/src/transformers/attention.rs index 768766b7a..2b839cd7c 100644 --- a/ml/src/transformers/attention.rs +++ b/ml/src/transformers/attention.rs @@ -2,11 +2,45 @@ //! //! This module provides basic attention mechanisms using modern Candle API patterns. +use candle_core::{Device, Tensor}; + +/// A causal attention mask for transformer self-attention. +/// +/// The mask is a 2-D `[seq_len, seq_len]` tensor where positions above the +/// diagonal are `f32::NEG_INFINITY` (masked out) and positions on or below +/// the diagonal are `0.0` (allowed). Adding this to raw attention logits +/// prevents the model from attending to future positions. +#[derive(Debug)] +pub struct AttentionMask { + pub mask: Tensor, +} + +impl AttentionMask { + /// Creates a lower-triangular causal mask of shape `[seq_len, seq_len]`. + /// + /// Positions `(i, j)` where `j <= i` are `0.0`; positions where `j > i` + /// are `f32::NEG_INFINITY`. + pub fn causal(seq_len: usize, device: &Device) -> Result { + let mut data = Vec::with_capacity(seq_len * seq_len); + for i in 0..seq_len { + for j in 0..seq_len { + if j <= i { + data.push(0.0f32); + } else { + data.push(f32::NEG_INFINITY); + } + } + } + let mask = Tensor::from_vec(data, (seq_len, seq_len), device)?; + Ok(Self { mask }) + } +} + #[cfg(test)] mod tests { + use super::*; use crate::tft::temporal_attention::AttentionConfig; - #[test] fn test_attention_config() { let config = AttentionConfig::default(); @@ -15,11 +49,24 @@ mod tests { assert_eq!(config.dropout_rate, 0.1); } - // TODO: Re-enable when AttentionMask is implemented - // #[test] - // fn test_attention_mask() { - // let device = Device::Cpu; - // let mask = AttentionMask::causal(4, &device)?; - // assert_eq!(mask.mask.dims(), &[4, 4]); - // } + #[test] + fn test_attention_mask() -> Result<(), candle_core::Error> { + let device = Device::Cpu; + let mask = AttentionMask::causal(4, &device)?; + assert_eq!(mask.mask.dims(), &[4, 4]); + + // Verify the diagonal and above-diagonal values + let flat = mask.mask.flatten_all()?.to_vec1::()?; + // Row 0: [0, -inf, -inf, -inf] + assert_eq!(*flat.get(0).unwrap_or(&f32::NAN), 0.0); + assert!(flat.get(1).unwrap_or(&0.0).is_infinite()); + // Row 1: [0, 0, -inf, -inf] + assert_eq!(*flat.get(4).unwrap_or(&f32::NAN), 0.0); + assert_eq!(*flat.get(5).unwrap_or(&f32::NAN), 0.0); + assert!(flat.get(6).unwrap_or(&0.0).is_infinite()); + // Row 3 (last): [0, 0, 0, 0] + assert_eq!(*flat.get(12).unwrap_or(&f32::NAN), 0.0); + assert_eq!(*flat.get(15).unwrap_or(&f32::NAN), 0.0); + Ok(()) + } } diff --git a/ml/src/transformers/mod.rs b/ml/src/transformers/mod.rs index 3879b1d64..a669f9ff6 100644 --- a/ml/src/transformers/mod.rs +++ b/ml/src/transformers/mod.rs @@ -30,10 +30,10 @@ // Core modules that compile successfully pub mod attention; -// Re-export core types that work (commented out until implemented) -// pub use attention::{ -// AttentionConfig, AttentionMask, CrossModalAttention, MultiHeadAttention, -// }; +// Re-export core types that are implemented +pub use attention::AttentionMask; +// TODO: Re-export remaining types when implemented: +// pub use attention::{AttentionConfig, CrossModalAttention, MultiHeadAttention}; /// Transformer model types optimized for different `HFT` use cases /// TransformerType component. diff --git a/services/ml_training_service/src/orchestrator.rs b/services/ml_training_service/src/orchestrator.rs index b70c5cb72..4bae1c070 100644 --- a/services/ml_training_service/src/orchestrator.rs +++ b/services/ml_training_service/src/orchestrator.rs @@ -860,8 +860,14 @@ impl TrainingOrchestrator { feature_vectors_225.len() ); - // Convert [f64; 225] to (FinancialFeatures, Vec) format for compatibility - // TODO: Once all models support [f64; 225], remove this conversion layer + // Convert [f64; 225] to (FinancialFeatures, Vec) format for compatibility. + // + // This conversion exists because UnifiedTrainable::train() accepts + // Vec<(FinancialFeatures, Vec)>. The FinancialFeatures fields below are + // stubs — only the Vec component is consumed by the underlying models. + // + // To remove this layer: change UnifiedTrainable::train() to accept &[Vec] + // directly and update all 4 model adapters (DQN, PPO, TFT, Mamba2). let training_samples: Vec<(FinancialFeatures, Vec)> = feature_vectors_225 .iter() .zip(ohlcv_bars.iter().skip(WARMUP_PERIOD))