# Box Embedding Training: Research-Based Best Practices

## Key Findings from Research Papers

### 1. BERE (Box Embeddings for Event Relation Extraction)

**Training Configuration:**
- **Optimizer**: AMSGrad (variant of Adam)
- **Learning Rate**: 0.001 (1e-3)
- **Epochs**: 100
- **Embedding Dimension**: 768 (from RoBERTa)
- **Loss Function**: Cross-entropy (multi-class classification)
- **Key Innovation**: No explicit constraint loss needed - constraints emerge from geometry

**Critical Insight**: BERE eliminates constraint violations by design. Traditional models need explicit constraint penalties, but box embeddings guarantee anti-symmetry and reduce conjunctive constraint violations by 85-88% compared to vector models.

### 2. BoxTE (Temporal Knowledge Graph Completion)

**Training Configuration:**
- **Optimizer**: Adam or SGD with momentum
- **Learning Rate**: 0.001 - 0.01 (dataset dependent)
- **Embedding Dimension**: 200-300 (default 200)
- **Loss Function**: Margin-based ranking loss
- **Margin Parameter**: 3-10 (dataset dependent)
- **Negative Sampling**: Self-adversarial (temperature α)
- **L2 Regularization**: 1e-4 to 1e-1
- **Batch Size**: 128-1024
- **Epochs**: Up to 1000 (with early stopping)

**Key Techniques:**
- **Tanh Mapping**: Restricts embeddings to [-1, 1] range
- **Power Norm**: L1 or L2 for distance computation
- **Relation-dependent transformations**: Each relation learns distinct temporal behavior

### 3. UKGE (Uncertain Knowledge Graph Embeddings)

**Training Configuration:**
- **Loss Function**: Mean Squared Error (MSE) for confidence scores
- **Optimizer**: Adagrad (default), Adam, or SGD
- **Learning Rate**: 0.001, 0.005, 0.01, 0.05, or 0.1 (highly dataset dependent)
- **PSL Integration**: Probabilistic Soft Logic for unseen facts
- **Regularization**: L2 weight decay

**Key Innovation**: Treats confidence scores as continuous values (0-1) rather than binary labels, using PSL to infer confidence for unseen facts.

## Recommended Training Configuration for Coreference

Based on the research, here's an optimal configuration for box embeddings in coreference resolution:

### Hyperparameters

```rust
TrainingConfig {
    learning_rate: 0.001,        // Start with 0.001 (AMSGrad standard)
    negative_weight: 0.5,        // Balance positive/negative pairs
    margin: 0.3,                 // For negative pairs
    regularization: 0.0001,      // L2 regularization (1e-4)
    epochs: 100,                 // With early stopping
    batch_size: 32,              // Start small, increase if memory allows
}
```

### Optimizer: AMSGrad

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

**AMSGrad Parameters:**
- β₁ = 0.9 (first moment decay)
- β₂ = 0.999 (second moment decay)
- ε = 1e-8 (numerical stability)

### Loss Function

**For Positive Pairs (should corefer):**
```
L_pos = -log(P(A|B)) - log(P(B|A)) + λ_reg * (Vol(A) + Vol(B))
```

**For Negative Pairs (shouldn't corefer):**
```
L_neg = max(0, margin - P(A|B)) * λ_neg
```

**Key Points:**
- Use numerical stability: clamp probabilities to [ε, 1-ε] before log
- Add distance penalty when boxes don't overlap (to encourage movement)
- Regularize box volumes to prevent unbounded growth

### Initialization Strategy

**Box Embeddings:**
- **Center (μ)**: Uniform in [-0.3, 0.3] (smaller range for better overlap)
- **Width (exp(δ))**: 0.5 (larger initial width so boxes can overlap)
- **For coreference groups**: Share similar centers to ensure initial overlap

**Why This Works:**
- Boxes need to overlap initially for gradients to flow
- Smaller centers + larger widths = better initial overlap
- Shared centers for coreferent entities = faster convergence

### Gradient Computation

**Finite Differences (Current Implementation):**
- Epsilon: 1e-3 to 1e-4 (larger when boxes are disjoint)
- Gradient clipping: [-10, 10] to prevent NaN/Inf
- Check for finite values before updates

**Better Approach (Analytical Gradients):**
- Derive ∂Vol/∂μ and ∂Vol/∂δ analytically
- Much faster (O(d) vs O(4d) forward passes)
- More accurate

### Learning Rate Scheduling

**Recommended Schedule:**
1. **Warmup**: Linear increase from 0.0001 to 0.001 over first 10% of epochs
2. **Stable**: Constant 0.001 for 60% of epochs
3. **Decay**: Cosine annealing to 0.0001 over final 30% of epochs

**Alternative**: Step decay (multiply by 0.1 every 30 epochs)

### Negative Sampling

**Self-Adversarial Sampling:**
- Sample negatives according to current model predictions
- Focus on hard negatives (high predicted score but wrong)
- Temperature α = 1.0

**Implementation:**
```rust
// Sample negative with probability ∝ exp(α * score)
let prob = (alpha * score).exp();
```

### Regularization

**L2 Weight Decay:**
- Applied to all parameters
- Coefficient: 1e-4 to 1e-3
- Prevents overfitting

**Box Volume Regularization:**
- Penalize large box volumes
- Prevents boxes from growing unbounded
- Coefficient: 1e-4

### Evaluation Metrics

**Coreference Metrics:**
- MUC, B³, CEAF, LEA, BLANC, CoNLL F1
- Constraint satisfaction rate (transitivity, anti-symmetry)

**Training Metrics:**
- Loss (should decrease)
- Gradient norms (should be non-zero)
- Box overlap for positive pairs (should increase)

## Implementation Recommendations

### 1. Replace Finite Differences with Analytical Gradients

**Current (Slow):**
- O(4d) forward passes per pair
- Numerical errors from epsilon choice

**Better (Fast):**
- Derive gradients analytically
- O(1) forward pass per pair
- Exact gradients

### 2. Use AMSGrad Optimizer

**Implementation:**
```rust
// Maintain max of squared gradients
v_t = max(v_{t-1}, β₂ * v_{t-1} + (1-β₂) * g_t²)
// Update with max instead of current
θ_t = θ_{t-1} - α * m_t / (sqrt(v_t) + ε)
```

### 3. Smart Initialization

**For Coreference:**
- Entities in same chain: shared center ± small noise
- Initial width: 0.5 (ensures overlap)
- Random entities: center in [-0.3, 0.3]

### 4. Learning Rate Warmup

**Why**: Stabilizes early training when parameters are far from optimal

**How**: Linear increase from 0.0001 to 0.001 over first 10 epochs

### 5. Gradient Clipping

**Norm-based clipping:**
```rust
let grad_norm = grad.iter().map(|x| x * x).sum::<f32>().sqrt();
if grad_norm > clip_norm {
    grad = grad * (clip_norm / grad_norm);
}
```

### 6. Early Stopping

**Monitor**: Validation loss or F1 score
**Patience**: 10-20 epochs without improvement
**Checkpoint**: Save best model based on validation

## Common Pitfalls and Solutions

### Problem: Loss Not Decreasing

**Causes:**
- Boxes initialized too far apart (no overlap → zero gradients)
- Learning rate too small
- Gradients too small (finite differences with wrong epsilon)

**Solutions:**
- ✅ Smart initialization (shared centers, larger widths)
- ✅ Increase learning rate to 0.01-0.05
- ✅ Use analytical gradients or larger epsilon (1e-2) when disjoint

### Problem: Loss Increasing

**Causes:**
- Learning rate too high
- Negative pairs dominating
- Regularization too strong

**Solutions:**
- ✅ Reduce learning rate to 0.0005
- ✅ Reduce negative_weight to 0.3-0.5
- ✅ Reduce regularization to 1e-5

### Problem: NaN/Inf Values

**Causes:**
- Gradient explosion
- Division by zero in volume computation
- Log of zero in loss

**Solutions:**
- ✅ Gradient clipping (norm ≤ 10)
- ✅ Clamp probabilities: [1e-8, 1-1e-8]
- ✅ Check for finite values before updates

### Problem: Slow Training

**Causes:**
- Finite differences (O(4d) forward passes)
- Large batch size
- Inefficient gradient computation

**Solutions:**
- ✅ Use analytical gradients
- ✅ Reduce batch size
- ✅ Use sparse matrix operations (if available)

## Next Steps for Our Implementation

1. **Implement AMSGrad optimizer** (replace SGD)
2. **Add learning rate warmup** (stabilize early training)
3. **Improve initialization** (shared centers for coref groups)
4. **Add analytical gradients** (replace finite differences)
5. **Implement self-adversarial negative sampling**
6. **Add early stopping** (prevent overfitting)
7. **Better evaluation** (constraint satisfaction metrics)

## 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)
- **BoxE**: "A Box Embedding Model for Knowledge Base Completion" (NeurIPS 2020)

