gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(ml): AttentionMask causal masking, document multi-asset DQN and conversion layer

Browse Source 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-23 01:03:53 +01:00
parent cfadd7d3dc
commit d4a0255595
5 changed files with 81 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
ml/src/observability/metrics.rs
View File

@@ -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();

11
ml/src/trainers/dqn/trainer.rs
View File

@@ -445,8 +445,15 @@ impl DQNTrainer {
None
};
// Multi-asset portfolio tracking (disabled by default - single-asset mode)
let multi_asset_portfolio: Option<Arc<crate::dqn::multi_asset::MultiAssetPortfolioTracker>> = 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<Arc<crate::dqn::multi_asset::MultiAssetPortfolioTracker>> = None;
// Stress testing for robustness validation
let stress_tester: Option<Arc<crate::dqn::stress_testing::DQNStressTester>> = if hyperparams.enable_stress_testing {

63
ml/src/transformers/attention.rs
View File

@@ -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<Self, candle_core::Error> {
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::<f32>()?;
// 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(())
}
}

8
ml/src/transformers/mod.rs
View File

@@ -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.

10
services/ml_training_service/src/orchestrator.rs
View File

@@ -860,8 +860,14 @@ impl TrainingOrchestrator {
feature_vectors_225.len()
);
// Convert [f64; 225] to (FinancialFeatures, Vec<f64>) format for compatibility
// TODO: Once all models support [f64; 225], remove this conversion layer
// Convert [f64; 225] to (FinancialFeatures, Vec<f64>) format for compatibility.
//
// This conversion exists because UnifiedTrainable::train() accepts
// Vec<(FinancialFeatures, Vec<f64>)>. The FinancialFeatures fields below are
// stubs — only the Vec<f64> component is consumed by the underlying models.
//
// To remove this layer: change UnifiedTrainable::train() to accept &[Vec<f64>]
// directly and update all 4 model adapters (DQN, PPO, TFT, Mamba2).
let training_samples: Vec<(FinancialFeatures, Vec<f64>)> = feature_vectors_225
.iter()
.zip(ohlcv_bars.iter().skip(WARMUP_PERIOD))