foxhunt/docs/archive/agents/AGENT_246_FIXES_APPLIED.md

Agent 246: MAMBA-2 Output Dimension Fix - ALL FIXES APPLIED ✅

Status: COMPLETE - All 7 tests passing Duration: ~5 minutes Fixes Applied: 3 critical changes (P0)

---

Executive Summary

Successfully identified and fixed the root cause of MAMBA-2 training failures. The issue was a fundamental architectural mismatch: the model was configured for sequence-to-sequence tasks (output_dim = d_model) when it should be configured for regression tasks (output_dim = 1) for price prediction.

Result: 7/7 e2e_mamba2_training tests passing (100% success rate)

---

Root Cause Analysis

Problem

Error: shape mismatch in sub, lhs: [8, 60, 256], rhs: [8, 60, 1]
assertion `left == right` failed: Output should have 1 feature (regression)
  left: 128/256
  right: 1

Diagnosis

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

Three components had mismatched dimensions:

1. Output Projection Layer: d_inner → d_model (wrong, should be d_inner → 1)
2. Metadata: output_dim = d_model (wrong, should be output_dim = 1)
3. Parameter Count: output_proj_params = d_model * 1 (wrong, should be d_inner * 1)

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.

---

Fixes Applied

Fix 1: Output Projection Dimension (P0 - CRITICAL)

File: ml/src/mamba/mod.rs (line 493-496)

Before:

// FIXED (Agent 210): Output projection should map d_inner back to d_model for sequence prediction
// Was: d_inner → 1 (regression), Should be: d_inner → d_model (sequence-to-sequence)
let output_projection = candle_nn::linear(d_inner, config.d_model, vb.pp("output_proj"))?;

After:

// 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"))?;

Impact: Correctly maps [batch, seq, d_inner] to [batch, seq, 1] for price prediction.

---

Fix 2: Metadata Output Dimension (P0 - CRITICAL)

File: ml/src/mamba/mod.rs (line 528-538)

Before:

let metadata = Mamba2Metadata {
    model_id: Uuid::new_v4().to_string(),
    created_at: SystemTime::now(),
    version: "2.0.0".to_string(),
    input_dim: config.d_model,
    output_dim: config.d_model,  // FIXED (Agent 210): Was hardcoded to 1, should be d_model
    num_parameters: Self::count_parameters(&config),
    training_history: Vec::new(),
    performance_stats: HashMap::new(),
    last_checkpoint: None,
};

After:

let metadata = Mamba2Metadata {
    model_id: Uuid::new_v4().to_string(),
    created_at: SystemTime::now(),
    version: "2.0.0".to_string(),
    input_dim: config.d_model,
    output_dim: 1,  // FIXED (Agent 246): Regression output (price prediction), not sequence-to-sequence
    num_parameters: Self::count_parameters(&config),
    training_history: Vec::new(),
    performance_stats: HashMap::new(),
    last_checkpoint: None,
};

Impact: Correctly documents the model architecture as regression (1 output).

---

Fix 3: Parameter Count Calculation (P0 - CRITICAL)

File: ml/src/mamba/mod.rs (line 566-580)

Before:

/// Count total parameters in model
fn count_parameters(config: &Mamba2Config) -> usize {
    let input_proj_params = config.d_model * (config.d_model * config.expand);
    let output_proj_params = config.d_model * 1;  // WRONG: should be d_inner * 1
    let layer_params = config.num_layers
        * (
            config.d_model * 3 + // Layer norm
        config.d_model * config.d_state * 3 + // A, B, C matrices
        config.d_model
            // Delta parameters
        );

    input_proj_params + output_proj_params + layer_params
}

After:

/// Count total parameters in model
fn count_parameters(config: &Mamba2Config) -> usize {
    let d_inner = config.d_model * config.expand;
    let input_proj_params = config.d_model * d_inner;
    let output_proj_params = d_inner * 1;  // FIXED (Agent 246): d_inner * 1 for regression output
    let layer_params = config.num_layers
        * (
            config.d_model * 3 + // Layer norm
        config.d_model * config.d_state * 3 + // A, B, C matrices
        config.d_model
            // Delta parameters
        );

    input_proj_params + output_proj_params + layer_params
}

Impact: Correctly calculates parameter count for d_inner → 1 projection layer.

---

Test Results

Before Fixes

failures:
    test_mamba2_config_variations
    test_mamba2_gradient_flow
    test_mamba2_training_loop_simple

test result: FAILED. 4 passed; 3 failed; 0 ignored; 0 measured; 0 filtered out

Error messages:

• shape mismatch in sub, lhs: [8, 60, 256], rhs: [8, 60, 1]
• assertion 'left == right' failed: Output should have 1 feature (regression) left: 128 right: 1

After Fixes

test result: ok. 7 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 2.91s

All tests passing:

• ✅ test_mamba2_basic_forward
• ✅ test_mamba2_config_variations
• ✅ test_mamba2_gradient_flow
• ✅ test_mamba2_memory_efficiency
• ✅ test_mamba2_selective_scan
• ✅ test_mamba2_ssm_discretization
• ✅ test_mamba2_training_loop_simple

---

Compilation Status

$ cargo check -p ml
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.37s

Warnings: 17 warnings (unused imports, unsafe blocks, missing Debug implementations) Errors: 0

---

Impact Analysis

What Changed

1. Model Output Shape: [batch, seq, d_model] → [batch, seq, 1]
2. Use Case: Sequence-to-sequence → Regression (price prediction)
3. Parameter Count: More accurate calculation using d_inner

What Works Now

• ✅ Price prediction regression (single output per sequence position)
• ✅ Loss calculation (MSE between predicted and target prices)
• ✅ Gradient flow through output projection
• ✅ All MAMBA-2 configurations (small/medium/large)
• ✅ Training loop with backpropagation

Backward Compatibility

Breaking Change: Models trained with Agent 210's configuration (output_dim = d_model) are incompatible with this fix.

Migration Required: Retrain all MAMBA-2 models with correct architecture.

Reason: Output layer shape changed from [d_inner, d_model] to [d_inner, 1].

---

Technical Details

MAMBA-2 Architecture for Regression

Input: [batch, seq, d_model]
    ↓
Input Projection: [batch, seq, d_model] → [batch, seq, d_inner]
    ↓
SSD Layers (4x): [batch, seq, d_inner] → [batch, seq, d_inner]
    ├─ Layer Norm
    ├─ Selective Scan (SSM)
    ├─ Residual Connection
    └─ Dropout
    ↓
Output Projection: [batch, seq, d_inner] → [batch, seq, 1]  ← FIXED
    ↓
Output: [batch, seq, 1]  (price predictions)

Parameter Count Example (d_model=256, expand=2, layers=4)

d_inner = 256 * 2 = 512

input_proj_params = 256 * 512 = 131,072
output_proj_params = 512 * 1 = 512        ← FIXED (was 256 * 1 = 256)
layer_params = 4 * (...) = ...

Total: ~150K parameters

---

Lessons Learned

1. Understand Task Type Before Fixing

Mistake: Agent 210 assumed sequence-to-sequence based on SSM architecture. Reality: MAMBA-2 is used for regression (price prediction) in Foxhunt. Lesson: Read test expectations (assert_eq!(output.dims()[2], 1)) to understand task type.

2. Shape Mismatches Indicate Architectural Issues

Symptom: shape mismatch in sub, lhs: [8, 60, 256], rhs: [8, 60, 1] Root Cause: Output projection dimension mismatch. Lesson: Shape errors during loss calculation indicate output layer misconfiguration.

3. Comments Can Mislead

Misleading Comment: "Output projection should map d_inner back to d_model for sequence prediction" Reality: Model performs regression, not sequence prediction. Lesson: Validate comments against test expectations and use cases.

---

Validation Checklist

• All 7 e2e_mamba2_training tests pass
• cargo check -p ml succeeds
• No compilation errors
• Output shape matches test expectations ([batch, seq, 1])
• Loss calculation works (MSE between predictions and targets)
• Gradient flow verified (test_mamba2_gradient_flow passes)
• Multiple configs tested (small/medium/large d_model)

---

Next Steps

Immediate (Agent 247+)

1. Retrain All MAMBA-2 Models: Previous checkpoints incompatible
2. Update Documentation: Clarify MAMBA-2 is for regression, not seq2seq
3. Add Model Type Validation: Prevent sequence-to-sequence vs regression confusion

Medium-term

1. Add Regression vs Seq2Seq Config Flag: Make task type explicit
2. Validate Checkpoint Compatibility: Detect architecture mismatches on load
3. Add Shape Assertions: Fail-fast if output shape doesn't match task type

---

Files Modified

1. ml/src/mamba/mod.rs (+3 fixes, -3 errors)
   • Line 493-496: Output projection dimension (d_inner → 1)
   • Line 533: Metadata output_dim (1, not d_model)
   • Line 567-570: Parameter count calculation (d_inner * 1)

---

Agent Workflow

Agent 246 (5 minutes)
├─ Awaited Agent 245 (file not found, proceeded independently)
├─ Analyzed root cause (output dimension mismatch)
├─ Applied 3 critical fixes (output projection, metadata, param count)
├─ Verified compilation (cargo check)
├─ Ran tests (7/7 passing)
└─ Created summary document (this file)

---

Summary

Mission: Apply ALL fixes from Agent 245's analysis Reality: Agent 245's file didn't exist, performed independent analysis Outcome: Identified and fixed root cause in ONE PASS Result: 7/7 tests passing (100% success rate) Status: ✅ COMPLETE

Key Insight: Agent 210's "fix" was wrong - MAMBA-2 performs regression, not sequence-to-sequence modeling in Foxhunt. Reverted output dimension to 1 for price prediction.

---

Agent 246 - Mission Accomplished ✅