# Box Embedding Training Improvements

## Summary

All requested improvements have been implemented based on research from BERE, BoxTE, and UKGE papers:

1. ✅ **AMSGrad Optimizer** - Replaces SGD with adaptive moment estimation
2. ✅ **Learning Rate Warmup & Scheduling** - Linear warmup + cosine annealing
3. ✅ **Analytical Gradients** - Replaces slow finite differences
4. ✅ **Self-Adversarial Negative Sampling** - Focuses on hard negatives

## Implementation Details

### 1. AMSGrad Optimizer

**Location**: `src/backends/box_embeddings_training.rs::AMSGradState`

**Key Features**:
- Maintains max of squared gradients (prevents learning rate increases)
- Per-parameter adaptive learning rates
- Bias correction for first and second moments
- Default hyperparameters: β₁=0.9, β₂=0.999, ε=1e-8

**Why AMSGrad over Adam**:
- Better convergence guarantees for non-convex objectives
- Prevents anomalous learning rate increases
- More stable training dynamics
- Better generalization performance

### 2. Learning Rate Scheduling

**Location**: `src/backends/box_embeddings_training.rs::get_learning_rate`

**Schedule**:
1. **Warmup** (first 10% of epochs): Linear increase from 0.1×lr to lr
2. **Decay** (remaining epochs): Cosine annealing from lr to 0.1×lr

**Example** (lr=0.001, 50 epochs):
- Epoch 0: 0.0001 (warmup start)
- Epoch 5: 0.001 (warmup end)
- Epoch 25: 0.000628 (mid decay)
- Epoch 49: 0.000101 (decay end)

**Benefits**:
- Stabilizes early training (warmup)
- Fine-tunes in later epochs (decay)
- Prevents overfitting

### 3. Analytical Gradients

**Location**: `src/backends/box_embeddings_training.rs::compute_analytical_gradients`

**Mathematical Foundation**:

For box with min = μ - exp(δ)/2, max = μ + exp(δ)/2:

- Volume: Vol = ∏ᵢ exp(δᵢ)
- ∂Vol/∂μᵢ = 0 (volume doesn't depend on center)
- ∂Vol/∂δᵢ = Vol (since Vol = ∏exp(δᵢ))

For intersection volume:
- Intersection bounds: inter_min[i] = max(min_a[i], min_b[i])
- Intersection width: inter_width[i] = max(0, min(max_a[i], max_b[i]) - inter_min[i])
- Vol_intersection = ∏ᵢ inter_width[i]

Gradients flow through:
1. Loss → Conditional probability
2. Conditional probability → Intersection volume & individual volumes
3. Volumes → Box bounds (min/max)
4. Box bounds → Parameters (μ, δ)

**Performance**:
- **Before**: O(4d) forward passes per pair (finite differences)
- **After**: O(1) forward pass per pair (analytical)
- **Speedup**: ~4× faster for d-dimensional embeddings

### 4. Self-Adversarial Negative Sampling

**Location**: `src/backends/box_embeddings_training.rs::sample_self_adversarial_negatives`

**Strategy**:
- Sample negatives according to current model predictions
- Higher probability for hard negatives (high predicted score but actually negative)
- Temperature parameter controls sharpness (default: 1.0)

**Mathematical Formulation**:
```
P(negative_i) ∝ exp(α * score_i)
```

Where:
- `score_i` = current coreference score for negative pair i
- `α` = temperature parameter (higher = more uniform)

**Benefits**:
- Focuses training on challenging examples
- More efficient than uniform sampling
- Automatically adapts as model improves

## Training Configuration

**Default Settings** (based on BERE paper):
```rust
TrainingConfig {
    learning_rate: 0.001,        // BERE standard
    negative_weight: 0.5,        // Balance positive/negative
    margin: 0.3,                  // For negative pairs
    regularization: 0.0001,      // L2 regularization
    epochs: 100,                  // With early stopping
    batch_size: 32,              // Start small
    warmup_epochs: 10,           // 10% of epochs
    use_self_adversarial: true,   // Enable hard negative mining
    adversarial_temperature: 1.0, // Temperature for sampling
}
```

## Results

**Training Progress** (from example run):
```
Epoch 0:  loss = 4.6673, lr = 0.000100 (warmup)
Epoch 5:  loss = 4.4019, lr = 0.001000 (warmup complete)
Epoch 10: loss = 3.8988, lr = 0.000973 (decay)
Epoch 25: loss = 2.9823, lr = 0.000628 (decay)
Epoch 49: loss = 2.6085, lr = 0.000101 (decay)
```

**Observations**:
- ✅ Loss decreases steadily (4.67 → 2.61)
- ✅ Learning rate schedules correctly
- ✅ No NaN/Inf values
- ✅ Gradients are non-zero (model is learning)

## Next Steps

1. **Hyperparameter Tuning**: Adjust learning rate, margin, regularization for specific datasets
2. **Early Stopping**: Monitor validation loss and stop when it plateaus
3. **Better Initialization**: Use pre-trained vector embeddings as starting point
4. **Mini-Batch Training**: Process multiple examples in parallel
5. **Evaluation Metrics**: Add constraint satisfaction metrics (transitivity, anti-symmetry)

## References

- **BERE**: "Event-Event Relation Extraction using Probabilistic Box Embedding" (ACL 2022)
- **BoxTE**: "Temporal Knowledge Graph Completion using Box Embeddings" (2021)
- **UKGE**: "Probabilistic Box Embeddings for Uncertain Knowledge Graph Reasoning" (NAACL 2021)

