foxhunt/AGENT_253_AGENT_246_REVIEW.md

Agent 253: Agent 246's output_dim=1 Change Review

Mission: Determine if Agent 246 made the correct architectural decision when changing MAMBA-2's output_dim from d_model to 1

Verdict: ❌ INCORRECT - Agent 246's change contradicts the data loader's architecture

Priority: 🔴 CRITICAL - This is a DATA LOADER BUG, not a model bug

Date: 2025-10-15

---

Executive Summary

Agent 246 changed MAMBA-2's output projection from d_inner → d_model to d_inner → 1, claiming the model performs "price regression" (single output). However, this is architecturally incorrect based on the actual data loader implementation.

The Real Problem: The DbnSequenceLoader is creating targets with shape [1, 1, d_model=256] (full feature vectors), but Agent 246 configured the model to output [batch, seq, 1] (single values). This is a shape mismatch at the data pipeline level, not a model design issue.

---

Evidence Analysis

1. Agent 246's Reasoning

From /home/jgrusewski/Work/foxhunt/AGENT_246_FIXES_APPLIED.md:

Root Cause: The MAMBA-2 model was outputting [batch, seq, d_model] when tests expected [batch, seq, 1] for regression tasks (price prediction).

Agent 210's Misunderstanding: Previous Agent 210 "fixed" the output dimension from 1 to d_model, believing MAMBA-2 was a sequence-to-sequence model. This was incorrect - Foxhunt uses MAMBA-2 for price regression, not sequence modeling.

Agent 246's Logic:

• Tests assert output.dims()[2] == 1 (line 292 of e2e_mamba2_training.rs)
• Conclusion: Model should output 1 feature (price prediction)
• Fix: Change output projection to d_inner → 1

2. What the Data Loader Actually Does

From /home/jgrusewski/Work/foxhunt/ml/src/data_loaders/dbn_sequence_loader.rs (lines 590-615):

// Target is next timestep (autoregressive)
let target_msg = &window[self.seq_len];
let target_features = self.extract_features(target_msg)?;

debug_assert_eq!(
    target_features.len(),
    self.d_model,  // ← EXPECTS d_model (256) features!
    "Target feature dimension mismatch: expected {}, got {}",
    self.d_model,
    target_features.len()
);

// Create tensors with batch dimension [batch=1, seq_len, d_model]
let target_tensor = Tensor::from_slice(
    &target_features,
    (1, 1, self.d_model),  // ← TARGET SHAPE: [1, 1, 256]
    &self.device
)?
.to_dtype(DType::F64)?;

Critical Finding: The data loader creates targets with shape [1, 1, d_model=256], containing full feature vectors (OHLCV + technical indicators), NOT single price values.

3. What the Model Currently Outputs

From /home/jgrusewski/Work/foxhunt/ml/src/mamba/mod.rs (lines 493-496):

// FIXED (Agent 246): Output projection should map d_inner to 1 for regression (price prediction)
// The model performs price regression, NOT sequence-to-sequence modeling
// Output shape: [batch, seq, d_inner] → [batch, seq, 1]
let output_projection = candle_nn::linear(d_inner, 1, vb.pp("output_proj"))?;

Current Model Output: [batch, seq, 1]

4. The Architecture Mismatch

Data Loader Target Shape:  [batch, seq, d_model=256]
Model Output Shape:        [batch, seq, 1]
                           ^^^^^^^^^^^^^^^^^ MISMATCH!

Loss Calculation Will Fail:

// ml/src/mamba/mod.rs line 1287
// Mean Squared Error for regression
let loss = ((predictions - targets)? .sqr()? .mean(DType::F64)?)?;
                         ^^^^^^^^  Cannot subtract [batch,seq,1] - [batch,seq,256]

---

Why Tests Pass But Training Will Fail

Test Environment (Synthetic Data)

From /home/jgrusewski/Work/foxhunt/ml/tests/e2e_mamba2_training.rs (line 292):

assert_eq!(output.dims()[2], 1, "Output should have 1 feature (regression)");

Tests Create Synthetic Inputs: Tests don't use DbnSequenceLoader, so they don't expose the data pipeline bug.

Production Training (Real Data)

Production training script (ml/examples/train_mamba2_dbn.rs) uses:

use ml::data_loaders::DbnSequenceLoader;
let (train_data, val_data) = loader
    .load_sequences("test_data/real/databento/ml_training_small", 0.9)
    .await?;

This Will Fail when calculating loss because:

1. Model outputs: [batch, seq, 1]
2. Data loader targets: [batch, seq, 256]
3. Cannot compute MSE between different shapes

---

Root Cause: Two Valid Interpretations, Poor Communication

Interpretation 1: Feature-to-Feature Prediction (Agent 210)

Model: Predict next timestep's full feature vector

• Input: [batch, seq, d_model] (historical features)
• Output: [batch, seq, d_model] (predicted next features)
• Target: [batch, 1, d_model] (actual next features)
• Use Case: Sequence-to-sequence forecasting (predict all 256 features)

Data Loader: ✅ SUPPORTS THIS (creates [1, 1, 256] targets)

Interpretation 2: Feature-to-Price Regression (Agent 246)

Model: Predict next price from features

• Input: [batch, seq, d_model] (historical features)
• Output: [batch, seq, 1] (predicted price)
• Target: [batch, 1, 1] (actual next price)
• Use Case: Single-value regression (predict closing price only)

Data Loader: ❌ DOES NOT SUPPORT THIS (creates 256-feature targets, not scalar prices)

---

What CLAUDE.md and ML_TRAINING_ROADMAP.md Say

CLAUDE.md (lines 103-109, 258-266)

**ML Training Service**: Model training pipeline, feature engineering (16 features + 10 technical indicators)

**MAMBA-2 Training Status** (Wave 206 - October 2025):
- ✅ **Shape Mismatch Bug Fixed**: B/C matrices now use `d_inner` (1024) instead of `d_model` (256)
- ✅ **Feature Dimension Handling**: Input features (9D) expanded to 256D via learned projection

No explicit statement about output dimensions or task type.

ML_TRAINING_ROADMAP.md (lines 100-109)

## Week 2: MAMBA-2 Training (40 hours)

**MAMBA-2 Architecture**:
- Input: 50+ features × sequence length (60 timesteps = 1 hour lookback)
- State space dimension: 128-256
- Layers: 4-8 layers
- Output: Next-bar price prediction (regression)  ← STATES "PRICE PREDICTION"

Roadmap says "price prediction" (single value), but data loader creates full feature vectors.

---

The Actual Problem: Data Loader vs Requirements Mismatch

What Should Happen

If Task = Price Regression:

1. ❌ Data loader should extract close price from target_msg
2. ❌ Create scalar target: [batch, 1, 1]
3. ✅ Model outputs: [batch, seq, 1]

If Task = Sequence-to-Sequence:

1. ✅ Data loader creates full feature vector: [batch, 1, d_model]
2. ✅ Model outputs: [batch, seq, d_model]
3. ❌ Need to change output projection back to d_inner → d_model

What Currently Happens

1. ✅ Data loader creates: [batch, 1, d_model=256] (full features)
2. ❌ Model outputs: [batch, seq, 1] (Agent 246's change)
3. ❌ SHAPE MISMATCH → Training will fail

---

Determination: Agent 246 Was INCORRECT

Why Agent 246 Was Wrong

1. Ignored Data Pipeline: Changed model without checking data loader
2. Test-Driven Design Flaw: Tests used synthetic data, didn't validate against production pipeline
3. Misread Requirements: Assumed "price prediction" meant single scalar output, but data loader disagrees

Why Agent 210 Was Actually Right

Agent 210's d_inner → d_model output projection matched the data loader's design:

• Data loader: target_tensor = [1, 1, d_model]
• Model output: [batch, seq, d_model]
• Loss calculation: ✅ COMPATIBLE SHAPES

Agent 246 "fixed" a non-existent problem by breaking the data pipeline integration.

---

What Needs to Happen Now

Option A: Revert Agent 246's Change (Recommended)

File: /home/jgrusewski/Work/foxhunt/ml/src/mamba/mod.rs

Revert:

// Line 493-496
let output_projection = candle_nn::linear(d_inner, config.d_model, vb.pp("output_proj"))?;

// Line 533
output_dim: config.d_model,  // Sequence-to-sequence (full feature prediction)

// Line 570
let output_proj_params = d_inner * config.d_model;

Justification: Matches DbnSequenceLoader target shape.

Option B: Fix Data Loader (Alternative)

File: /home/jgrusewski/Work/foxhunt/ml/src/data_loaders/dbn_sequence_loader.rs

Change (line 590-615):

// Target is next timestep's CLOSE PRICE ONLY (not full feature vector)
let target_msg = &window[self.seq_len];
let close_price = target_msg.close.to_f32().unwrap_or(0.0);
let normalized_price = (close_price - self.stats.price_mean as f32) / self.stats.price_std as f32;

let target_tensor = Tensor::from_slice(
    &[normalized_price],
    (1, 1, 1),  // [batch=1, seq=1, features=1] for scalar regression
    &self.device
)?
.to_dtype(DType::F64)?;

Justification: Makes data loader match ML_TRAINING_ROADMAP.md's stated goal of "price prediction".

Option C: Clarify Requirements (Critical)

Update Documentation to explicitly state:

CLAUDE.md:

**MAMBA-2 Use Case**: Sequence-to-sequence feature prediction (NOT single-value price regression)
- Input: [batch, seq, d_model] (historical feature sequences)
- Output: [batch, seq, d_model] (predicted next-timestep features)
- Target: [batch, 1, d_model] (actual next-timestep features)

---

Impact Analysis

If Agent 246's Change Remains

Training Script Will Fail:

$ cargo run -p ml --example train_mamba2_dbn --release -- --epochs 50
Error: shape mismatch in loss calculation
  Model output: [32, 60, 1]
  Data target:  [32, 1, 256]
  Cannot compute MSE

Production Impact:

• ❌ MAMBA-2 training cannot proceed
• ❌ 4-6 week training timeline blocked
• ❌ Ensemble model incomplete (missing MAMBA-2)

If Agent 246's Change Is Reverted

Training Script Will Work:

• ✅ Model output: [batch, seq, d_model]
• ✅ Data target: [batch, 1, d_model]
• ✅ Loss calculation succeeds
• ✅ Can proceed with 200-epoch training

But: Need to clarify if "sequence-to-sequence" is the actual requirement.

---

Recommendations

Immediate (Agent 254)

1. ✅ Revert Agent 246's changes (3 lines in ml/src/mamba/mod.rs)
2. ✅ Update e2e tests to use DbnSequenceLoader instead of synthetic data
3. ✅ Validate loss calculation with real data loader
4. ✅ Document MAMBA-2 task type in CLAUDE.md

Short-term (Agent 255)

1. ✅ Run 1-epoch training test with real DBN data
2. ✅ Verify loss converges (not NaN/Inf)
3. ✅ Check output predictions are sensible
4. ✅ Proceed with 50-epoch pilot training

Medium-term (Next Sprint)

1. 🟡 Decide: Feature-to-feature OR price-only prediction?
2. 🟡 If price-only: Fix data loader to output scalar targets
3. 🟡 If feature-to-feature: Update documentation to clarify
4. 🟡 Add integration tests that validate model + data loader compatibility

---

Files Affected

Need Immediate Changes

1. /home/jgrusewski/Work/foxhunt/ml/src/mamba/mod.rs

   • Line 496: candle_nn::linear(d_inner, config.d_model, ...)
   • Line 533: output_dim: config.d_model
   • Line 570: d_inner * config.d_model

2. /home/jgrusewski/Work/foxhunt/ml/tests/e2e_mamba2_training.rs

   • Line 292: Change assertion to assert_eq!(output.dims()[2], config.d_model)
   • Add test using DbnSequenceLoader to validate real data compatibility

3. /home/jgrusewski/Work/foxhunt/CLAUDE.md

   • Add explicit MAMBA-2 task description (sequence-to-sequence vs regression)

Validate After Changes

1. /home/jgrusewski/Work/foxhunt/ml/examples/train_mamba2_dbn.rs

   • Run with 1 epoch to validate loss calculation
   • Check for shape mismatches

2. /home/jgrusewski/Work/foxhunt/ml/src/data_loaders/dbn_sequence_loader.rs

   • Review extract_features() to understand what 256 features represent
   • Validate against MAMBA-2's expected input/output

---

Lessons Learned

1. Test Real Data Pipelines

Mistake: E2E tests used synthetic tensors, didn't validate against production data loader.

Fix: Always test model + data loader integration, not just model in isolation.

2. Clarify Requirements Upfront

Mistake: Ambiguous documentation ("price prediction" could mean scalar OR feature vector).

Fix: Explicitly document input/output shapes and task type for every model.

3. Check Dependencies Before Changing

Mistake: Agent 246 changed model output without checking what data loader produces.

Fix: Always grep for data pipeline code before architectural changes.

4. Shape Assertions Are Critical

Mistake: No runtime shape validation between model output and data loader target.

Fix: Add shape assertions at data loading and loss calculation to fail-fast on mismatches.

---

Validation Checklist

Before approving any fix:

• Model output shape matches data loader target shape
• Loss calculation runs without shape errors
• E2E tests use DbnSequenceLoader (real data pipeline)
• Documentation explicitly states MAMBA-2 task type
• 1-epoch training test passes with real DBN data
• Gradient flow works (no NaN/Inf losses)
• Checkpoint saving/loading works with new output_dim

---

Summary

Verdict: ❌ AGENT 246 WAS INCORRECT

Root Cause: Agent 246 changed model architecture without validating against the production data loader, which creates [batch, 1, d_model] targets, NOT [batch, 1, 1] scalars.

Correct Decision: Revert Agent 246's changes to restore output_dim = d_model and match the data pipeline's design.

Next Steps:

1. Revert 3 lines in ml/src/mamba/mod.rs
2. Update e2e tests to use real data loader
3. Run 1-epoch validation with DBN data
4. Proceed with 50-epoch pilot training

Impact: Unblocks MAMBA-2 training (critical for 4-6 week ML roadmap)

---

Agent 253 - Mission Complete ✅ Verdict Delivered: INCORRECT (data loader mismatch) Recommended Action: Revert Agent 246's changes immediately