Training Guide

This guide covers everything you need to know about training JEPA models effectively.

Training Overview

JEPA training follows these key principles:

Self-Supervised Learning: No labeled data required
Context-Target Prediction: Learn by predicting masked regions
Joint Embedding Space: Shared representations for context and targets
Scalable Architecture: Works from small to very large models

Training Process

The JEPA training loop consists of:

Data Loading: Load and preprocess input data
Masking: Create context and target regions
Encoding: Encode context and targets separately
Prediction: Predict target embeddings from context
Loss Computation: Compare predicted and actual target embeddings
Optimization: Update model parameters

Basic Training

Simple Training Script

import torch
from torch.utils.data import DataLoader

from jepa.models import JEPA
from jepa.models.encoder import Encoder
from jepa.models.predictor import Predictor
from jepa.trainer import create_trainer

encoder = Encoder(hidden_dim=256)
predictor = Predictor(hidden_dim=256)
model = JEPA(encoder=encoder, predictor=predictor)

train_dataset = ...  # yields (state_t, state_t1)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

trainer = create_trainer(model, learning_rate=5e-4)
trainer.train(train_loader, num_epochs=50)

# With Weights & Biases logging
trainer_wandb = create_trainer(
    model,
    learning_rate=5e-4,
    logger="wandb",
    logger_project="jepa-training",
    logger_run_name="baseline",
)

# Composite logging configuration
logger_config = {
    "console": {"enabled": True, "level": "INFO"},
    "wandb": {"enabled": True, "project": "jepa-training", "tags": ["demo"]},
    "tensorboard": {"enabled": True, "log_dir": "./logs/tensorboard"},
}

trainer_multi = create_trainer(model, logger=logger_config)

CLI Training

# Basic training
jepa-train --config my_config.yaml

# With specific parameters
python -m jepa.cli train \
  --config my_config.yaml \
  --epochs 100 \
  --batch-size 64 \
  --learning-rate 0.001

Advanced Training Techniques

Distributed Training

# Multi-GPU training
torchrun --nproc_per_node=4 python -m jepa.cli train \
  --config config.yaml \
  --distributed true

Programmatically:

trainer = create_trainer(
    model,
    distributed=True,
    world_size=int(os.environ["WORLD_SIZE"]),
    local_rank=int(os.environ.get("LOCAL_RANK", 0)),
)

Mixed Precision Training

training:
  mixed_precision: true
  grad_scaler: true
  amp_backend: "native"  # or "apex"

Gradient Accumulation

training:
  batch_size: 32           # Effective batch size
  micro_batch_size: 8      # Actual batch size per step
  gradient_accumulation_steps: 4  # 32/8 = 4

Learning Rate Scheduling

training:
  scheduler:
    type: "cosine"
    warmup_epochs: 10
    min_lr: 1e-6
    
  # Or step scheduling
  scheduler:
    type: "step"
    step_size: 30
    gamma: 0.1

Loss Functions

Standard JEPA Loss

def jepa_loss(predicted_targets, actual_targets, mask=None):
    """
    Standard JEPA prediction loss
    """
    loss = F.mse_loss(predicted_targets, actual_targets, reduction='none')
    if mask is not None:
        loss = loss * mask.unsqueeze(-1)
    return loss.mean()

Contrastive Loss

training:
  loss:
    type: "contrastive"
    temperature: 0.1
    negative_samples: 64

Multi-Scale Loss

training:
  loss:
    type: "multiscale"
    scales: [1, 2, 4]
    weights: [1.0, 0.5, 0.25]

Optimization Strategies

Optimizer Selection

training:
  optimizer:
    type: "adamw"
    learning_rate: 0.001
    weight_decay: 1e-4
    betas: [0.9, 0.999]

Learning Rate Finding

from jepa.trainer.utils import find_learning_rate

# Find optimal learning rate
trainer = JEPATrainer(config)
optimal_lr = find_learning_rate(trainer, min_lr=1e-6, max_lr=1e-1)
print(f"Optimal LR: {optimal_lr}")

Gradient Clipping

training:
  gradient_clip_norm: 1.0     # Clip by norm
  gradient_clip_value: 0.5    # Clip by value

Training Monitoring

Logging Configuration

logging:
  backends: ["wandb", "tensorboard", "console"]
  log_frequency: 10
  
  wandb:
    project: "jepa-training"
    tags: ["experiment-1"]
    
  tensorboard:
    log_dir: "logs/tensorboard"

Key Metrics to Track

Training Loss: Should decrease over time
Validation Loss: Should decrease without overfitting
Learning Rate: Track scheduling
Gradient Norm: Monitor for gradient explosion
Memory Usage: Ensure efficient GPU utilization

Early Stopping

training:
  early_stopping:
    patience: 10
    min_delta: 1e-4
    monitor: "val_loss"
    mode: "min"

Checkpointing and Resuming

Automatic Checkpointing

training:
  save_frequency: 10        # Save every 10 epochs
  save_top_k: 3            # Keep best 3 checkpoints
  checkpoint_dir: "checkpoints/"

Resume Training

python -m jepa.cli train \
  --config my_config.yaml \
  --resume checkpoints/epoch_50.pth

Custom Checkpointing

# Save custom checkpoint
trainer.save_checkpoint("my_checkpoint.pth")

# Load checkpoint (restores optimizer/scheduler automatically)
trainer.load_checkpoint("my_checkpoint.pth")

Data-Specific Training

Vision Training

model:
  encoder_type: "vision_transformer"
  patch_size: 16
  
data:
  transforms:
    resize: [224, 224]
    normalize: true
    augmentation:
      horizontal_flip: 0.5
      rotation: 15

NLP Training

model:
  encoder_type: "transformer"
  vocab_size: 50000
  
data:
  tokenizer: "bert-base-uncased"
  max_length: 512
  masking:
    mask_ratio: 0.15
    random_token_prob: 0.1

Time Series Training

model:
  encoder_type: "temporal_cnn"
  
data:
  window_size: 100
  stride: 50
  masking:
    strategy: "contiguous"
    mask_length: 20

Performance Optimization

Memory Optimization

training:
  # Use gradient checkpointing
  gradient_checkpointing: true
  
  # Optimize data loading
  data:
    num_workers: 8
    pin_memory: true
    prefetch_factor: 2

Speed Optimization

# Use compiled models (PyTorch 2.0+)
model = torch.compile(model)

# Enable optimized attention
torch.backends.cuda.enable_flash_sdp(True)

Batch Size Optimization

# Find optimal batch size
from jepa.trainer.utils import find_optimal_batch_size

optimal_batch_size = find_optimal_batch_size(
    model, 
    dataloader, 
    device="cuda"
)

Hyperparameter Tuning

Grid Search

# hyperparams.yaml
sweep:
  learning_rate: [0.001, 0.01, 0.1]
  batch_size: [32, 64, 128]
  encoder_dim: [256, 512, 1024]

Random Search

import wandb

# Initialize sweep
sweep_config = {
    'method': 'random',
    'parameters': {
        'learning_rate': {'values': [0.001, 0.01, 0.1]},
        'batch_size': {'values': [32, 64, 128]}
    }
}

sweep_id = wandb.sweep(sweep_config)
wandb.agent(sweep_id, function=train_model)

Troubleshooting

Common Issues

Loss not decreasing

Check learning rate (try lower values)
Verify data preprocessing
Check for gradient clipping

GPU out of memory

Reduce batch size
Enable gradient checkpointing
Use mixed precision

Training too slow

Increase batch size
Use more workers
Enable model compilation

Model overfitting

Add dropout
Reduce model size
Use data augmentation

Debugging Tips

# Debug mode
config.debug = True
config.training.epochs = 1
config.training.batch_size = 2

# Log gradients
for name, param in model.named_parameters():
    if param.grad is not None:
        print(f"{name}: {param.grad.norm()}")

Best Practices

Start Simple: Begin with small models and datasets
Monitor Closely: Watch loss curves and metrics
Save Often: Regular checkpointing prevents data loss
Experiment: Try different architectures and hyperparameters
Document: Keep detailed logs of experiments
Validate: Always use validation data for model selection

Advanced Topics

Custom Training Loops

class CustomJEPATrainer(JEPATrainer):
    def training_step(self, batch, batch_idx):
        # Custom training logic
        context, targets, mask = self.prepare_batch(batch)
        
        # Encode
        context_emb = self.model.encoder(context)
        target_emb = self.model.encoder(targets)
        
        # Predict
        predicted_targets = self.model.predictor(context_emb, mask)
        
        # Compute loss
        loss = self.criterion(predicted_targets, target_emb)
        
        return loss

Multi-Task Training

training:
  tasks:
    - name: "reconstruction"
      weight: 1.0
      loss: "mse"
    - name: "contrastive"
      weight: 0.5
      loss: "infonce"

Examples

For complete training examples, see:

Vision Training
NLP Training
Time Series Training
Multi-Modal Training