jeffreys-flow-sampling

name	jeffreys-flow-sampling
description	Jeffreys Flow framework for robust Boltzmann generators and rare event sampling. Addresses mode collapse in multi-modal distributions using Jeffreys divergence + Parallel Tempering distillation. Use when: sampling rough energy landscapes, Boltzmann generators, rare events, quantum thermal states, path integral Monte Carlo, avoiding KL divergence mode collapse.

Jeffreys Flow Sampling

Robust generative framework for sampling physical systems with rough energy landscapes, based on the Jeffreys Flow method that mitigates mode collapse in Boltzmann generators.

Core Problem

Mode Collapse in Boltzmann Generators:

• Standard Boltzmann generators use reverse KL divergence
• KL divergence causes catastrophic mode collapse
• Missing modes in multi-modal distributions
• Problem in: rough energy landscapes, metastable states, rare events

Solution: Jeffreys Flow

Uses symmetric Jeffreys divergence instead of KL divergence, combined with Parallel Tempering trajectory distillation.

Key Methods

1. Jeffreys Divergence

J(P, Q) = D_KL(P|Q) + D_KL(Q|P)  (symmetric)

• Balanced: local target-seeking + global mode coverage
• Prevents mode collapse by symmetric formulation
• Corrects inherent inaccuracies through distillation

2. Parallel Tempering Distillation

1. Run Parallel Tempering to collect empirical samples
2. Distill PT trajectory data using Jeffreys divergence
3. Train generator on distilled empirical reference
4. Achieve robust mode coverage

3. Application Domains

Domain	Use Case
Statistical Mechanics	Rare event sampling in rough landscapes
Quantum Systems	Path Integral Monte Carlo for thermal states
Machine Learning	Multi-modal distribution generation
Molecular Dynamics	Metastable state sampling
Stochastic Optimization	Replica Exchange SGLD bias correction

Activation Keywords

• Jeffreys Flow
• Boltzmann generator
• rare event sampling
• mode collapse
• Parallel Tempering
• multi-modal sampling
• quantum thermal states
• 稀有事件采样
• 多模态分布

Tools Used

• exec: Run Python simulation scripts
• write: Create analysis scripts, save results
• read: Load reference materials, existing models

Instructions for Agents

Step 1: Problem Identification

Identify sampling challenges:

1. Is energy landscape rough/multi-modal?
2. Are rare events critical?
3. Is mode collapse observed with KL divergence?
4. Is Parallel Tempering feasible?

Step 2: Framework Setup

# Basic Jeffreys Flow structure
class JeffreysFlow:
    def __init__(self, energy_fn, temperatures):
        self.energy = energy_fn
        self.temps = temperatures
        
    def parallel_tempering(self, n_steps):
        """Collect empirical samples via PT"""
        # Run replica exchange at multiple temperatures
        # Store trajectory data
        pass
        
    def jeffreys_divergence(self, p_samples, q_samples):
        """Compute symmetric J divergence"""
        # D_KL(P|Q) + D_KL(Q|P)
        pass
        
    def train_generator(self, empirical_data):
        """Distill PT data with Jeffreys divergence"""
        pass

Step 3: Implementation Considerations

Temperature Selection:

• Use geometric progression: T_i = T_0 * r^(i)
• Ensure sufficient overlap between replicas
• Higher T for barrier crossing, lower T for precision

Distillation Process:

• Collect sufficient PT trajectory samples
• Use Jeffreys divergence as training objective
• Monitor mode coverage during training

Validation:

• Check all modes are captured
• Verify importance sampling weights
• Compare with ground truth where available

Step 4: Quantum Systems Application

For Path Integral Monte Carlo:

# Quantum thermal state sampling
def quantum_thermal_sampling(potential, temperature, n_particles):
    """
    Jeffreys Flow for quantum thermal states
    
    Key: Sample configurations in imaginary time path integral
    """
    # Path discretization
    # Ring polymer representation
    # Jeffreys Flow for multi-modal path distribution
    pass

Error Handling

Mode Collapse Still Occurs

• Increase temperature range in PT
• Increase empirical sample count
• Check energy landscape connectivity
• Consider hierarchical tempering

Slow Mixing

• Optimize replica exchange rate
• Adjust temperature spacing
• Use adaptive tempering schedules

Importance Sampling Bias

• Use Replica Exchange SGLD with Jeffreys correction
• Apply systematic bias correction from PT reference

Examples

Example 1: Multi-modal Distribution

User: "Sample from a 2D potential with 5 deep wells, Boltzmann generator keeps missing 2 wells"

Agent:
1. Identify: multi-modal with mode collapse risk
2. Setup: Jeffreys Flow with PT at 5 temperatures
3. Run: PT to collect empirical samples
4. Train: generator with Jeffreys divergence
5. Verify: all 5 wells captured

Example 2: Quantum Thermal State

User: "Sample quantum thermal states for a double-well potential at low temperature"

Agent:
1. Identify: quantum system, rare tunneling events
2. Setup: Path Integral representation + Jeffreys Flow
3. Run: PT across imaginary time configurations
4. Train: generator for path distribution
5. Extract: tunneling rate from importance weights

Related Papers

• Jeffreys Flow: Robust Boltzmann Generators (arxiv:2604.05303)
• Boltzmann Generators (Noé et al.)
• Parallel Tempering methods
• Path Integral Monte Carlo literature

Related Skills

• statistical-sampling-methods
• quantum-monte-carlo
• generative-models-physics

Notes

• Jeffreys divergence = symmetrized KL divergence
• Key innovation: distillation from PT empirical data
• Prevents catastrophic mode collapse
• Scalable to high-dimensional systems
• Applicable to quantum and classical systems