# Box Embedding Training: Verification Results

## Implementation Summary

We've implemented a complete training system for box embeddings with:

### Mathematical Framework

1. **Reparameterization**: `min = μ - exp(δ)/2`, `max = μ + exp(δ)/2`
   - Ensures boxes are always valid (min ≤ max)
   - No constraint violations during training

2. **Loss Function**:
   - **Positive pairs**: `L = -log(P(A|B)) - log(P(B|A)) + λ_reg * (Vol(A) + Vol(B))`
   - **Negative pairs**: `L = max(0, margin - P(A|B)) * λ_neg`
   - Numerical stability: clamp probabilities to [ε, 1-ε] before log

3. **Gradient Computation**:
   - Finite differences: `∂L/∂θ ≈ (L(θ + ε) - L(θ)) / ε`
   - Gradient clipping: limit to [-10, 10] to prevent NaN/Inf
   - NaN detection: skip updates if gradients are invalid

4. **Optimization**:
   - SGD with learning rate 0.1-0.5
   - Delta clamping: width between 0.01 and 10.0
   - Gradient accumulation: collect all gradients before update

### Training Pipeline

1. **Data Format**: `TrainingExample` with entities and coreference chains
2. **Initialization**: Random boxes (center in [-1, 1], width = 0.2)
3. **Training Loop**: 
   - Build positive/negative pairs
   - Compute gradients
   - Accumulate and apply updates
4. **Inference**: Convert to `BoxEmbedding` and use `BoxCorefResolver`

## Current Status

✅ **Compiles and runs**  
✅ **No NaN/Inf errors** (with safety checks)  
⚠️ **Loss not decreasing** (needs investigation)

### Issues Identified

1. **Gradients may be too small**: Finite differences with ε=1e-3 might be too large
2. **Learning rate**: May need tuning
3. **Initialization**: Random boxes might not be in good starting position
4. **Entity ID matching**: Test entities need to match training entity IDs

## Next Steps

1. **Better initialization**: Start from vector embeddings (if available)
2. **Gradient debugging**: Print gradient magnitudes to verify they're non-zero
3. **Loss curve analysis**: Check if loss is actually changing (might be very small changes)
4. **Automatic differentiation**: Replace finite differences with autograd for accuracy
5. **Learning rate scheduling**: Adaptive LR based on loss

## Verification Test

The example runs successfully:
- Creates 3 training examples
- Initializes 6 boxes
- Trains for 20-50 epochs
- Loss computed (though not decreasing yet)
- Inference works (though coreference not resolved - expected with current training)

This demonstrates the **complete training pipeline** is functional, though optimization needs refinement.

