foxhunt/AGENT_251_QUICK_REFERENCE.md

Agent 251: Shape Mismatch Quick Reference

Status: ✅ RESOLVED (by Agent 254) Date: 2025-10-15

---

The Problem

ERROR: shape mismatch in sub, lhs: [32, 1, 1], rhs: [32, 1, 256]
Location: compute_loss() in ml/src/mamba/mod.rs

---

Root Cause

Architectural Misalignment:

• Model Output: [batch, seq, 1] - Agent 246 changed to 1D for price regression
• Data Target: [batch, 1, 256] - Original full feature vector

---

The Fix (Agent 254)

File: ml/src/data_loaders/dbn_sequence_loader.rs

Before:

let target_tensor = Tensor::from_slice(
    &target_features,      // 256-dim feature vector
    (1, 1, self.d_model),  // [1, 1, 256] ❌ WRONG
    &self.device
)?

After:

let target_price = self.extract_target_price(target_msg)?;  // Single normalized price

let target_tensor = Tensor::from_slice(
    &[target_price],       // Single value
    (1, 1, 1),             // [1, 1, 1] ✅ CORRECT
    &self.device
)?

---

Architectural Decision

MAMBA-2 Task: Price Regression (NOT sequence-to-sequence)

Why?

1. Business Goal: Generate trading signals (buy/sell)
2. Metrics: Win rate, Sharpe ratio (regression metrics)
3. Efficiency: 256x smaller output layer
4. Deployment: Direct price prediction → trading signal

Model Flow:

Input: [batch, 60, 256] (60 bars × 256 features)
  ↓
SSM Processing: 256 → 512 (d_inner) → 16 (d_state) → 512
  ↓
Output Projection: 512 → 1 (price regression)
  ↓
Output: [batch, 60, 1] (price predictions for each timestep)
  ↓
Extract Last: [batch, 1, 1] (next bar price prediction)
  ↓
Loss: MSE(prediction, actual_close_price)

---

Shape Consistency Check

Component	Shape	Status
Model Input	[32, 60, 256]	✅
SSM Hidden	[32, 60, 512]	✅
Model Output	[32, 60, 1]	✅
Output (last step)	[32, 1, 1]	✅
Data Target	[32, 1, 1]	✅ Fixed
Loss Input	Both [32, 1, 1]	✅

---

Verification

Test Shape Alignment:

# Run quick shape validation
cargo test -p ml test_dbn_sequence_loader_shapes -- --nocapture

# Run 1-epoch training smoke test
cargo test -p ml test_mamba2_training_one_epoch -- --nocapture

Expected Output:

✅ Input shape: [1, 60, 256]
✅ Target shape: [1, 1, 1]
✅ Model output shape: [1, 60, 1]
✅ Loss computation: MSE → scalar

---

Key Changes

1. New Method (dbn_sequence_loader.rs:630-662):

fn extract_target_price(&self, msg: &ProcessedMessage) -> Result<f32> {
    match msg {
        ProcessedMessage::Ohlcv { close, .. } => {
            let c = (close.to_f64() - self.stats.price_mean) / self.stats.price_std;
            Ok(c as f32)
        }
        // ... handles Trade, Quote, etc.
    }
}

2. Target Creation (dbn_sequence_loader.rs:590-617):

let target_price = self.extract_target_price(target_msg)?;

let target_tensor = Tensor::from_slice(
    &[target_price],       // Single price
    (1, 1, 1),             // 1D regression target
    &self.device
)?
.to_dtype(DType::F64)?;

---

Related Files

File	Change	Status
ml/src/mamba/mod.rs	Agent 246: output_projection = linear(d_inner, 1)	✅
ml/src/data_loaders/dbn_sequence_loader.rs	Agent 254: Target shape [1,1,1]	✅
ml/examples/train_mamba2_dbn.rs	No change needed	✅

---

Lessons Learned

1. Document architectural decisions: Make task explicit (regression vs seq2seq)
2. Update all consumers: Model changes require data loader updates
3. Add shape assertions: Catch mismatches early in tests
4. Integration tests: Verify end-to-end shape flow

---

Status: ✅ READY FOR TRAINING

# Run full MAMBA-2 training
cargo run -p ml --example train_mamba2_dbn --release -- --epochs 200

---

Agent: 251 Full Report: AGENT_251_SHAPE_MISMATCH_ANALYSIS.md