Configuration Guide

This guide covers all aspects of configuring JEPA for your specific use case.

Configuration Overview

JEPA uses YAML configuration files to specify:

• Model architecture parameters

• Training hyperparameters

• Data loading settings

• Logging configuration

• Experiment tracking

Configuration Structure

A typical configuration file has these sections:

# Model architecture
model:
  encoder_dim: 512
  predictor_dim: 256
  num_layers: 6
  
# Training parameters
training:
  epochs: 100
  batch_size: 64
  learning_rate: 0.001
  
# Data settings
data:
  dataset_path: "data/"
  train_split: 0.8
  
# Logging setup
logging:
  backends: ["wandb", "tensorboard"]
  level: "INFO"

Model Configuration

Core Architecture

model:
  # Encoder settings
  encoder_dim: 512           # Hidden dimension
  encoder_layers: 6          # Number of layers
  encoder_heads: 8           # Attention heads (for Transformer)
  encoder_type: "transformer" # Options: transformer, cnn, mlp
  
  # Predictor settings
  predictor_dim: 256         # Predictor hidden size
  predictor_layers: 3        # Predictor depth
  
  # Output settings
  output_dim: 128            # Final embedding dimension
  dropout: 0.1               # Dropout rate

Advanced Architecture Options

Transformer Encoder:

model:
  encoder_type: "transformer"
  encoder_dim: 512
  encoder_layers: 12
  encoder_heads: 8
  feedforward_dim: 2048
  positional_encoding: true
  layer_norm_eps: 1e-6

CNN Encoder:

model:
  encoder_type: "cnn"
  channels: [64, 128, 256, 512]
  kernel_sizes: [3, 3, 3, 3]
  strides: [2, 2, 2, 2]
  pooling: "adaptive"

MLP Encoder:

model:
  encoder_type: "mlp"
  hidden_dims: [512, 256, 128]
  activation: "relu"
  batch_norm: true

Training Configuration

Basic Training Settings

training:
  epochs: 100               # Training epochs
  batch_size: 64            # Batch size
  learning_rate: 0.001      # Initial learning rate
  weight_decay: 1e-4        # L2 regularization
  gradient_clip: 1.0        # Gradient clipping
  
  # Validation settings
  val_frequency: 5          # Validate every N epochs
  val_patience: 10          # Early stopping patience

Optimizer Configuration

training:
  optimizer:
    type: "adamw"           # Options: adam, adamw, sgd, rmsprop
    learning_rate: 0.001
    betas: [0.9, 0.999]     # Adam betas
    weight_decay: 1e-4
    eps: 1e-8

Learning Rate Scheduling

training:
  scheduler:
    type: "cosine"          # Options: cosine, step, exponential, plateau
    warmup_epochs: 10       # Warmup period
    min_lr: 1e-6           # Minimum learning rate
    
    # For step scheduler
    step_size: 30
    gamma: 0.1
    
    # For plateau scheduler
    patience: 5
    factor: 0.5

Data Configuration

Dataset Settings

data:
  # Data paths
  dataset_path: "data/train"
  val_path: "data/val"      # Optional separate validation
  test_path: "data/test"    # Optional test set
  
  # Data splits
  train_split: 0.8          # Train/val split if no separate val
  random_seed: 42           # For reproducible splits
  
  # Loading settings
  batch_size: 64            # Can override training batch_size
  num_workers: 4            # DataLoader workers
  pin_memory: true          # GPU optimization
  shuffle: true             # Shuffle training data

Data Transforms

data:
  transforms:
    # Image transforms
    resize: [224, 224]
    normalize:
      mean: [0.485, 0.456, 0.406]
      std: [0.229, 0.224, 0.225]
    augmentation:
      horizontal_flip: 0.5
      rotation: 15
      color_jitter: 0.1
      
    # Text transforms
    tokenizer: "bert-base-uncased"
    max_length: 512
    truncation: true
    
    # Time series transforms
    window_size: 100
    stride: 50
    normalization: "z-score"

Context/Target Configuration

data:
  masking:
    strategy: "random"      # Options: random, block, structured
    mask_ratio: 0.15        # Fraction to mask
    block_size: 16          # For block masking
    
    # Context/target regions
    context_ratio: 0.85     # Fraction for context
    target_overlap: false   # Allow context/target overlap

Logging Configuration

Basic Logging

logging:
  level: "INFO"             # DEBUG, INFO, WARNING, ERROR
  backends: ["console"]     # Active logging backends
  log_dir: "logs"          # Output directory
  
  # Console settings
  console:
    format: "%(asctime)s - %(levelname)s - %(message)s"
    colored: true

Weights & Biases

logging:
  backends: ["wandb", "console"]
  
  wandb:
    project: "jepa-experiments"
    entity: "my-team"
    name: "experiment-1"     # Run name
    tags: ["baseline", "transformer"]
    notes: "Initial experiment"
    
    # What to log
    log_frequency: 10        # Log every N steps
    log_gradients: false     # Log gradient histograms
    log_model: false         # Save model artifacts

TensorBoard

logging:
  backends: ["tensorboard", "console"]
  
  tensorboard:
    log_dir: "logs/tensorboard"
    log_frequency: 10
    log_images: true         # Log sample images
    log_histograms: true     # Log parameter histograms

Multi-Backend Logging

logging:
  backends: ["wandb", "tensorboard", "console"]
  
  # Shared settings
  log_frequency: 10
  save_frequency: 50        # Save checkpoints every N epochs
  
  # Backend-specific settings
  wandb:
    project: "jepa"
  tensorboard:
    log_dir: "logs/tb"
  console:
    level: "INFO"

Environment-Specific Configs

Development Configuration

# dev_config.yaml
model:
  encoder_dim: 128          # Smaller for faster training
  encoder_layers: 2
  
training:
  epochs: 5                 # Quick testing
  batch_size: 16            # Fit on smaller GPUs
  
logging:
  backends: ["console"]     # Simple logging
  level: "DEBUG"            # Verbose output

Production Configuration

# prod_config.yaml
model:
  encoder_dim: 1024         # Large model
  encoder_layers: 24
  
training:
  epochs: 1000              # Long training
  batch_size: 128           # Large batches
  gradient_clip: 1.0        # Stability
  
logging:
  backends: ["wandb", "tensorboard"]
  wandb:
    project: "production-runs"

Loading and Merging Configs

Loading from File

from config.config import load_config

config = load_config("my_config.yaml")

Merging Configurations

from config.config import load_config, merge_configs

base_config = load_config("base_config.yaml")
override_config = load_config("overrides.yaml")

final_config = merge_configs(base_config, override_config)

Command Line Overrides

Override config values from command line:

python -m jepa.cli train \
  --config config/base.yaml \
  --learning-rate 0.01 \
  --batch-size 128 \
  --epochs 50

Validation and Debugging

Config Validation

Validate your configuration:

python -m jepa.cli config --validate my_config.yaml

View Current Config

python -m jepa.cli config --show my_config.yaml

Create Template

Generate a template configuration:

python -m jepa.cli config --create-template new_config.yaml

Best Practices

Use Version Control Track configuration files in git

Environment Variables Use environment variables for sensitive data:

logging:
  wandb:
    api_key: "${WANDB_API_KEY}"

Configuration Inheritance Create base configurations and extend them:

# child_config.yaml
base_config: "base_config.yaml"

# Override specific values
training:
  learning_rate: 0.01

Comments and Documentation Document your configurations:

model:
  encoder_dim: 512  # Chosen based on ablation study

Naming Conventions Use descriptive configuration names:

• vision_large.yaml

• nlp_debug.yaml

• production_v2.yaml

Troubleshooting

Config Not Found Check file paths are relative to working directory

Validation Errors Use --validate flag to check configuration

Memory Issues Reduce batch_size or encoder_dim

Slow Training Increase batch_size or num_workers

NaN Loss Reduce learning_rate or add gradient_clip

Examples

See the examples directory for complete configuration examples for different use cases.