Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1 @@
+checkpoints/*

inl_llm/__init__.py

@@ -1,26 +1,26 @@
-"""
-Integrator Neural Language Model (INL-LLM)
-
-A novel language model architecture based on integrator dynamics and learnable equilibrium.
-
-All optimizations enabled by default (Level 1 + 2):
-- Low-rank embeddings (-87% params)
-- Shared controllers (-96% params)
-- Hierarchical equilibrium (-98% params)
-- Adaptive early stopping (+50% speed)
-- Gradient checkpointing (-65% memory)
-- Sparse excitation (10x less compute)
-
-Author: Boris Peyriguère
-"""
-
-__version__ = "2.0.0"
-__author__ = "Boris Peyriguère"
-
-# Simple API
-from .models import create_model, IntegratorLanguageModel
-
-__all__ = [
-    'create_model',  # Main API
-    'IntegratorLanguageModel',  # Main class
-]
+"""
+Integrator Neural Language Model (INL-LLM)
+
+A novel language model architecture based on integrator dynamics and learnable equilibrium.
+
+All optimizations enabled by default (Level 1 + 2):
+- Low-rank embeddings (-87% params)
+- Shared controllers (-96% params)
+- Hierarchical equilibrium (-98% params)
+- Adaptive early stopping (+50% speed)
+- Gradient checkpointing (-65% memory)
+- Sparse excitation (10x less compute)
+
+Author: Boris Peyriguère
+"""
+
+__version__ = "2.0.0"
+__author__ = "Boris Peyriguère"
+
+# Simple API
+from .models import create_model, IntegratorLanguageModel
+
+__all__ = [
+    'create_model',  # Main API
+    'IntegratorLanguageModel',  # Main class
+]

inl_llm/core/__init__.py

@@ -1,20 +1,20 @@
-"""
-Core components of INL-LLM architecture.
-
-Includes:
-- IntegratorNeuronLayer: Base integrator dynamics
-- IntegratorLoss: Loss functions with variance weighting
-- Schedulers: Equilibrium-exploration cycle schedulers
-"""
-
-from .integrator_neuron_layer import IntegratorNeuronLayer, IntegratorModel
-from .integrator_losses import IntegratorLoss, compute_convergence_metrics
-from .integrator_scheduler_v2 import create_cycle_scheduler
-
-__all__ = [
-    'IntegratorNeuronLayer',
-    'IntegratorModel',
-    'IntegratorLoss',
-    'compute_convergence_metrics',
-    'create_cycle_scheduler'
-]
+"""
+Core components of INL-LLM architecture.
+
+Includes:
+- IntegratorNeuronLayer: Base integrator dynamics
+- IntegratorLoss: Loss functions with variance weighting
+- Schedulers: Equilibrium-exploration cycle schedulers
+"""
+
+from .integrator_neuron_layer import IntegratorNeuronLayer, IntegratorModel
+from .integrator_losses import IntegratorLoss, compute_convergence_metrics
+from .integrator_scheduler_v2 import create_cycle_scheduler
+
+__all__ = [
+    'IntegratorNeuronLayer',
+    'IntegratorModel',
+    'IntegratorLoss',
+    'compute_convergence_metrics',
+    'create_cycle_scheduler'
+]

inl_llm/core/adaptive_budget_allocator.py

@@ -0,0 +1,835 @@
+"""
+Adaptive Budget Allocator for INL Architecture (ULTRA-OPTIMIZED v2)
+
+This module implements dynamic iteration budget allocation across layers:
+- Global budget pool (e.g., 125 iterations total for 25 layers)
+- Adaptive allocation based on layer complexity and convergence speed
+- Bio-inspired: Different brain regions process at different speeds
+
+Key Features:
+✅ Budget-aware: Total compute stays constant
+✅ Adaptive: Simple layers use fewer iterations, complex layers use more
+✅ Convergence-driven: Stop early when layer has converged
+✅ Multiple strategies: uniform, complexity-based, learned allocation
+
+NEW ULTRA-OPTIMIZED FEATURES (v2):
+🚀 Multi-Criteria Convergence: delta + velocity + error magnitude
+🚀 Budget Redistribution Pool: Unused budget → next layers
+🚀 Phase-Aware Allocation: Equilibrium vs Exploration phase
+🚀 Layer-Position Specialization: Early/Mid/Late layer patterns
+🚀 Loss-Component Tracking: L_speed, L_energy, L_mean awareness
+🚀 Gradient Magnitude Tracking: Allocate more to actively learning layers
+
+Author: Boris Peyriguère
+"""
+
+import torch
+import torch.nn as nn
+from typing import Dict, List, Optional, Tuple, Literal, Any
+import math
+
+
+class AdaptiveBudgetAllocator(nn.Module):
+    """
+    Manages iteration budget allocation across layers (ULTRA-OPTIMIZED v2).
+
+    Strategies:
+    - 'uniform': Equal iterations per layer (baseline)
+    - 'learned': Learnable per-layer budget allocation
+    - 'dynamic': Runtime allocation based on convergence speed
+    - 'hybrid': Combination of learned + dynamic (RECOMMENDED)
+
+    NEW v2 Features:
+    - Multi-criteria convergence detection
+    - Budget redistribution pool
+    - Phase-aware allocation (equilibrium/exploration)
+    - Layer position specialization (early/mid/late)
+    """
+
+    def __init__(
+        self,
+        num_layers: int,
+        total_budget: int,
+        strategy: Literal['uniform', 'learned', 'dynamic', 'hybrid'] = 'hybrid',
+        min_iterations_per_layer: int = 2,
+        max_iterations_per_layer: int = 15,
+        convergence_threshold: float = 1e-3,
+        warmup_iterations: int = 3,
+        # NEW v2 parameters
+        use_multi_criteria_convergence: bool = True,
+        use_budget_redistribution: bool = True,
+        use_phase_aware: bool = True,
+        use_layer_specialization: bool = True,
+        use_loss_tracking: bool = True,
+        use_gradient_tracking: bool = True,
+        velocity_threshold: float = 1e-3,
+        error_threshold: float = 1e-2,
+        redistribution_window: int = 3
+    ):
+        """
+        Args:
+            num_layers: Number of layers in the model
+            total_budget: Total iteration budget (e.g., 125 for 25 layers × 5 avg)
+            strategy: Allocation strategy
+            min_iterations_per_layer: Minimum iterations per layer
+            max_iterations_per_layer: Maximum iterations per layer
+            convergence_threshold: Threshold for early stopping (delta norm)
+            warmup_iterations: Minimum iterations before checking convergence
+
+            NEW v2 Args:
+            use_multi_criteria_convergence: Use delta + velocity + error for convergence
+            use_budget_redistribution: Redistribute unused budget to next layers
+            use_phase_aware: Adapt to equilibrium/exploration phase
+            use_layer_specialization: Early/mid/late layer patterns
+            use_loss_tracking: Track L_speed, L_energy, L_mean per layer
+            use_gradient_tracking: Track gradient magnitudes for allocation
+            velocity_threshold: Convergence threshold for velocity magnitude
+            error_threshold: Convergence threshold for error magnitude
+            redistribution_window: How many next layers to share unused budget with
+        """
+        super().__init__()
+
+        self.num_layers = num_layers
+        self.total_budget = total_budget
+        self.strategy = strategy
+        self.min_iterations = min_iterations_per_layer
+        self.max_iterations = max_iterations_per_layer
+        self.convergence_threshold = convergence_threshold
+        self.warmup_iterations = warmup_iterations
+
+        # NEW v2 feature flags
+        self.use_multi_criteria = use_multi_criteria_convergence
+        self.use_redistribution = use_budget_redistribution
+        self.use_phase_aware = use_phase_aware
+        self.use_layer_specialization = use_layer_specialization
+        self.use_loss_tracking = use_loss_tracking
+        self.use_gradient_tracking = use_gradient_tracking
+        self.velocity_threshold = velocity_threshold
+        self.error_threshold = error_threshold
+        self.redistribution_window = redistribution_window
+
+        # Learnable budget allocation (if using learned or hybrid strategy)
+        if strategy in ['learned', 'hybrid']:
+            # Initialize to uniform allocation, will be learned
+            initial_allocation = torch.ones(num_layers) / num_layers
+            self.budget_weights = nn.Parameter(initial_allocation)
+        else:
+            self.register_buffer('budget_weights', torch.ones(num_layers) / num_layers)
+
+        # Original statistics tracking
+        self.register_buffer('layer_iterations_history', torch.zeros(num_layers))
+        self.register_buffer('layer_convergence_speed', torch.ones(num_layers))
+        self.register_buffer('update_count', torch.zeros(1))
+
+        # NEW v2: Multi-criteria convergence tracking
+        self.register_buffer('layer_velocity_history', torch.zeros(num_layers))
+        self.register_buffer('layer_error_history', torch.zeros(num_layers))
+
+        # NEW v2: Phase tracking
+        self.current_phase = 'equilibrium'  # or 'exploration'
+        self.phase_multipliers = {
+            'equilibrium': 0.8,  # Use 20% less iterations in equilibrium (fast convergence)
+            'exploration': 1.2   # Use 20% more iterations in exploration (need stability)
+        }
+
+        # NEW v2: Layer position specialization patterns
+        if self.use_layer_specialization:
+            self.layer_position_weights = self._compute_layer_position_weights()
+
+        # NEW v2: Budget redistribution pool (shared across forward pass)
+        self.budget_pool = 0.0
+        self.register_buffer('unused_budget_history', torch.zeros(num_layers))
+
+        # NEW v2: Loss component tracking
+        if self.use_loss_tracking:
+            self.register_buffer('layer_L_speed', torch.zeros(num_layers))
+            self.register_buffer('layer_L_energy', torch.zeros(num_layers))
+            self.register_buffer('layer_L_mean', torch.zeros(num_layers))
+
+        # NEW v2: Gradient magnitude tracking
+        if self.use_gradient_tracking:
+            self.register_buffer('layer_grad_magnitude', torch.ones(num_layers))
+
+    def _compute_layer_position_weights(self) -> torch.Tensor:
+        """
+        Compute position-based weights for layer specialization.
+
+        Bio-inspired pattern:
+        - Early layers (0-33%): Fast processing, fewer iterations (0.8x)
+        - Middle layers (34-66%): Complex processing, more iterations (1.2x)
+        - Late layers (67-100%): Refinement, medium iterations (1.0x)
+
+        Returns:
+            Tensor of shape [num_layers] with position weights
+        """
+        weights = torch.ones(self.num_layers)
+
+        third = self.num_layers // 3
+
+        # Early layers: faster
+        weights[:third] = 0.8
+
+        # Middle layers: slower (more complex)
+        weights[third:2*third] = 1.2
+
+        # Late layers: medium
+        weights[2*third:] = 1.0
+
+        return weights
+
+    def get_layer_budget(self, layer_idx: int, training: bool = True, bonus_budget: float = 0.0) -> int:
+        """
+        Get iteration budget for a specific layer (ULTRA-OPTIMIZED v2).
+
+        NEW v2: Applies multiple adjustments:
+        - Phase-aware multiplier (equilibrium vs exploration)
+        - Layer position specialization (early/mid/late)
+        - Gradient magnitude adjustment
+        - Loss component adjustment
+        - Budget redistribution bonus
+
+        Args:
+            layer_idx: Layer index
+            training: Whether in training mode
+            bonus_budget: Bonus iterations from budget redistribution pool
+
+        Returns:
+            Number of iterations allocated to this layer
+        """
+        # Base budget calculation (original strategies)
+        if self.strategy == 'uniform':
+            base_budget = self.total_budget // self.num_layers
+
+        elif self.strategy == 'learned':
+            weights = torch.softmax(self.budget_weights, dim=0)
+            base_budget = (weights[layer_idx] * self.total_budget).item()
+
+        elif self.strategy == 'dynamic':
+            speed = self.layer_convergence_speed[layer_idx].item()
+            relative_budget = (1.0 / (speed + 0.1))
+            total_relative = sum(1.0 / (self.layer_convergence_speed[i].item() + 0.1)
+                                for i in range(self.num_layers))
+            fraction = relative_budget / total_relative
+            base_budget = fraction * self.total_budget
+
+        elif self.strategy == 'hybrid':
+            weights = torch.softmax(self.budget_weights, dim=0)
+            learned_budget = weights[layer_idx] * self.total_budget
+
+            if self.update_count.item() > 10:
+                speed = self.layer_convergence_speed[layer_idx].item()
+                speed_factor = 1.0 / (speed + 0.1)
+                avg_speed_factor = sum(1.0 / (self.layer_convergence_speed[i].item() + 0.1)
+                                      for i in range(self.num_layers)) / self.num_layers
+                adjustment = speed_factor / avg_speed_factor
+                learned_budget = learned_budget * adjustment
+
+            base_budget = learned_budget.item()
+        else:
+            raise ValueError(f"Unknown strategy: {self.strategy}")
+
+        # NEW v2: Apply phase-aware multiplier
+        if self.use_phase_aware:
+            phase_mult = self.phase_multipliers.get(self.current_phase, 1.0)
+            base_budget *= phase_mult
+
+        # NEW v2: Apply layer position specialization
+        if self.use_layer_specialization:
+            pos_weight = self.layer_position_weights[layer_idx].item()
+            base_budget *= pos_weight
+
+        # NEW v2: Apply gradient magnitude adjustment (after warmup)
+        if self.use_gradient_tracking and self.update_count.item() > 10:
+            grad_mag = self.layer_grad_magnitude[layer_idx].item()
+            avg_grad = self.layer_grad_magnitude.mean().item()
+            if avg_grad > 1e-8:
+                grad_adjustment = grad_mag / avg_grad
+                # Clip to reasonable range [0.8, 1.3]
+                grad_adjustment = max(0.8, min(1.3, grad_adjustment))
+                base_budget *= grad_adjustment
+
+        # NEW v2: Apply loss component adjustment (high L_speed = needs more iterations)
+        if self.use_loss_tracking and self.update_count.item() > 10:
+            L_speed = self.layer_L_speed[layer_idx].item()
+            L_energy = self.layer_L_energy[layer_idx].item()
+
+            # High speed loss = slow convergence = more iterations needed
+            avg_speed = self.layer_L_speed.mean().item()
+            if avg_speed > 1e-8:
+                speed_adjustment = 1.0 + 0.2 * (L_speed / avg_speed - 1.0)
+                speed_adjustment = max(0.9, min(1.2, speed_adjustment))
+                base_budget *= speed_adjustment
+
+        # NEW v2: Add bonus from redistribution pool
+        base_budget += bonus_budget
+
+        # Final budget with bounds
+        budget = int(base_budget)
+        return max(self.min_iterations, min(self.max_iterations, budget))
+
+    def check_convergence(
+        self,
+        x_current: torch.Tensor,
+        x_prev: torch.Tensor,
+        iteration: int,
+        v_current: Optional[torch.Tensor] = None,
+        mu: Optional[torch.Tensor] = None
+    ) -> Tuple[bool, Dict[str, float]]:
+        """
+        Check if layer has converged (ULTRA-OPTIMIZED v2 with multi-criteria).
+
+        NEW v2: Multi-criteria convergence detection:
+        1. Delta norm: ||x_current - x_prev|| < threshold (original)
+        2. Velocity magnitude: ||v|| < velocity_threshold (new)
+        3. Error magnitude: ||x - mu|| < error_threshold (new)
+
+        All criteria must be satisfied for convergence (AND logic).
+
+        Args:
+            x_current: Current state [batch_size, d_model]
+            x_prev: Previous state [batch_size, d_model]
+            iteration: Current iteration number
+            v_current: Current velocity [batch_size, d_model] (optional, for multi-criteria)
+            mu: Learned equilibrium [batch_size, d_model] or scalar (optional, for error check)
+
+        Returns:
+            converged: True if converged, False otherwise
+            metrics: Dictionary with convergence metrics
+        """
+        if iteration < self.warmup_iterations:
+            return False, {'delta': 0.0, 'velocity': 0.0, 'error': 0.0}
+
+        metrics = {}
+
+        # Criterion 1: Delta norm (original)
+        delta = torch.norm(x_current - x_prev, dim=-1).mean()
+        metrics['delta'] = delta.item()
+        delta_converged = delta.item() < self.convergence_threshold
+
+        # If multi-criteria is disabled, return early
+        if not self.use_multi_criteria:
+            return delta_converged, metrics
+
+        # Criterion 2: Velocity magnitude (NEW v2)
+        velocity_converged = True
+        if v_current is not None:
+            v_mag = torch.norm(v_current, dim=-1).mean()
+            metrics['velocity'] = v_mag.item()
+            velocity_converged = v_mag.item() < self.velocity_threshold
+        else:
+            metrics['velocity'] = 0.0
+
+        # Criterion 3: Error magnitude (NEW v2)
+        error_converged = True
+        if mu is not None:
+            error = torch.norm(x_current - mu, dim=-1).mean()
+            metrics['error'] = error.item()
+            error_converged = error.item() < self.error_threshold
+        else:
+            metrics['error'] = 0.0
+
+        # ALL criteria must be satisfied (AND logic)
+        converged = delta_converged and velocity_converged and error_converged
+
+        return converged, metrics
+
+    def update_statistics(
+        self,
+        layer_idx: int,
+        iterations_used: int,
+        final_delta: float,
+        budget_allocated: int = 0,
+        final_velocity: float = 0.0,
+        final_error: float = 0.0,
+        loss_components: Optional[Dict[str, float]] = None,
+        grad_magnitude: Optional[float] = None
+    ):
+        """
+        Update layer statistics after processing (ULTRA-OPTIMIZED v2).
+
+        NEW v2: Tracks additional metrics:
+        - Velocity magnitude
+        - Error magnitude
+        - Unused budget
+        - Loss components (L_speed, L_energy, L_mean)
+        - Gradient magnitude
+
+        Args:
+            layer_idx: Layer index
+            iterations_used: Number of iterations actually used
+            final_delta: Final convergence delta (smaller = faster convergence)
+            budget_allocated: Budget that was allocated (NEW v2)
+            final_velocity: Final velocity magnitude (NEW v2)
+            final_error: Final error magnitude (NEW v2)
+            loss_components: Dict with L_speed, L_energy, L_mean (NEW v2)
+            grad_magnitude: Gradient magnitude for this layer (NEW v2)
+        """
+        alpha = 0.9  # Exponential moving average
+
+        # Original statistics
+        self.layer_iterations_history[layer_idx] = (
+            alpha * self.layer_iterations_history[layer_idx] +
+            (1 - alpha) * iterations_used
+        )
+
+        speed = 1.0 / (final_delta + 1e-6)
+        self.layer_convergence_speed[layer_idx] = (
+            alpha * self.layer_convergence_speed[layer_idx] +
+            (1 - alpha) * speed
+        )
+
+        # NEW v2: Track velocity
+        self.layer_velocity_history[layer_idx] = (
+            alpha * self.layer_velocity_history[layer_idx] +
+            (1 - alpha) * final_velocity
+        )
+
+        # NEW v2: Track error
+        self.layer_error_history[layer_idx] = (
+            alpha * self.layer_error_history[layer_idx] +
+            (1 - alpha) * final_error
+        )
+
+        # NEW v2: Track unused budget
+        if budget_allocated > 0:
+            unused = budget_allocated - iterations_used
+            self.unused_budget_history[layer_idx] = (
+                alpha * self.unused_budget_history[layer_idx] +
+                (1 - alpha) * unused
+            )
+
+        # NEW v2: Track loss components
+        if self.use_loss_tracking and loss_components is not None:
+            if 'L_speed' in loss_components:
+                self.layer_L_speed[layer_idx] = (
+                    alpha * self.layer_L_speed[layer_idx] +
+                    (1 - alpha) * loss_components['L_speed']
+                )
+            if 'L_energy' in loss_components:
+                self.layer_L_energy[layer_idx] = (
+                    alpha * self.layer_L_energy[layer_idx] +
+                    (1 - alpha) * loss_components['L_energy']
+                )
+            if 'L_mean' in loss_components:
+                self.layer_L_mean[layer_idx] = (
+                    alpha * self.layer_L_mean[layer_idx] +
+                    (1 - alpha) * loss_components['L_mean']
+                )
+
+        # NEW v2: Track gradient magnitude
+        if self.use_gradient_tracking and grad_magnitude is not None:
+            self.layer_grad_magnitude[layer_idx] = (
+                alpha * self.layer_grad_magnitude[layer_idx] +
+                (1 - alpha) * grad_magnitude
+            )
+
+        self.update_count += 1
+
+    def set_phase(self, phase: str):
+        """
+        Set training phase for phase-aware budget allocation.
+
+        Args:
+            phase: 'equilibrium' or 'exploration'
+        """
+        if phase not in ['equilibrium', 'exploration']:
+            raise ValueError(f"Unknown phase: {phase}. Use 'equilibrium' or 'exploration'.")
+        self.current_phase = phase
+
+    def reset_budget_pool(self):
+        """
+        Reset the budget redistribution pool (call at start of forward pass).
+        """
+        self.budget_pool = 0.0
+
+    def add_to_budget_pool(self, unused_iterations: int):
+        """
+        Add unused iterations to redistribution pool.
+
+        Args:
+            unused_iterations: Number of unused iterations from a layer
+        """
+        if self.use_redistribution:
+            self.budget_pool += unused_iterations
+
+    def get_redistribution_bonus(self, layer_idx: int) -> float:
+        """
+        Get bonus iterations from redistribution pool for a layer.
+
+        Distributes pool evenly across next N layers (redistribution_window).
+
+        Args:
+            layer_idx: Current layer index
+
+        Returns:
+            Bonus iterations from pool
+        """
+        if not self.use_redistribution or self.budget_pool <= 0:
+            return 0.0
+
+        # Distribute to next N layers
+        remaining_layers = self.num_layers - layer_idx
+        if remaining_layers <= 0:
+            return 0.0
+
+        # Distribute pool across min(remaining_layers, window)
+        window = min(remaining_layers, self.redistribution_window)
+        bonus = self.budget_pool / window
+
+        # Deduct from pool
+        self.budget_pool -= bonus
+
+        return bonus
+
+    def get_all_budgets(self, training: bool = True) -> List[int]:
+        """
+        Get budget allocation for all layers.
+
+        Args:
+            training: Whether in training mode
+
+        Returns:
+            List of iteration budgets for each layer
+        """
+        budgets = [self.get_layer_budget(i, training) for i in range(self.num_layers)]
+
+        # Ensure total doesn't exceed budget (adjust if needed)
+        total = sum(budgets)
+        if total > self.total_budget:
+            # Scale down proportionally
+            scale = self.total_budget / total
+            budgets = [max(self.min_iterations, int(b * scale)) for b in budgets]
+
+        return budgets
+
+    def get_statistics(self) -> Dict[str, Any]:
+        """
+        Get current allocation statistics (ULTRA-OPTIMIZED v2).
+
+        NEW v2: Includes all new metrics tracked by v2 allocator.
+
+        Returns:
+            Dictionary with comprehensive statistics
+        """
+        budgets = self.get_all_budgets(training=False)
+
+        stats = {
+            # Original statistics
+            'layer_budgets': torch.tensor(budgets),
+            'layer_iterations_history': self.layer_iterations_history.clone(),
+            'layer_convergence_speed': self.layer_convergence_speed.clone(),
+            'budget_weights': torch.softmax(self.budget_weights, dim=0) if self.strategy in ['learned', 'hybrid'] else self.budget_weights,
+            'total_budget': torch.tensor(self.total_budget),
+            'updates': self.update_count.clone(),
+
+            # NEW v2: Multi-criteria convergence tracking
+            'layer_velocity_history': self.layer_velocity_history.clone(),
+            'layer_error_history': self.layer_error_history.clone(),
+
+            # NEW v2: Phase information
+            'current_phase': self.current_phase,
+            'phase_multipliers': self.phase_multipliers,
+
+            # NEW v2: Budget redistribution
+            'unused_budget_history': self.unused_budget_history.clone(),
+            'current_budget_pool': self.budget_pool,
+        }
+
+        # NEW v2: Layer position weights (if enabled)
+        if self.use_layer_specialization:
+            stats['layer_position_weights'] = self.layer_position_weights.clone()
+
+        # NEW v2: Loss component tracking (if enabled)
+        if self.use_loss_tracking:
+            stats['layer_L_speed'] = self.layer_L_speed.clone()
+            stats['layer_L_energy'] = self.layer_L_energy.clone()
+            stats['layer_L_mean'] = self.layer_L_mean.clone()
+
+        # NEW v2: Gradient magnitude tracking (if enabled)
+        if self.use_gradient_tracking:
+            stats['layer_grad_magnitude'] = self.layer_grad_magnitude.clone()
+
+        # NEW v2: Feature flags summary
+        stats['v2_features'] = {
+            'multi_criteria_convergence': self.use_multi_criteria,
+            'budget_redistribution': self.use_redistribution,
+            'phase_aware': self.use_phase_aware,
+            'layer_specialization': self.use_layer_specialization,
+            'loss_tracking': self.use_loss_tracking,
+            'gradient_tracking': self.use_gradient_tracking
+        }
+
+        return stats
+
+    def __repr__(self) -> str:
+        budgets = self.get_all_budgets(training=False)
+        avg_budget = sum(budgets) / len(budgets)
+        min_budget = min(budgets)
+        max_budget = max(budgets)
+
+        # NEW v2: Count enabled features
+        enabled_features = []
+        if self.use_multi_criteria:
+            enabled_features.append("multi-criteria")
+        if self.use_redistribution:
+            enabled_features.append("redistribution")
+        if self.use_phase_aware:
+            enabled_features.append("phase-aware")
+        if self.use_layer_specialization:
+            enabled_features.append("layer-spec")
+        if self.use_loss_tracking:
+            enabled_features.append("loss-track")
+        if self.use_gradient_tracking:
+            enabled_features.append("grad-track")
+
+        features_str = ", ".join(enabled_features) if enabled_features else "none"
+
+        return (
+            f"AdaptiveBudgetAllocator-v2(\n"
+            f"  strategy={self.strategy},\n"
+            f"  num_layers={self.num_layers},\n"
+            f"  total_budget={self.total_budget},\n"
+            f"  avg_budget_per_layer={avg_budget:.1f},\n"
+            f"  budget_range=[{min_budget}, {max_budget}],\n"
+            f"  convergence_threshold={self.convergence_threshold:.1e},\n"
+            f"  phase={self.current_phase},\n"
+            f"  v2_features=[{features_str}]\n"
+            f")"
+        )
+
+
+class BudgetAwareINLLayer(nn.Module):
+    """
+    Wrapper for INL layers that respects budget allocation (ULTRA-OPTIMIZED v2).
+
+    Handles:
+    - Dynamic iteration count based on budget
+    - Early stopping when converged (multi-criteria)
+    - Statistics tracking for budget allocator
+    - Budget redistribution pool management
+
+    NEW v2 Features:
+    - Multi-criteria convergence checking
+    - Budget redistribution to next layers
+    - Loss component extraction
+    - Gradient magnitude tracking
+    """
+
+    def __init__(
+        self,
+        inl_layer: nn.Module,
+        layer_idx: int,
+        budget_allocator: Optional[AdaptiveBudgetAllocator] = None
+    ):
+        """
+        Args:
+            inl_layer: The base INL layer to wrap
+            layer_idx: Index of this layer
+            budget_allocator: Budget allocator (if None, uses default iterations)
+        """
+        super().__init__()
+
+        self.inl_layer = inl_layer
+        self.layer_idx = layer_idx
+        self.budget_allocator = budget_allocator
+
+    def forward(
+        self,
+        h: torch.Tensor,
+        x_init: torch.Tensor,
+        v_init: torch.Tensor,
+        default_iterations: int = 5,
+        return_trajectory: bool = False,
+        mu: Optional[torch.Tensor] = None,
+        loss_components: Optional[Dict[str, float]] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor, Dict]:
+        """
+        Forward pass with budget-aware iteration control (ULTRA-OPTIMIZED v2).
+
+        NEW v2: Includes multi-criteria convergence, budget redistribution,
+        and comprehensive statistics tracking.
+
+        Args:
+            h: Context embedding [batch_size * seq_len, d_model]
+            x_init: Initial state [batch_size * seq_len, d_model]
+            v_init: Initial velocity [batch_size * seq_len, d_model]
+            default_iterations: Default iterations if no budget allocator
+            return_trajectory: Whether to return full trajectory
+            mu: Learned equilibrium (for error-based convergence) (NEW v2)
+            loss_components: Loss components dict (L_speed, L_energy, L_mean) (NEW v2)
+
+        Returns:
+            x_final: Final state
+            v_final: Final velocity
+            info: Dictionary with statistics
+        """
+        # NEW v2: Get budget with redistribution bonus
+        if self.budget_allocator is not None:
+            bonus = self.budget_allocator.get_redistribution_bonus(self.layer_idx)
+            max_iters = self.budget_allocator.get_layer_budget(
+                self.layer_idx,
+                training=self.training,
+                bonus_budget=bonus
+            )
+        else:
+            max_iters = default_iterations
+
+        # Run iterations
+        x, v = x_init, v_init
+        x_prev = x_init
+
+        if return_trajectory:
+            x_traj = [x.clone()]
+            v_traj = [v.clone()]
+
+        actual_iterations = 0
+        converged = False
+        convergence_metrics = {}
+
+        for iteration in range(max_iters):
+            # One integration step
+            x_next, v_next, aux = self.inl_layer(h, x, v, step=iteration)
+
+            # NEW v2: Check convergence with multi-criteria (if budget allocator available)
+            if self.budget_allocator is not None and iteration >= self.budget_allocator.warmup_iterations:
+                converged, convergence_metrics = self.budget_allocator.check_convergence(
+                    x_next, x, iteration,
+                    v_current=v_next,  # NEW v2: velocity for multi-criteria
+                    mu=mu  # NEW v2: equilibrium for error-based check
+                )
+                if converged and not self.training:
+                    # Early stop during inference
+                    x, v = x_next, v_next
+                    actual_iterations = iteration + 1
+                    break
+
+            x_prev = x
+            x, v = x_next, v_next
+            actual_iterations = iteration + 1
+
+            if return_trajectory:
+                x_traj.append(x.clone())
+                v_traj.append(v.clone())
+
+        # NEW v2: Add unused budget to redistribution pool
+        if self.budget_allocator is not None:
+            unused = max_iters - actual_iterations
+            self.budget_allocator.add_to_budget_pool(unused)
+
+        # NEW v2: Update statistics with all new metrics (during training)
+        if self.training and self.budget_allocator is not None:
+            final_delta = torch.norm(x - x_prev, dim=-1).mean().item()
+            final_velocity = torch.norm(v, dim=-1).mean().item() if v is not None else 0.0
+            final_error = torch.norm(x - mu, dim=-1).mean().item() if mu is not None else 0.0
+
+            # Extract gradient magnitude if possible
+            grad_mag = None
+            if x.requires_grad and x.grad is not None:
+                grad_mag = torch.norm(x.grad, dim=-1).mean().item()
+
+            self.budget_allocator.update_statistics(
+                self.layer_idx,
+                actual_iterations,
+                final_delta,
+                budget_allocated=max_iters,  # NEW v2
+                final_velocity=final_velocity,  # NEW v2
+                final_error=final_error,  # NEW v2
+                loss_components=loss_components,  # NEW v2
+                grad_magnitude=grad_mag  # NEW v2
+            )
+
+        # Prepare output info
+        info = {
+            'iterations_used': actual_iterations,
+            'max_iterations': max_iters,
+            'converged': converged,
+            'layer_idx': self.layer_idx,
+            'convergence_metrics': convergence_metrics  # NEW v2
+        }
+
+        if return_trajectory:
+            info['x_trajectory'] = torch.stack(x_traj, dim=1)
+            info['v_trajectory'] = torch.stack(v_traj, dim=1)
+
+        return x, v, info
+
+
+def create_budget_allocator(
+    num_layers: int,
+    avg_iterations_per_layer: int = 5,
+    strategy: str = 'hybrid',
+    **kwargs
+) -> AdaptiveBudgetAllocator:
+    """
+    Helper function to create a budget allocator.
+
+    Args:
+        num_layers: Number of layers
+        avg_iterations_per_layer: Average iterations per layer (determines total budget)
+        strategy: Allocation strategy
+        **kwargs: Additional arguments for AdaptiveBudgetAllocator
+
+    Returns:
+        Configured AdaptiveBudgetAllocator
+    """
+    total_budget = num_layers * avg_iterations_per_layer
+
+    return AdaptiveBudgetAllocator(
+        num_layers=num_layers,
+        total_budget=total_budget,
+        strategy=strategy,
+        **kwargs
+    )
+
+
+if __name__ == '__main__':
+    print("=" * 70)
+    print("ADAPTIVE BUDGET ALLOCATOR - Test")
+    print("=" * 70)
+
+    # Create allocator
+    allocator = create_budget_allocator(
+        num_layers=25,
+        avg_iterations_per_layer=5,
+        strategy='hybrid'
+    )
+
+    print(f"\n{allocator}")
+
+    # Test budget allocation
+    print("\n📊 Initial Budget Allocation:")
+    budgets = allocator.get_all_budgets()
+    for i, budget in enumerate(budgets):
+        print(f"  Layer {i:2d}: {budget:2d} iterations")
+
+    print(f"\n✅ Total budget: {sum(budgets)} / {allocator.total_budget}")
+
+    # Simulate some updates
+    print("\n🔄 Simulating convergence updates...")
+    for i in range(25):
+        # Simulate: early layers converge faster, later layers slower
+        convergence_speed = 1.0 if i < 10 else 0.5
+        final_delta = 0.001 * convergence_speed
+        iterations = 4 if i < 10 else 7
+
+        allocator.update_statistics(i, iterations, final_delta)
+
+    # Check updated allocation
+    print("\n📊 Updated Budget Allocation (after learning):")
+    budgets_updated = allocator.get_all_budgets()
+    for i, budget in enumerate(budgets_updated):
+        change = "+" if budget > budgets[i] else ("-" if budget < budgets[i] else " ")
+        print(f"  Layer {i:2d}: {budget:2d} iterations {change}")
+
+    print(f"\n✅ Total budget: {sum(budgets_updated)} / {allocator.total_budget}")
+
+    # Show statistics
+    print("\n📈 Statistics:")
+    stats = allocator.get_statistics()
+    print(f"  Updates: {stats['updates'].item():.0f}")
+    print(f"  Convergence speeds (first 5 layers): {stats['layer_convergence_speed'][:5].tolist()}")
+    print(f"  Convergence speeds (last 5 layers): {stats['layer_convergence_speed'][-5:].tolist()}")
+
+    print("\n" + "=" * 70)
+    print("✅ ADAPTIVE BUDGET ALLOCATOR WORKING!")
+    print("=" * 70)

inl_llm/core/integrator_losses.py

@@ -1,352 +1,352 @@
-"""
-Adaptive Loss Functions for IntegratorNeuronLayer Training
-
-Implements:
-- L_task: Main task loss (MSE or CE)
-- L_mean: Soft constraint to encourage convergence towards target
-- L_speed: Penalizes slow convergence in early iterations
-- L_energy: Regularizes velocity to prevent wild oscillations
-"""
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from typing import Dict, Optional
-
-
-class IntegratorLoss(nn.Module):
-    """
-    Combined loss function with adaptive weighting for curriculum learning.
-    """
-
-    def __init__(
-        self,
-        target_value: float = 5.0,
-        lambda_mean_init: float = 1.0,
-        lambda_speed: float = 0.1,
-        lambda_energy: float = 0.01,
-        energy_p: float = 2.0,
-        annealing_schedule: str = 'exponential',
-        annealing_factor: float = 0.1,
-        annealing_epochs: int = 100,
-        variance_weighted: bool = True,
-        exploration_phase: bool = False,
-        exploration_lambda_mean: float = 0.05,
-        exploration_lambda_energy: float = 0.001,
-        task_loss_type: str = 'mse'  # 'mse' for regression, 'ce' for classification (LM)
-    ):
-        """
-        Args:
-            target_value: Target value for convergence (default 5.0)
-            lambda_mean_init: Initial weight for L_mean (will be annealed)
-            lambda_speed: Weight for L_speed (convergence speed penalty)
-            lambda_energy: Weight for L_energy (velocity regularization)
-            energy_p: Power for energy loss (2.0 = L2, 1.0 = L1)
-            annealing_schedule: 'exponential' or 'linear'
-            annealing_factor: Target factor for lambda_mean after annealing
-            annealing_epochs: Number of epochs to anneal over
-            variance_weighted: Use variance-weighted regularization
-            exploration_phase: Current phase (equilibrium=False, exploration=True)
-            exploration_lambda_mean: Lambda mean during exploration phase
-            exploration_lambda_energy: Lambda energy during exploration phase
-            task_loss_type: 'mse' for regression, 'ce' for classification/language modeling
-        """
-        super().__init__()
-
-        self.target_value = target_value
-        self.lambda_mean_init = lambda_mean_init
-        self.lambda_speed = lambda_speed
-        self.lambda_energy = lambda_energy
-        self.energy_p = energy_p
-        self.annealing_schedule = annealing_schedule
-        self.annealing_factor = annealing_factor
-        self.annealing_epochs = annealing_epochs
-
-        # Phase control and variance weighting
-        self.variance_weighted = variance_weighted
-        self.exploration_phase = exploration_phase
-        self.exploration_lambda_mean = exploration_lambda_mean
-        self.exploration_lambda_energy = exploration_lambda_energy
-
-        # Task loss type: MSE for regression, CrossEntropy for classification (language models)
-        self.task_loss_type = task_loss_type
-        if task_loss_type == 'mse':
-            self.task_loss = nn.MSELoss()
-        elif task_loss_type == 'ce':
-            self.task_loss = nn.CrossEntropyLoss()
-        else:
-            raise ValueError(f"Unknown task_loss_type: {task_loss_type}. Use 'mse' or 'ce'.")
-
-    def get_lambda_mean(self, epoch: int) -> float:
-        """
-        Compute current lambda_mean based on annealing schedule.
-
-        Args:
-            epoch: Current training epoch
-
-        Returns:
-            Current lambda_mean value
-        """
-        if epoch >= self.annealing_epochs:
-            return self.lambda_mean_init * self.annealing_factor
-
-        progress = epoch / self.annealing_epochs
-
-        if self.annealing_schedule == 'exponential':
-            # Exponential decay: lambda_mean = init * (factor)^progress
-            lambda_mean = self.lambda_mean_init * (self.annealing_factor ** progress)
-        elif self.annealing_schedule == 'linear':
-            # Linear decay: lambda_mean = init * (1 - progress * (1 - factor))
-            lambda_mean = self.lambda_mean_init * (1 - progress * (1 - self.annealing_factor))
-        else:
-            raise ValueError(f"Unknown annealing schedule: {self.annealing_schedule}")
-
-        return lambda_mean
-
-    def compute_L_task(
-        self,
-        predictions: torch.Tensor,
-        targets: torch.Tensor
-    ) -> torch.Tensor:
-        """
-        Main task loss: MSE between final prediction and target.
-
-        Args:
-            predictions: Model predictions [batch_size, output_dim]
-            targets: Ground truth targets [batch_size, output_dim]
-
-        Returns:
-            Scalar loss
-        """
-        return self.task_loss(predictions, targets)
-
-    def compute_L_mean(
-        self,
-        x_final: torch.Tensor,
-        epoch: int,
-        learned_mu: Optional[torch.Tensor] = None
-    ) -> torch.Tensor:
-        """
-        Mean constraint loss: encourages batch mean to be close to target.
-        Supports variance-weighted regularization and learned mu.
-
-        Args:
-            x_final: Final state x_T [batch_size, output_dim]
-            epoch: Current epoch for annealing
-            learned_mu: Learned equilibrium attractor, if None uses target_value
-
-        Returns:
-            Scalar loss
-        """
-        # Use exploration phase lambda if in exploration mode
-        if self.exploration_phase:
-            lambda_mean = self.exploration_lambda_mean
-        else:
-            lambda_mean = self.get_lambda_mean(epoch)
-
-        # Use learned mu if provided, otherwise fixed target
-        target = learned_mu if learned_mu is not None else self.target_value
-
-        # Variance-weighted regularization
-        if self.variance_weighted:
-            # Compute per-neuron variance across batch
-            x_var = torch.var(x_final, dim=0, keepdim=False)  # [output_dim]
-            # Weight inversely proportional to variance (stable neurons penalized less)
-            weights = 1.0 / (1.0 + x_var)  # [output_dim]
-            # Normalize weights
-            weights = weights / weights.sum() * weights.numel()
-            # Weighted penalty
-            deviations = (x_final - target) ** 2  # [batch_size, output_dim]
-            loss = lambda_mean * (weights * deviations.mean(dim=0)).mean()
-        else:
-            # Uniform weighting
-            batch_mean = x_final.mean(dim=0)  # [output_dim]
-            loss = lambda_mean * ((batch_mean - target) ** 2).mean()
-
-        return loss
-
-    def compute_L_speed(
-        self,
-        x_trajectory: torch.Tensor
-    ) -> torch.Tensor:
-        """
-        Speed loss: penalizes deviation from target in early iterations.
-
-        Uses weighted sum: w_t = exp(-t / tau) to prioritize early steps.
-
-        Args:
-            x_trajectory: Trajectory of states [batch_size, T+1, output_dim]
-
-        Returns:
-            Scalar loss
-        """
-        T = x_trajectory.shape[1] - 1  # Exclude initial state
-        if T == 0:
-            return torch.tensor(0.0, device=x_trajectory.device)
-
-        # Exponentially decaying weights: prioritize early iterations
-        tau = T / 3.0  # Decay constant
-        t_indices = torch.arange(1, T + 1, device=x_trajectory.device, dtype=torch.float32)
-        weights = torch.exp(-t_indices / tau)
-        weights = weights / weights.sum()  # Normalize
-
-        # Compute weighted deviation from target
-        deviations = torch.abs(x_trajectory[:, 1:, :] - self.target_value)  # [B, T, output_dim]
-        weighted_dev = (deviations * weights.view(1, -1, 1)).sum(dim=1)  # [B, output_dim]
-
-        loss = self.lambda_speed * weighted_dev.mean()
-        return loss
-
-    def compute_L_energy(
-        self,
-        v_trajectory: torch.Tensor
-    ) -> torch.Tensor:
-        """
-        Energy loss: regularizes velocity to prevent oscillations.
-
-        Args:
-            v_trajectory: Trajectory of velocities [batch_size, T+1, output_dim]
-
-        Returns:
-            Scalar loss
-        """
-        # Average absolute velocity over time
-        energy = torch.abs(v_trajectory) ** self.energy_p  # [B, T+1, output_dim]
-        loss = self.lambda_energy * energy.mean()
-        return loss
-
-    def forward(
-        self,
-        predictions: torch.Tensor,
-        targets: torch.Tensor,
-        trajectory: Optional[Dict[str, torch.Tensor]] = None,
-        epoch: int = 0,
-        learned_mu: Optional[torch.Tensor] = None
-    ) -> Dict[str, torch.Tensor]:
-        """
-        Compute total loss and components.
-
-        Args:
-            predictions: Final predictions [batch_size, output_dim]
-            targets: Ground truth [batch_size, output_dim]
-            trajectory: Optional trajectory dict with 'x', 'v', and 'aux'
-            epoch: Current epoch for annealing
-            learned_mu: Learned equilibrium attractor (v2)
-
-        Returns:
-            Dictionary with total loss and components
-        """
-        losses = {}
-
-        # Main task loss
-        L_task = self.compute_L_task(predictions, targets)
-        losses['L_task'] = L_task
-
-        total_loss = L_task
-
-        # Auxiliary losses (require trajectory)
-        if trajectory is not None:
-            x_traj = trajectory['x']  # [B, T+1, output_dim]
-            v_traj = trajectory['v']  # [B, T+1, output_dim]
-
-            # Mean constraint loss (v2: with learned_mu and variance weighting)
-            x_final = x_traj[:, -1, :]  # [B, output_dim]
-            L_mean = self.compute_L_mean(x_final, epoch, learned_mu)
-            losses['L_mean'] = L_mean
-            total_loss = total_loss + L_mean
-
-            # Speed loss
-            L_speed = self.compute_L_speed(x_traj)
-            losses['L_speed'] = L_speed
-            total_loss = total_loss + L_speed
-
-            # Energy loss (reduced during exploration phase)
-            lambda_energy = self.exploration_lambda_energy if self.exploration_phase else self.lambda_energy
-            # Temporarily override for this computation
-            original_lambda_energy = self.lambda_energy
-            self.lambda_energy = lambda_energy
-            L_energy = self.compute_L_energy(v_traj)
-            self.lambda_energy = original_lambda_energy  # Restore
-            losses['L_energy'] = L_energy
-            total_loss = total_loss + L_energy
-
-        losses['total'] = total_loss
-
-        # Report current phase lambda values
-        if self.exploration_phase:
-            losses['lambda_mean'] = torch.tensor(self.exploration_lambda_mean)
-        else:
-            losses['lambda_mean'] = torch.tensor(self.get_lambda_mean(epoch))
-
-        return losses
-
-    def set_exploration_phase(self, is_exploration: bool):
-        """
-        Set the current training phase.
-
-        Args:
-            is_exploration: True for exploration phase, False for equilibrium phase
-        """
-        self.exploration_phase = is_exploration
-
-
-def compute_convergence_metrics(
-    x_trajectory: torch.Tensor,
-    target_value: float = 5.0,
-    epsilon: float = 0.1
-) -> Dict[str, float]:
-    """
-    Compute metrics about convergence behavior.
-
-    Args:
-        x_trajectory: Trajectory [batch_size, T+1, output_dim]
-        target_value: Target value for convergence
-        epsilon: Tolerance for "converged" check
-
-    Returns:
-        Dictionary with metrics:
-            - time_to_converge: Average time steps to reach epsilon-ball
-            - final_rmse: RMSE at final time step
-            - final_mean: Mean value at final time step
-            - final_std: Std dev at final time step
-            - fraction_converged: Fraction of samples within epsilon at end
-    """
-    batch_size, T_plus_1, output_dim = x_trajectory.shape
-    T = T_plus_1 - 1
-
-    # Final time step statistics
-    x_final = x_trajectory[:, -1, :]  # [B, output_dim]
-    final_rmse = torch.sqrt(((x_final - target_value) ** 2).mean()).item()
-    final_mean = x_final.mean().item()
-    final_std = x_final.std().item()
-
-    # Fraction converged at final step
-    is_converged = torch.abs(x_final - target_value) <= epsilon
-    fraction_converged = is_converged.float().mean().item()
-
-    # Time to converge (first time within epsilon-ball)
-    # [B, T, output_dim]
-    deviations = torch.abs(x_trajectory[:, 1:, :] - target_value)  # Skip initial state
-    within_epsilon = deviations <= epsilon  # [B, T, output_dim]
-
-    # For each sample, find first time it's converged (across all output dims)
-    within_epsilon_all = within_epsilon.all(dim=-1)  # [B, T]
-
-    # Find first True index for each batch element
-    time_to_converge_list = []
-    for b in range(batch_size):
-        converged_times = torch.where(within_epsilon_all[b])[0]
-        if len(converged_times) > 0:
-            time_to_converge_list.append(converged_times[0].item() + 1)  # +1 because we skipped initial
-        else:
-            time_to_converge_list.append(T)  # Never converged
-
-    avg_time_to_converge = sum(time_to_converge_list) / len(time_to_converge_list)
-
-    return {
-        'time_to_converge': avg_time_to_converge,
-        'final_rmse': final_rmse,
-        'final_mean': final_mean,
-        'final_std': final_std,
-        'fraction_converged': fraction_converged
-    }
+"""
+Adaptive Loss Functions for IntegratorNeuronLayer Training
+
+Implements:
+- L_task: Main task loss (MSE or CE)
+- L_mean: Soft constraint to encourage convergence towards target
+- L_speed: Penalizes slow convergence in early iterations
+- L_energy: Regularizes velocity to prevent wild oscillations
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, Optional
+
+
+class IntegratorLoss(nn.Module):
+    """
+    Combined loss function with adaptive weighting for curriculum learning.
+    """
+
+    def __init__(
+        self,
+        target_value: float = 5.0,
+        lambda_mean_init: float = 1.0,
+        lambda_speed: float = 0.1,
+        lambda_energy: float = 0.01,
+        energy_p: float = 2.0,
+        annealing_schedule: str = 'exponential',
+        annealing_factor: float = 0.1,
+        annealing_epochs: int = 100,
+        variance_weighted: bool = True,
+        exploration_phase: bool = False,
+        exploration_lambda_mean: float = 0.05,
+        exploration_lambda_energy: float = 0.001,
+        task_loss_type: str = 'mse'  # 'mse' for regression, 'ce' for classification (LM)
+    ):
+        """
+        Args:
+            target_value: Target value for convergence (default 5.0)
+            lambda_mean_init: Initial weight for L_mean (will be annealed)
+            lambda_speed: Weight for L_speed (convergence speed penalty)
+            lambda_energy: Weight for L_energy (velocity regularization)
+            energy_p: Power for energy loss (2.0 = L2, 1.0 = L1)
+            annealing_schedule: 'exponential' or 'linear'
+            annealing_factor: Target factor for lambda_mean after annealing
+            annealing_epochs: Number of epochs to anneal over
+            variance_weighted: Use variance-weighted regularization
+            exploration_phase: Current phase (equilibrium=False, exploration=True)
+            exploration_lambda_mean: Lambda mean during exploration phase
+            exploration_lambda_energy: Lambda energy during exploration phase
+            task_loss_type: 'mse' for regression, 'ce' for classification/language modeling
+        """
+        super().__init__()
+
+        self.target_value = target_value
+        self.lambda_mean_init = lambda_mean_init
+        self.lambda_speed = lambda_speed
+        self.lambda_energy = lambda_energy
+        self.energy_p = energy_p
+        self.annealing_schedule = annealing_schedule
+        self.annealing_factor = annealing_factor
+        self.annealing_epochs = annealing_epochs
+
+        # Phase control and variance weighting
+        self.variance_weighted = variance_weighted
+        self.exploration_phase = exploration_phase
+        self.exploration_lambda_mean = exploration_lambda_mean
+        self.exploration_lambda_energy = exploration_lambda_energy
+
+        # Task loss type: MSE for regression, CrossEntropy for classification (language models)
+        self.task_loss_type = task_loss_type
+        if task_loss_type == 'mse':
+            self.task_loss = nn.MSELoss()
+        elif task_loss_type == 'ce':
+            self.task_loss = nn.CrossEntropyLoss()
+        else:
+            raise ValueError(f"Unknown task_loss_type: {task_loss_type}. Use 'mse' or 'ce'.")
+
+    def get_lambda_mean(self, epoch: int) -> float:
+        """
+        Compute current lambda_mean based on annealing schedule.
+
+        Args:
+            epoch: Current training epoch
+
+        Returns:
+            Current lambda_mean value
+        """
+        if epoch >= self.annealing_epochs:
+            return self.lambda_mean_init * self.annealing_factor
+
+        progress = epoch / self.annealing_epochs
+
+        if self.annealing_schedule == 'exponential':
+            # Exponential decay: lambda_mean = init * (factor)^progress
+            lambda_mean = self.lambda_mean_init * (self.annealing_factor ** progress)
+        elif self.annealing_schedule == 'linear':
+            # Linear decay: lambda_mean = init * (1 - progress * (1 - factor))
+            lambda_mean = self.lambda_mean_init * (1 - progress * (1 - self.annealing_factor))
+        else:
+            raise ValueError(f"Unknown annealing schedule: {self.annealing_schedule}")
+
+        return lambda_mean
+
+    def compute_L_task(
+        self,
+        predictions: torch.Tensor,
+        targets: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Main task loss: MSE between final prediction and target.
+
+        Args:
+            predictions: Model predictions [batch_size, output_dim]
+            targets: Ground truth targets [batch_size, output_dim]
+
+        Returns:
+            Scalar loss
+        """
+        return self.task_loss(predictions, targets)
+
+    def compute_L_mean(
+        self,
+        x_final: torch.Tensor,
+        epoch: int,
+        learned_mu: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Mean constraint loss: encourages batch mean to be close to target.
+        Supports variance-weighted regularization and learned mu.
+
+        Args:
+            x_final: Final state x_T [batch_size, output_dim]
+            epoch: Current epoch for annealing
+            learned_mu: Learned equilibrium attractor, if None uses target_value
+
+        Returns:
+            Scalar loss
+        """
+        # Use exploration phase lambda if in exploration mode
+        if self.exploration_phase:
+            lambda_mean = self.exploration_lambda_mean
+        else:
+            lambda_mean = self.get_lambda_mean(epoch)
+
+        # Use learned mu if provided, otherwise fixed target
+        target = learned_mu if learned_mu is not None else self.target_value
+
+        # Variance-weighted regularization
+        if self.variance_weighted:
+            # Compute per-neuron variance across batch
+            x_var = torch.var(x_final, dim=0, keepdim=False)  # [output_dim]
+            # Weight inversely proportional to variance (stable neurons penalized less)
+            weights = 1.0 / (1.0 + x_var)  # [output_dim]
+            # Normalize weights
+            weights = weights / weights.sum() * weights.numel()
+            # Weighted penalty
+            deviations = (x_final - target) ** 2  # [batch_size, output_dim]
+            loss = lambda_mean * (weights * deviations.mean(dim=0)).mean()
+        else:
+            # Uniform weighting
+            batch_mean = x_final.mean(dim=0)  # [output_dim]
+            loss = lambda_mean * ((batch_mean - target) ** 2).mean()
+
+        return loss
+
+    def compute_L_speed(
+        self,
+        x_trajectory: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Speed loss: penalizes deviation from target in early iterations.
+
+        Uses weighted sum: w_t = exp(-t / tau) to prioritize early steps.
+
+        Args:
+            x_trajectory: Trajectory of states [batch_size, T+1, output_dim]
+
+        Returns:
+            Scalar loss
+        """
+        T = x_trajectory.shape[1] - 1  # Exclude initial state
+        if T == 0:
+            return torch.tensor(0.0, device=x_trajectory.device)
+
+        # Exponentially decaying weights: prioritize early iterations
+        tau = T / 3.0  # Decay constant
+        t_indices = torch.arange(1, T + 1, device=x_trajectory.device, dtype=torch.float32)
+        weights = torch.exp(-t_indices / tau)
+        weights = weights / weights.sum()  # Normalize
+
+        # Compute weighted deviation from target
+        deviations = torch.abs(x_trajectory[:, 1:, :] - self.target_value)  # [B, T, output_dim]
+        weighted_dev = (deviations * weights.view(1, -1, 1)).sum(dim=1)  # [B, output_dim]
+
+        loss = self.lambda_speed * weighted_dev.mean()
+        return loss
+
+    def compute_L_energy(
+        self,
+        v_trajectory: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Energy loss: regularizes velocity to prevent oscillations.
+
+        Args:
+            v_trajectory: Trajectory of velocities [batch_size, T+1, output_dim]
+
+        Returns:
+            Scalar loss
+        """
+        # Average absolute velocity over time
+        energy = torch.abs(v_trajectory) ** self.energy_p  # [B, T+1, output_dim]
+        loss = self.lambda_energy * energy.mean()
+        return loss
+
+    def forward(
+        self,
+        predictions: torch.Tensor,
+        targets: torch.Tensor,
+        trajectory: Optional[Dict[str, torch.Tensor]] = None,
+        epoch: int = 0,
+        learned_mu: Optional[torch.Tensor] = None
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Compute total loss and components.
+
+        Args:
+            predictions: Final predictions [batch_size, output_dim]
+            targets: Ground truth [batch_size, output_dim]
+            trajectory: Optional trajectory dict with 'x', 'v', and 'aux'
+            epoch: Current epoch for annealing
+            learned_mu: Learned equilibrium attractor (v2)
+
+        Returns:
+            Dictionary with total loss and components
+        """
+        losses = {}
+
+        # Main task loss
+        L_task = self.compute_L_task(predictions, targets)
+        losses['L_task'] = L_task
+
+        total_loss = L_task
+
+        # Auxiliary losses (require trajectory)
+        if trajectory is not None:
+            x_traj = trajectory['x']  # [B, T+1, output_dim]
+            v_traj = trajectory['v']  # [B, T+1, output_dim]
+
+            # Mean constraint loss (v2: with learned_mu and variance weighting)
+            x_final = x_traj[:, -1, :]  # [B, output_dim]
+            L_mean = self.compute_L_mean(x_final, epoch, learned_mu)
+            losses['L_mean'] = L_mean
+            total_loss = total_loss + L_mean
+
+            # Speed loss
+            L_speed = self.compute_L_speed(x_traj)
+            losses['L_speed'] = L_speed
+            total_loss = total_loss + L_speed
+
+            # Energy loss (reduced during exploration phase)
+            lambda_energy = self.exploration_lambda_energy if self.exploration_phase else self.lambda_energy
+            # Temporarily override for this computation
+            original_lambda_energy = self.lambda_energy
+            self.lambda_energy = lambda_energy
+            L_energy = self.compute_L_energy(v_traj)
+            self.lambda_energy = original_lambda_energy  # Restore
+            losses['L_energy'] = L_energy
+            total_loss = total_loss + L_energy
+
+        losses['total'] = total_loss
+
+        # Report current phase lambda values
+        if self.exploration_phase:
+            losses['lambda_mean'] = torch.tensor(self.exploration_lambda_mean)
+        else:
+            losses['lambda_mean'] = torch.tensor(self.get_lambda_mean(epoch))
+
+        return losses
+
+    def set_exploration_phase(self, is_exploration: bool):
+        """
+        Set the current training phase.
+
+        Args:
+            is_exploration: True for exploration phase, False for equilibrium phase
+        """
+        self.exploration_phase = is_exploration
+
+
+def compute_convergence_metrics(
+    x_trajectory: torch.Tensor,
+    target_value: float = 5.0,
+    epsilon: float = 0.1
+) -> Dict[str, float]:
+    """
+    Compute metrics about convergence behavior.
+
+    Args:
+        x_trajectory: Trajectory [batch_size, T+1, output_dim]
+        target_value: Target value for convergence
+        epsilon: Tolerance for "converged" check
+
+    Returns:
+        Dictionary with metrics:
+            - time_to_converge: Average time steps to reach epsilon-ball
+            - final_rmse: RMSE at final time step
+            - final_mean: Mean value at final time step
+            - final_std: Std dev at final time step
+            - fraction_converged: Fraction of samples within epsilon at end
+    """
+    batch_size, T_plus_1, output_dim = x_trajectory.shape
+    T = T_plus_1 - 1
+
+    # Final time step statistics
+    x_final = x_trajectory[:, -1, :]  # [B, output_dim]
+    final_rmse = torch.sqrt(((x_final - target_value) ** 2).mean()).item()
+    final_mean = x_final.mean().item()
+    final_std = x_final.std().item()
+
+    # Fraction converged at final step
+    is_converged = torch.abs(x_final - target_value) <= epsilon
+    fraction_converged = is_converged.float().mean().item()
+
+    # Time to converge (first time within epsilon-ball)
+    # [B, T, output_dim]
+    deviations = torch.abs(x_trajectory[:, 1:, :] - target_value)  # Skip initial state
+    within_epsilon = deviations <= epsilon  # [B, T, output_dim]
+
+    # For each sample, find first time it's converged (across all output dims)
+    within_epsilon_all = within_epsilon.all(dim=-1)  # [B, T]
+
+    # Find first True index for each batch element
+    time_to_converge_list = []
+    for b in range(batch_size):
+        converged_times = torch.where(within_epsilon_all[b])[0]
+        if len(converged_times) > 0:
+            time_to_converge_list.append(converged_times[0].item() + 1)  # +1 because we skipped initial
+        else:
+            time_to_converge_list.append(T)  # Never converged
+
+    avg_time_to_converge = sum(time_to_converge_list) / len(time_to_converge_list)
+
+    return {
+        'time_to_converge': avg_time_to_converge,
+        'final_rmse': final_rmse,
+        'final_mean': final_mean,
+        'final_std': final_std,
+        'fraction_converged': fraction_converged
+    }

inl_llm/core/integrator_neuron_layer.py

@@ -1,552 +1,552 @@
-"""
-IntegratorNeuronLayer (INL) - Learnable Dynamics Architecture
-
-This module implements a neural network layer with learnable integrator/velocity dynamics.
-Key features:
-- Initial convergence towards 5 (configurable target)
-- Learnable controller parameters (alpha, beta, gating)
-- Soft constraints allowing deviation when data requires it
-- Deterministic and fully differentiable
-"""
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import math
-from typing import Optional, Tuple, Dict, Any
-
-# Optional: safetensors support for fast/secure model saving
-try:
-    from safetensors.torch import save_file, load_file
-    SAFETENSORS_AVAILABLE = True
-except ImportError:
-    SAFETENSORS_AVAILABLE = False
-
-
-class IntegratorNeuronLayer(nn.Module):
-    """
-    Implements learnable integrator dynamics with velocity control.
-
-    Equations:
-        error = x_t - mu
-        alpha = alpha_base * exp(-kappa * ||error||)  [if dynamic_alpha=True]
-        v_{t+1} = alpha * v_t + (1 - alpha) * v_cand - beta * error + harmonic_noise
-        x_{t+1} = x_t + (dt * velocity_scale) * g * v_{t+1}
-
-    where alpha_base, beta, g, v_cand are context-dependent learnable parameters
-    computed by a fused MLP controller from inputs [h, x, v].
-    """
-
-    def __init__(
-        self,
-        hidden_dim: int,
-        output_dim: int = 1,
-        target_value: float = 5.0,
-        dt: float = 0.1,
-        hidden_controller: int = 64,
-        init_alpha: float = 0.8,
-        init_beta: float = 0.5,
-        init_gate: float = 0.5,
-        velocity_scale: float = 1.0,
-        excitation_amplitude: float = 0.03,
-        learnable_mu: bool = True,
-        dynamic_alpha: bool = True,
-        alpha_kappa: float = 1.0
-    ):
-        """
-        Args:
-            hidden_dim: Dimension of context embedding h_t
-            output_dim: Dimension of state x (typically 1 for scalar prediction)
-            target_value: Initial target value (default 5.0)
-            dt: Time step for integration
-            hidden_controller: Hidden size for controller MLPs
-            init_alpha: Initial inertia coefficient
-            init_beta: Initial correction coefficient
-            init_gate: Initial gating value
-            velocity_scale: Scale factor for velocity
-            excitation_amplitude: Amplitude of deterministic harmonic noise
-            learnable_mu: Use learnable equilibrium attractor
-            dynamic_alpha: Use dynamic integration gain (α-control)
-            alpha_kappa: Sensitivity parameter for dynamic alpha
-        """
-        super().__init__()
-
-        # Validate hyperparameters
-        if hidden_dim <= 0:
-            raise ValueError(f"hidden_dim must be positive, got {hidden_dim}")
-        if output_dim <= 0:
-            raise ValueError(f"output_dim must be positive, got {output_dim}")
-        if dt <= 0:
-            raise ValueError(f"dt must be positive, got {dt}")
-        if hidden_controller <= 0:
-            raise ValueError(f"hidden_controller must be positive, got {hidden_controller}")
-        if not 0 <= init_alpha <= 1:
-            raise ValueError(f"init_alpha must be in [0, 1], got {init_alpha}")
-        if init_beta < 0:
-            raise ValueError(f"init_beta must be non-negative, got {init_beta}")
-        if not 0 <= init_gate <= 1:
-            raise ValueError(f"init_gate must be in [0, 1], got {init_gate}")
-        if velocity_scale <= 0:
-            raise ValueError(f"velocity_scale must be positive, got {velocity_scale}")
-        if excitation_amplitude < 0:
-            raise ValueError(f"excitation_amplitude must be non-negative, got {excitation_amplitude}")
-        if alpha_kappa < 0:
-            raise ValueError(f"alpha_kappa must be non-negative, got {alpha_kappa}")
-
-        self.hidden_dim = hidden_dim
-        self.output_dim = output_dim
-        self.dt = dt
-        self.velocity_scale = velocity_scale
-        self.dynamic_alpha = dynamic_alpha
-        self.alpha_kappa = alpha_kappa
-
-        # Pre-compute constant for performance
-        self._dt_velocity_scale = dt * velocity_scale
-
-        # Learnable equilibrium attractor
-        if learnable_mu:
-            self.mu = nn.Parameter(torch.full((output_dim,), target_value))
-            self.learnable_mu = True
-        else:
-            self.register_buffer('mu', torch.full((output_dim,), target_value))
-            self.learnable_mu = False
-
-        # Deterministic harmonic excitation
-        # Store as buffer so it can be modified dynamically (e.g., by scheduler)
-        self.register_buffer('excitation_amplitude', torch.tensor(excitation_amplitude, dtype=torch.float32))
-        # Learnable frequency and phase per dimension (deterministic initialization)
-        # Use deterministic initialization for reproducibility
-        gen = torch.Generator()
-        gen.manual_seed(42)  # Fixed seed for reproducibility
-        self.excitation_gamma = nn.Parameter(torch.randn(output_dim, generator=gen) * 0.1 + 1.0)
-        self.excitation_phi = nn.Parameter(torch.randn(output_dim, generator=gen) * 2 * math.pi)
-
-        # Fused controller MLP - outputs all 4 parameters at once for GPU efficiency
-        # Uses 3 separate inputs to avoid concat overhead
-        # Input: h (hidden_dim), x (output_dim), v (output_dim)
-        self.controller_h = nn.Linear(hidden_dim, hidden_controller)
-        self.controller_x = nn.Linear(output_dim, hidden_controller)
-        self.controller_v = nn.Linear(output_dim, hidden_controller)
-        self.controller_mlp = nn.Sequential(
-            nn.ReLU(),
-            nn.Linear(hidden_controller, 4 * output_dim),  # 4x output for all params
-        )
-
-        # Store output_dim for splitting
-        self._controller_output_dim = output_dim
-
-        # Initialize controller input layers
-        with torch.no_grad():
-            nn.init.xavier_uniform_(self.controller_h.weight)
-            nn.init.xavier_uniform_(self.controller_x.weight)
-            nn.init.xavier_uniform_(self.controller_v.weight)
-            self.controller_h.bias.zero_()
-            self.controller_x.bias.zero_()
-            self.controller_v.bias.zero_()
-
-            # Initialize output layer to produce desired initial values
-            bias = self.controller_mlp[-1].bias
-            alpha_bias = bias[0*output_dim:1*output_dim]
-            beta_bias = bias[1*output_dim:2*output_dim]
-            gate_bias = bias[2*output_dim:3*output_dim]
-            v_cand_bias = bias[3*output_dim:4*output_dim]
-
-            alpha_bias.fill_(self._inverse_sigmoid(init_alpha))
-            beta_bias.fill_(self._inverse_softplus(init_beta))
-            gate_bias.fill_(self._inverse_sigmoid(init_gate))
-            v_cand_bias.fill_(0.0)
-
-            # Small random initialization for symmetry breaking
-            self.controller_mlp[-1].weight.normal_(0.0, 0.01)
-
-    @staticmethod
-    def _inverse_sigmoid(y: float) -> float:
-        """Inverse of sigmoid function for initialization."""
-        y = max(min(y, 0.999), 0.001)  # Clamp to avoid inf
-        return torch.tensor(y / (1 - y)).log().item()
-
-    @staticmethod
-    def _inverse_softplus(y: float) -> float:
-        """Inverse of softplus function for initialization."""
-        y = max(y, 0.001)
-        return torch.tensor(y).expm1().log().item()
-
-    def forward(
-        self,
-        h: torch.Tensor,
-        x: torch.Tensor,
-        v: torch.Tensor,
-        step: int = 0,
-        return_aux: bool = True
-    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Dict[str, torch.Tensor]]]:
-        """
-        Forward pass computing one integration step.
-
-        Args:
-            h: Context embedding [batch_size, hidden_dim]
-            x: Current state [batch_size, output_dim]
-            v: Current velocity [batch_size, output_dim]
-            step: Current iteration step for deterministic excitation
-            return_aux: If False, skip creating aux dict (performance optimization)
-
-        Returns:
-            x_next: Next state [batch_size, output_dim]
-            v_next: Next velocity [batch_size, output_dim]
-            aux: Dictionary with controller parameters for monitoring (None if return_aux=False)
-        """
-        # Process inputs separately then sum (avoids concat overhead)
-        # Fuse additions for better performance
-        controller_hidden = self.controller_h(h)
-        controller_hidden = controller_hidden + self.controller_x(x)
-        controller_hidden = controller_hidden + self.controller_v(v)
-
-        # Compute all controller parameters in one forward pass (GPU efficient)
-        controller_output = self.controller_mlp(controller_hidden)
-
-        # Split into individual parameters using torch.split (more efficient than slicing)
-        alpha_base_raw, beta_raw, gate_raw, v_cand = torch.split(
-            controller_output, self._controller_output_dim, dim=1
-        )
-
-        # Apply activations (fused when possible with inplace for memory efficiency)
-        alpha_base = torch.sigmoid(alpha_base_raw)
-        beta = F.softplus(beta_raw)
-        gate = torch.sigmoid(gate_raw)
-        # v_cand has no activation (linear output)
-
-        # Compute error once (used in both alpha and velocity update)
-        error = x - self.mu
-
-        # Dynamic integration gain (α-control)
-        if self.dynamic_alpha:
-            # Only compute when needed (avoid torch.where overhead)
-            imbalance = torch.norm(error, dim=-1, keepdim=True)
-            alpha = alpha_base * torch.exp(-self.alpha_kappa * imbalance)
-        else:
-            alpha = alpha_base
-
-        # Update velocity with error correction term
-        v_next = alpha * v + (1 - alpha) * v_cand - beta * error
-
-        # Add deterministic harmonic excitation (only if amplitude > 0)
-        if self.excitation_amplitude.item() > 0:
-            # Deterministic noise based on iteration step
-            t = float(step)
-            # harmonic_noise shape: [output_dim]
-            harmonic_noise = self.excitation_amplitude * torch.sin(
-                self.excitation_gamma * t + self.excitation_phi
-            )
-            # Broadcast to [batch_size, output_dim] - implicit broadcasting is efficient
-            v_next = v_next + harmonic_noise
-
-        # Update state with gated velocity (use pre-computed constant)
-        x_next = x + self._dt_velocity_scale * gate * v_next
-
-        # Return auxiliary info for monitoring/loss (only if requested)
-        if return_aux:
-            aux = {
-                'alpha': alpha,
-                'alpha_base': alpha_base,
-                'beta': beta,
-                'gate': gate,
-                'v_cand': v_cand,
-                'error': error,
-                'mu': self.mu
-            }
-        else:
-            aux = None
-
-        return x_next, v_next, aux
-
-    def init_state(self, batch_size: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Initialize state x and velocity v.
-
-        Args:
-            batch_size: Batch size
-            device: Device to create tensors on
-
-        Returns:
-            x0: Initial state [batch_size, output_dim] initialized to learned mu
-            v0: Initial velocity [batch_size, output_dim] initialized to 0
-        """
-        # Initialize to current learned equilibrium, ensure correct device
-        # Move to device before expand for efficiency
-        mu_on_device = self.mu.to(device)
-        x0 = mu_on_device.unsqueeze(0).expand(batch_size, -1)
-        v0 = torch.zeros((batch_size, self.output_dim), device=device)
-        return x0, v0
-
-    def reset_parameters(self) -> None:
-        """
-        Reset all learnable parameters to their initial values.
-        Standard PyTorch method for parameter reinitialization.
-        """
-        # Reset controller layers
-        nn.init.xavier_uniform_(self.controller_h.weight)
-        nn.init.xavier_uniform_(self.controller_x.weight)
-        nn.init.xavier_uniform_(self.controller_v.weight)
-        self.controller_h.bias.zero_()
-        self.controller_x.bias.zero_()
-        self.controller_v.bias.zero_()
-
-        # Reset output layer with proper initialization
-        output_dim = self._controller_output_dim
-        bias = self.controller_mlp[-1].bias
-
-        # Use stored init values if available, otherwise use defaults
-        init_alpha = getattr(self, '_init_alpha', 0.8)
-        init_beta = getattr(self, '_init_beta', 0.5)
-        init_gate = getattr(self, '_init_gate', 0.5)
-
-        bias[0*output_dim:1*output_dim].fill_(self._inverse_sigmoid(init_alpha))
-        bias[1*output_dim:2*output_dim].fill_(self._inverse_softplus(init_beta))
-        bias[2*output_dim:3*output_dim].fill_(self._inverse_sigmoid(init_gate))
-        bias[3*output_dim:4*output_dim].zero_()
-
-        self.controller_mlp[-1].weight.normal_(0.0, 0.01)
-
-        # Reset excitation parameters
-        with torch.no_grad():
-            gen = torch.Generator()
-            gen.manual_seed(42)
-            self.excitation_gamma.copy_(torch.randn(output_dim, generator=gen) * 0.1 + 1.0)
-            self.excitation_phi.copy_(torch.randn(output_dim, generator=gen) * 2 * math.pi)
-
-    def __repr__(self) -> str:
-        """String representation for debugging."""
-        # Use .item() for scalar tensors in repr (acceptable in non-critical path)
-        exc_amp = self.excitation_amplitude.item() if self.excitation_amplitude.numel() == 1 else self.excitation_amplitude
-        return (
-            f"{self.__class__.__name__}(\n"
-            f"  hidden_dim={self.hidden_dim}, output_dim={self.output_dim},\n"
-            f"  dt={self.dt}, velocity_scale={self.velocity_scale},\n"
-            f"  excitation_amplitude={exc_amp:.4f},\n"
-            f"  learnable_mu={self.learnable_mu}, dynamic_alpha={self.dynamic_alpha},\n"
-            f"  alpha_kappa={self.alpha_kappa}\n"
-            f")"
-        )
-
-
-class IntegratorModel(nn.Module):
-    """
-    Complete model: Backbone + IntegratorNeuronLayer + Readout
-    """
-
-    def __init__(
-        self,
-        input_dim: int,
-        hidden_dim: int = 128,
-        num_layers: int = 2,
-        num_iterations: int = 10,
-        output_dim: int = 1,
-        target_value: float = 5.0,
-        **inl_kwargs: Any
-    ):
-        """
-        Args:
-            input_dim: Input feature dimension
-            hidden_dim: Hidden dimension for backbone and INL
-            num_layers: Number of layers in backbone MLP
-            num_iterations: Number of integration steps T
-            output_dim: Output dimension (1 for scalar regression)
-            target_value: Target value for convergence (default 5.0)
-            **inl_kwargs: Additional arguments for IntegratorNeuronLayer
-        """
-        super().__init__()
-
-        # Validate hyperparameters
-        if input_dim <= 0:
-            raise ValueError(f"input_dim must be positive, got {input_dim}")
-        if hidden_dim <= 0:
-            raise ValueError(f"hidden_dim must be positive, got {hidden_dim}")
-        if num_layers <= 0:
-            raise ValueError(f"num_layers must be positive, got {num_layers}")
-        if num_iterations <= 0:
-            raise ValueError(f"num_iterations must be positive, got {num_iterations}")
-        if output_dim <= 0:
-            raise ValueError(f"output_dim must be positive, got {output_dim}")
-
-        self.input_dim = input_dim
-        self.hidden_dim = hidden_dim
-        self.num_iterations = num_iterations
-        self.output_dim = output_dim
-
-        # Backbone: simple MLP (can be replaced with Transformer)
-        layers = []
-        current_dim = input_dim
-        for _ in range(num_layers):
-            layers.extend([
-                nn.Linear(current_dim, hidden_dim),
-                nn.ReLU(),
-                nn.LayerNorm(hidden_dim)
-            ])
-            current_dim = hidden_dim
-        self.backbone = nn.Sequential(*layers)
-
-        # Integrator Neuron Layer
-        self.inl = IntegratorNeuronLayer(
-            hidden_dim=hidden_dim,
-            output_dim=output_dim,
-            target_value=target_value,
-            **inl_kwargs
-        )
-
-        # Readout layer
-        self.readout = nn.Linear(output_dim, output_dim)
-        # Initialize readout to identity transformation (no bias shift)
-        # Since x is already initialized to target_value, we just pass it through
-        with torch.no_grad():
-            # Only set diagonal if square matrix
-            if self.readout.weight.shape[0] == self.readout.weight.shape[1]:
-                self.readout.weight.fill_(0.0)
-                self.readout.weight.diagonal().fill_(1.0)
-            else:
-                # For non-square, use Xavier/Glorot initialization
-                nn.init.xavier_uniform_(self.readout.weight)
-            self.readout.bias.fill_(0.0)  # No bias - x already at target_value
-
-    def _run_dynamics(
-        self,
-        inputs: torch.Tensor,
-        return_trajectory: bool = False
-    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Dict[str, torch.Tensor]]]:
-        """
-        Internal method to run INL dynamics.
-
-        Args:
-            inputs: Input features [batch_size, input_dim]
-            return_trajectory: If True, return full trajectory and aux info
-
-        Returns:
-            x: Final state [batch_size, output_dim]
-            v: Final velocity [batch_size, output_dim]
-            trajectory: Optional dict with trajectory info if return_trajectory=True
-        """
-        batch_size = inputs.shape[0]
-        device = inputs.device
-
-        # Compute context from backbone
-        h = self.backbone(inputs)  # [B, hidden_dim]
-
-        # Initialize state and velocity
-        x, v = self.inl.init_state(batch_size, device)
-
-        # Store trajectory if requested (pre-allocate for efficiency)
-        if return_trajectory:
-            # Pre-allocate tensors with empty (no initialization overhead)
-            x_traj = torch.empty(batch_size, self.num_iterations + 1, self.output_dim, device=device)
-            v_traj = torch.empty(batch_size, self.num_iterations + 1, self.output_dim, device=device)
-            x_traj[:, 0] = x
-            v_traj[:, 0] = v
-            # For aux, we still need a list (dict values vary)
-            aux_traj = []
-
-        # Run integration steps
-        for t in range(self.num_iterations):
-            # Skip aux creation if not needed (performance)
-            x, v, aux = self.inl(h, x, v, step=t, return_aux=return_trajectory)
-
-            if return_trajectory:
-                # Store directly in pre-allocated tensors (no detach needed, done at the end)
-                x_traj[:, t + 1] = x
-                v_traj[:, t + 1] = v
-                # Only store essential aux info (skip redundant fields)
-                aux_traj.append({
-                    'alpha': aux['alpha'].detach(),
-                    'beta': aux['beta'].detach(),
-                    'error': aux['error'].detach()
-                })
-
-        if return_trajectory:
-            trajectory = {
-                'x': x_traj.detach(),  # Already stacked, just detach
-                'v': v_traj.detach(),  # Already stacked, just detach
-                'aux': aux_traj
-            }
-            return x, v, trajectory
-
-        return x, v, None
-
-    def forward(
-        self,
-        inputs: torch.Tensor,
-        return_trajectory: bool = False
-    ) -> Tuple[torch.Tensor, Optional[Dict[str, torch.Tensor]]]:
-        """
-        Forward pass through complete model.
-
-        Args:
-            inputs: Input features [batch_size, input_dim]
-            return_trajectory: If True, return full trajectory and aux info
-
-        Returns:
-            output: Final prediction [batch_size, output_dim]
-            trajectory: Optional dict with trajectory info if return_trajectory=True
-        """
-        x, v, trajectory = self._run_dynamics(inputs, return_trajectory)
-        output = self.readout(x)
-
-        if return_trajectory:
-            return output, trajectory
-
-        return output, None
-
-    def get_final_state(self, inputs: torch.Tensor) -> torch.Tensor:
-        """Get final state x_T before readout."""
-        x, _, _ = self._run_dynamics(inputs, return_trajectory=False)
-        return x
-
-    def get_learned_mu(self) -> Optional[torch.Tensor]:
-        """
-        Get the learned equilibrium attractor.
-
-        Returns:
-            Learned mu tensor if learnable_mu enabled, else None
-        """
-        if hasattr(self.inl, 'learnable_mu') and self.inl.learnable_mu:
-            return self.inl.mu
-        return None
-
-    def save_safetensors(self, path: str) -> None:
-        """
-        Save model state dict using safetensors format.
-
-        Args:
-            path: Path to save file (e.g., 'model.safetensors')
-
-        Requires: pip install safetensors
-        """
-        if not SAFETENSORS_AVAILABLE:
-            raise ImportError(
-                "safetensors not installed. Install with: pip install safetensors"
-            )
-
-        save_file(self.state_dict(), path)
-
-    def load_safetensors(self, path: str, strict: bool = True) -> None:
-        """
-        Load model state dict from safetensors format.
-
-        Args:
-            path: Path to safetensors file
-            strict: Whether to strictly enforce matching keys
-
-        Requires: pip install safetensors
-        """
-        if not SAFETENSORS_AVAILABLE:
-            raise ImportError(
-                "safetensors not installed. Install with: pip install safetensors"
-            )
-
-        state_dict = load_file(path)
-        self.load_state_dict(state_dict, strict=strict)
-
-    def __repr__(self) -> str:
-        """String representation for debugging."""
-        return (
-            f"{self.__class__.__name__}(\n"
-            f"  input_dim={self.input_dim}, hidden_dim={self.hidden_dim},\n"
-            f"  output_dim={self.output_dim}, num_iterations={self.num_iterations}\n"
-            f")"
-        )
+"""
+IntegratorNeuronLayer (INL) - Learnable Dynamics Architecture
+
+This module implements a neural network layer with learnable integrator/velocity dynamics.
+Key features:
+- Initial convergence towards 5 (configurable target)
+- Learnable controller parameters (alpha, beta, gating)
+- Soft constraints allowing deviation when data requires it
+- Deterministic and fully differentiable
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Optional, Tuple, Dict, Any
+
+# Optional: safetensors support for fast/secure model saving
+try:
+    from safetensors.torch import save_file, load_file
+    SAFETENSORS_AVAILABLE = True
+except ImportError:
+    SAFETENSORS_AVAILABLE = False
+
+
+class IntegratorNeuronLayer(nn.Module):
+    """
+    Implements learnable integrator dynamics with velocity control.
+
+    Equations:
+        error = x_t - mu
+        alpha = alpha_base * exp(-kappa * ||error||)  [if dynamic_alpha=True]
+        v_{t+1} = alpha * v_t + (1 - alpha) * v_cand - beta * error + harmonic_noise
+        x_{t+1} = x_t + (dt * velocity_scale) * g * v_{t+1}
+
+    where alpha_base, beta, g, v_cand are context-dependent learnable parameters
+    computed by a fused MLP controller from inputs [h, x, v].
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        output_dim: int = 1,
+        target_value: float = 5.0,
+        dt: float = 0.1,
+        hidden_controller: int = 64,
+        init_alpha: float = 0.8,
+        init_beta: float = 0.5,
+        init_gate: float = 0.5,
+        velocity_scale: float = 1.0,
+        excitation_amplitude: float = 0.03,
+        learnable_mu: bool = True,
+        dynamic_alpha: bool = True,
+        alpha_kappa: float = 1.0
+    ):
+        """
+        Args:
+            hidden_dim: Dimension of context embedding h_t
+            output_dim: Dimension of state x (typically 1 for scalar prediction)
+            target_value: Initial target value (default 5.0)
+            dt: Time step for integration
+            hidden_controller: Hidden size for controller MLPs
+            init_alpha: Initial inertia coefficient
+            init_beta: Initial correction coefficient
+            init_gate: Initial gating value
+            velocity_scale: Scale factor for velocity
+            excitation_amplitude: Amplitude of deterministic harmonic noise
+            learnable_mu: Use learnable equilibrium attractor
+            dynamic_alpha: Use dynamic integration gain (α-control)
+            alpha_kappa: Sensitivity parameter for dynamic alpha
+        """
+        super().__init__()
+
+        # Validate hyperparameters
+        if hidden_dim <= 0:
+            raise ValueError(f"hidden_dim must be positive, got {hidden_dim}")
+        if output_dim <= 0:
+            raise ValueError(f"output_dim must be positive, got {output_dim}")
+        if dt <= 0:
+            raise ValueError(f"dt must be positive, got {dt}")
+        if hidden_controller <= 0:
+            raise ValueError(f"hidden_controller must be positive, got {hidden_controller}")
+        if not 0 <= init_alpha <= 1:
+            raise ValueError(f"init_alpha must be in [0, 1], got {init_alpha}")
+        if init_beta < 0:
+            raise ValueError(f"init_beta must be non-negative, got {init_beta}")
+        if not 0 <= init_gate <= 1:
+            raise ValueError(f"init_gate must be in [0, 1], got {init_gate}")
+        if velocity_scale <= 0:
+            raise ValueError(f"velocity_scale must be positive, got {velocity_scale}")
+        if excitation_amplitude < 0:
+            raise ValueError(f"excitation_amplitude must be non-negative, got {excitation_amplitude}")
+        if alpha_kappa < 0:
+            raise ValueError(f"alpha_kappa must be non-negative, got {alpha_kappa}")
+
+        self.hidden_dim = hidden_dim
+        self.output_dim = output_dim
+        self.dt = dt
+        self.velocity_scale = velocity_scale
+        self.dynamic_alpha = dynamic_alpha
+        self.alpha_kappa = alpha_kappa
+
+        # Pre-compute constant for performance
+        self._dt_velocity_scale = dt * velocity_scale
+
+        # Learnable equilibrium attractor
+        if learnable_mu:
+            self.mu = nn.Parameter(torch.full((output_dim,), target_value))
+            self.learnable_mu = True
+        else:
+            self.register_buffer('mu', torch.full((output_dim,), target_value))
+            self.learnable_mu = False
+
+        # Deterministic harmonic excitation
+        # Store as buffer so it can be modified dynamically (e.g., by scheduler)
+        self.register_buffer('excitation_amplitude', torch.tensor(excitation_amplitude, dtype=torch.float32))
+        # Learnable frequency and phase per dimension (deterministic initialization)
+        # Use deterministic initialization for reproducibility
+        gen = torch.Generator()
+        gen.manual_seed(42)  # Fixed seed for reproducibility
+        self.excitation_gamma = nn.Parameter(torch.randn(output_dim, generator=gen) * 0.1 + 1.0)
+        self.excitation_phi = nn.Parameter(torch.randn(output_dim, generator=gen) * 2 * math.pi)
+
+        # Fused controller MLP - outputs all 4 parameters at once for GPU efficiency
+        # Uses 3 separate inputs to avoid concat overhead
+        # Input: h (hidden_dim), x (output_dim), v (output_dim)
+        self.controller_h = nn.Linear(hidden_dim, hidden_controller)
+        self.controller_x = nn.Linear(output_dim, hidden_controller)
+        self.controller_v = nn.Linear(output_dim, hidden_controller)
+        self.controller_mlp = nn.Sequential(
+            nn.ReLU(),
+            nn.Linear(hidden_controller, 4 * output_dim),  # 4x output for all params
+        )
+
+        # Store output_dim for splitting
+        self._controller_output_dim = output_dim
+
+        # Initialize controller input layers
+        with torch.no_grad():
+            nn.init.xavier_uniform_(self.controller_h.weight)
+            nn.init.xavier_uniform_(self.controller_x.weight)
+            nn.init.xavier_uniform_(self.controller_v.weight)
+            self.controller_h.bias.zero_()
+            self.controller_x.bias.zero_()
+            self.controller_v.bias.zero_()
+
+            # Initialize output layer to produce desired initial values
+            bias = self.controller_mlp[-1].bias
+            alpha_bias = bias[0*output_dim:1*output_dim]
+            beta_bias = bias[1*output_dim:2*output_dim]
+            gate_bias = bias[2*output_dim:3*output_dim]
+            v_cand_bias = bias[3*output_dim:4*output_dim]
+
+            alpha_bias.fill_(self._inverse_sigmoid(init_alpha))
+            beta_bias.fill_(self._inverse_softplus(init_beta))
+            gate_bias.fill_(self._inverse_sigmoid(init_gate))
+            v_cand_bias.fill_(0.0)
+
+            # Small random initialization for symmetry breaking
+            self.controller_mlp[-1].weight.normal_(0.0, 0.01)
+
+    @staticmethod
+    def _inverse_sigmoid(y: float) -> float:
+        """Inverse of sigmoid function for initialization."""
+        y = max(min(y, 0.999), 0.001)  # Clamp to avoid inf
+        return torch.tensor(y / (1 - y)).log().item()
+
+    @staticmethod
+    def _inverse_softplus(y: float) -> float:
+        """Inverse of softplus function for initialization."""
+        y = max(y, 0.001)
+        return torch.tensor(y).expm1().log().item()
+
+    def forward(
+        self,
+        h: torch.Tensor,
+        x: torch.Tensor,
+        v: torch.Tensor,
+        step: int = 0,
+        return_aux: bool = True
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Dict[str, torch.Tensor]]]:
+        """
+        Forward pass computing one integration step.
+
+        Args:
+            h: Context embedding [batch_size, hidden_dim]
+            x: Current state [batch_size, output_dim]
+            v: Current velocity [batch_size, output_dim]
+            step: Current iteration step for deterministic excitation
+            return_aux: If False, skip creating aux dict (performance optimization)
+
+        Returns:
+            x_next: Next state [batch_size, output_dim]
+            v_next: Next velocity [batch_size, output_dim]
+            aux: Dictionary with controller parameters for monitoring (None if return_aux=False)
+        """
+        # Process inputs separately then sum (avoids concat overhead)
+        # Fuse additions for better performance
+        controller_hidden = self.controller_h(h)
+        controller_hidden = controller_hidden + self.controller_x(x)
+        controller_hidden = controller_hidden + self.controller_v(v)
+
+        # Compute all controller parameters in one forward pass (GPU efficient)
+        controller_output = self.controller_mlp(controller_hidden)
+
+        # Split into individual parameters using torch.split (more efficient than slicing)
+        alpha_base_raw, beta_raw, gate_raw, v_cand = torch.split(
+            controller_output, self._controller_output_dim, dim=1
+        )
+
+        # Apply activations (fused when possible with inplace for memory efficiency)
+        alpha_base = torch.sigmoid(alpha_base_raw)
+        beta = F.softplus(beta_raw)
+        gate = torch.sigmoid(gate_raw)
+        # v_cand has no activation (linear output)
+
+        # Compute error once (used in both alpha and velocity update)
+        error = x - self.mu
+
+        # Dynamic integration gain (α-control)
+        if self.dynamic_alpha:
+            # Only compute when needed (avoid torch.where overhead)
+            imbalance = torch.norm(error, dim=-1, keepdim=True)
+            alpha = alpha_base * torch.exp(-self.alpha_kappa * imbalance)
+        else:
+            alpha = alpha_base
+
+        # Update velocity with error correction term
+        v_next = alpha * v + (1 - alpha) * v_cand - beta * error
+
+        # Add deterministic harmonic excitation (only if amplitude > 0)
+        if self.excitation_amplitude.item() > 0:
+            # Deterministic noise based on iteration step
+            t = float(step)
+            # harmonic_noise shape: [output_dim]
+            harmonic_noise = self.excitation_amplitude * torch.sin(
+                self.excitation_gamma * t + self.excitation_phi
+            )
+            # Broadcast to [batch_size, output_dim] - implicit broadcasting is efficient
+            v_next = v_next + harmonic_noise
+
+        # Update state with gated velocity (use pre-computed constant)
+        x_next = x + self._dt_velocity_scale * gate * v_next
+
+        # Return auxiliary info for monitoring/loss (only if requested)
+        if return_aux:
+            aux = {
+                'alpha': alpha,
+                'alpha_base': alpha_base,
+                'beta': beta,
+                'gate': gate,
+                'v_cand': v_cand,
+                'error': error,
+                'mu': self.mu
+            }
+        else:
+            aux = None
+
+        return x_next, v_next, aux
+
+    def init_state(self, batch_size: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Initialize state x and velocity v.
+
+        Args:
+            batch_size: Batch size
+            device: Device to create tensors on
+
+        Returns:
+            x0: Initial state [batch_size, output_dim] initialized to learned mu
+            v0: Initial velocity [batch_size, output_dim] initialized to 0
+        """
+        # Initialize to current learned equilibrium, ensure correct device
+        # Move to device before expand for efficiency
+        mu_on_device = self.mu.to(device)
+        x0 = mu_on_device.unsqueeze(0).expand(batch_size, -1)
+        v0 = torch.zeros((batch_size, self.output_dim), device=device)
+        return x0, v0
+
+    def reset_parameters(self) -> None:
+        """
+        Reset all learnable parameters to their initial values.
+        Standard PyTorch method for parameter reinitialization.
+        """
+        # Reset controller layers
+        nn.init.xavier_uniform_(self.controller_h.weight)
+        nn.init.xavier_uniform_(self.controller_x.weight)
+        nn.init.xavier_uniform_(self.controller_v.weight)
+        self.controller_h.bias.zero_()
+        self.controller_x.bias.zero_()
+        self.controller_v.bias.zero_()
+
+        # Reset output layer with proper initialization
+        output_dim = self._controller_output_dim
+        bias = self.controller_mlp[-1].bias
+
+        # Use stored init values if available, otherwise use defaults
+        init_alpha = getattr(self, '_init_alpha', 0.8)
+        init_beta = getattr(self, '_init_beta', 0.5)
+        init_gate = getattr(self, '_init_gate', 0.5)
+
+        bias[0*output_dim:1*output_dim].fill_(self._inverse_sigmoid(init_alpha))
+        bias[1*output_dim:2*output_dim].fill_(self._inverse_softplus(init_beta))
+        bias[2*output_dim:3*output_dim].fill_(self._inverse_sigmoid(init_gate))
+        bias[3*output_dim:4*output_dim].zero_()
+
+        self.controller_mlp[-1].weight.normal_(0.0, 0.01)
+
+        # Reset excitation parameters
+        with torch.no_grad():
+            gen = torch.Generator()
+            gen.manual_seed(42)
+            self.excitation_gamma.copy_(torch.randn(output_dim, generator=gen) * 0.1 + 1.0)
+            self.excitation_phi.copy_(torch.randn(output_dim, generator=gen) * 2 * math.pi)
+
+    def __repr__(self) -> str:
+        """String representation for debugging."""
+        # Use .item() for scalar tensors in repr (acceptable in non-critical path)
+        exc_amp = self.excitation_amplitude.item() if self.excitation_amplitude.numel() == 1 else self.excitation_amplitude
+        return (
+            f"{self.__class__.__name__}(\n"
+            f"  hidden_dim={self.hidden_dim}, output_dim={self.output_dim},\n"
+            f"  dt={self.dt}, velocity_scale={self.velocity_scale},\n"
+            f"  excitation_amplitude={exc_amp:.4f},\n"
+            f"  learnable_mu={self.learnable_mu}, dynamic_alpha={self.dynamic_alpha},\n"
+            f"  alpha_kappa={self.alpha_kappa}\n"
+            f")"
+        )
+
+
+class IntegratorModel(nn.Module):
+    """
+    Complete model: Backbone + IntegratorNeuronLayer + Readout
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int = 128,
+        num_layers: int = 2,
+        num_iterations: int = 10,
+        output_dim: int = 1,
+        target_value: float = 5.0,
+        **inl_kwargs: Any
+    ):
+        """
+        Args:
+            input_dim: Input feature dimension
+            hidden_dim: Hidden dimension for backbone and INL
+            num_layers: Number of layers in backbone MLP
+            num_iterations: Number of integration steps T
+            output_dim: Output dimension (1 for scalar regression)
+            target_value: Target value for convergence (default 5.0)
+            **inl_kwargs: Additional arguments for IntegratorNeuronLayer
+        """
+        super().__init__()
+
+        # Validate hyperparameters
+        if input_dim <= 0:
+            raise ValueError(f"input_dim must be positive, got {input_dim}")
+        if hidden_dim <= 0:
+            raise ValueError(f"hidden_dim must be positive, got {hidden_dim}")
+        if num_layers <= 0:
+            raise ValueError(f"num_layers must be positive, got {num_layers}")
+        if num_iterations <= 0:
+            raise ValueError(f"num_iterations must be positive, got {num_iterations}")
+        if output_dim <= 0:
+            raise ValueError(f"output_dim must be positive, got {output_dim}")
+
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.num_iterations = num_iterations
+        self.output_dim = output_dim
+
+        # Backbone: simple MLP (can be replaced with Transformer)
+        layers = []
+        current_dim = input_dim
+        for _ in range(num_layers):
+            layers.extend([
+                nn.Linear(current_dim, hidden_dim),
+                nn.ReLU(),
+                nn.LayerNorm(hidden_dim)
+            ])
+            current_dim = hidden_dim
+        self.backbone = nn.Sequential(*layers)
+
+        # Integrator Neuron Layer
+        self.inl = IntegratorNeuronLayer(
+            hidden_dim=hidden_dim,
+            output_dim=output_dim,
+            target_value=target_value,
+            **inl_kwargs
+        )
+
+        # Readout layer
+        self.readout = nn.Linear(output_dim, output_dim)
+        # Initialize readout to identity transformation (no bias shift)
+        # Since x is already initialized to target_value, we just pass it through
+        with torch.no_grad():
+            # Only set diagonal if square matrix
+            if self.readout.weight.shape[0] == self.readout.weight.shape[1]:
+                self.readout.weight.fill_(0.0)
+                self.readout.weight.diagonal().fill_(1.0)
+            else:
+                # For non-square, use Xavier/Glorot initialization
+                nn.init.xavier_uniform_(self.readout.weight)
+            self.readout.bias.fill_(0.0)  # No bias - x already at target_value
+
+    def _run_dynamics(
+        self,
+        inputs: torch.Tensor,
+        return_trajectory: bool = False
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Dict[str, torch.Tensor]]]:
+        """
+        Internal method to run INL dynamics.
+
+        Args:
+            inputs: Input features [batch_size, input_dim]
+            return_trajectory: If True, return full trajectory and aux info
+
+        Returns:
+            x: Final state [batch_size, output_dim]
+            v: Final velocity [batch_size, output_dim]
+            trajectory: Optional dict with trajectory info if return_trajectory=True
+        """
+        batch_size = inputs.shape[0]
+        device = inputs.device
+
+        # Compute context from backbone
+        h = self.backbone(inputs)  # [B, hidden_dim]
+
+        # Initialize state and velocity
+        x, v = self.inl.init_state(batch_size, device)
+
+        # Store trajectory if requested (pre-allocate for efficiency)
+        if return_trajectory:
+            # Pre-allocate tensors with empty (no initialization overhead)
+            x_traj = torch.empty(batch_size, self.num_iterations + 1, self.output_dim, device=device)
+            v_traj = torch.empty(batch_size, self.num_iterations + 1, self.output_dim, device=device)
+            x_traj[:, 0] = x
+            v_traj[:, 0] = v
+            # For aux, we still need a list (dict values vary)
+            aux_traj = []
+
+        # Run integration steps
+        for t in range(self.num_iterations):
+            # Skip aux creation if not needed (performance)
+            x, v, aux = self.inl(h, x, v, step=t, return_aux=return_trajectory)
+
+            if return_trajectory:
+                # Store directly in pre-allocated tensors (no detach needed, done at the end)
+                x_traj[:, t + 1] = x
+                v_traj[:, t + 1] = v
+                # Only store essential aux info (skip redundant fields)
+                aux_traj.append({
+                    'alpha': aux['alpha'].detach(),
+                    'beta': aux['beta'].detach(),
+                    'error': aux['error'].detach()
+                })
+
+        if return_trajectory:
+            trajectory = {
+                'x': x_traj.detach(),  # Already stacked, just detach
+                'v': v_traj.detach(),  # Already stacked, just detach
+                'aux': aux_traj
+            }
+            return x, v, trajectory
+
+        return x, v, None
+
+    def forward(
+        self,
+        inputs: torch.Tensor,
+        return_trajectory: bool = False
+    ) -> Tuple[torch.Tensor, Optional[Dict[str, torch.Tensor]]]:
+        """
+        Forward pass through complete model.
+
+        Args:
+            inputs: Input features [batch_size, input_dim]
+            return_trajectory: If True, return full trajectory and aux info
+
+        Returns:
+            output: Final prediction [batch_size, output_dim]
+            trajectory: Optional dict with trajectory info if return_trajectory=True
+        """
+        x, v, trajectory = self._run_dynamics(inputs, return_trajectory)
+        output = self.readout(x)
+
+        if return_trajectory:
+            return output, trajectory
+
+        return output, None
+
+    def get_final_state(self, inputs: torch.Tensor) -> torch.Tensor:
+        """Get final state x_T before readout."""
+        x, _, _ = self._run_dynamics(inputs, return_trajectory=False)
+        return x
+
+    def get_learned_mu(self) -> Optional[torch.Tensor]:
+        """
+        Get the learned equilibrium attractor.
+
+        Returns:
+            Learned mu tensor if learnable_mu enabled, else None
+        """
+        if hasattr(self.inl, 'learnable_mu') and self.inl.learnable_mu:
+            return self.inl.mu
+        return None
+
+    def save_safetensors(self, path: str) -> None:
+        """
+        Save model state dict using safetensors format.
+
+        Args:
+            path: Path to save file (e.g., 'model.safetensors')
+
+        Requires: pip install safetensors
+        """
+        if not SAFETENSORS_AVAILABLE:
+            raise ImportError(
+                "safetensors not installed. Install with: pip install safetensors"
+            )
+
+        save_file(self.state_dict(), path)
+
+    def load_safetensors(self, path: str, strict: bool = True) -> None:
+        """
+        Load model state dict from safetensors format.
+
+        Args:
+            path: Path to safetensors file
+            strict: Whether to strictly enforce matching keys
+
+        Requires: pip install safetensors
+        """
+        if not SAFETENSORS_AVAILABLE:
+            raise ImportError(
+                "safetensors not installed. Install with: pip install safetensors"
+            )
+
+        state_dict = load_file(path)
+        self.load_state_dict(state_dict, strict=strict)
+
+    def __repr__(self) -> str:
+        """String representation for debugging."""
+        return (
+            f"{self.__class__.__name__}(\n"
+            f"  input_dim={self.input_dim}, hidden_dim={self.hidden_dim},\n"
+            f"  output_dim={self.output_dim}, num_iterations={self.num_iterations}\n"
+            f")"
+        )

inl_llm/core/integrator_scheduler_v2.py

@@ -1,426 +1,426 @@
-"""
-INL-LLM: Equilibrium-Exploration Cycle Scheduler
-
-Implements rhythmic training phases that alternate between:
-- Equilibrium Phase: Strong stability constraint, low excitation (stabilization)
-- Exploration Phase: Weak stability constraint, high excitation (discovery)
-
-This deterministic cycling encourages structured exploration without randomness.
-"""
-
-from typing import Dict, NamedTuple
-import torch
-from copy import deepcopy
-
-
-class PhaseConfig(NamedTuple):
-    """Configuration for a training phase."""
-    name: str
-    lambda_mean: float
-    excitation_amplitude: float
-    duration_epochs: int
-
-
-class EquilibriumExplorationScheduler:
-    """
-    Manages equilibrium-exploration cycles for v2 training.
-
-    Example cycle:
-    - Equilibrium (10 epochs): lambda_mean=0.5, excitation=0.0
-    - Exploration (20 epochs): lambda_mean=0.05, excitation=0.05
-    - Repeat...
-    """
-
-    def __init__(
-        self,
-        equilibrium_config: Dict = None,
-        exploration_config: Dict = None,
-        num_cycles: int = 5,
-        warmup_epochs: int = 10
-    ):
-        """
-        Args:
-            equilibrium_config: Config for equilibrium phase
-            exploration_config: Config for exploration phase
-            num_cycles: Number of complete cycles to perform
-            warmup_epochs: Initial warmup before cycling starts
-        """
-        # Default equilibrium phase: stabilization
-        if equilibrium_config is None:
-            equilibrium_config = {
-                'lambda_mean': 0.5,
-                'excitation_amplitude': 0.0,
-                'duration_epochs': 10
-            }
-
-        # Default exploration phase: discovery
-        if exploration_config is None:
-            exploration_config = {
-                'lambda_mean': 0.05,
-                'excitation_amplitude': 0.05,
-                'duration_epochs': 20
-            }
-
-        self.equilibrium_phase = PhaseConfig(
-            name='equilibrium',
-            **equilibrium_config
-        )
-        self.exploration_phase = PhaseConfig(
-            name='exploration',
-            **exploration_config
-        )
-
-        self.num_cycles = num_cycles
-        self.warmup_epochs = warmup_epochs
-
-        # Build phase schedule
-        self.schedule = self._build_schedule()
-
-        # Current state
-        self.current_epoch = 0
-        self.current_phase = None
-
-    def _build_schedule(self):
-        """Build the complete phase schedule and epoch-to-phase mapping."""
-        schedule = []
-        epoch_to_phase = {}
-
-        # Warmup: equilibrium phase
-        if self.warmup_epochs > 0:
-            schedule.append({
-                'name': 'warmup',
-                'phase': self.equilibrium_phase,
-                'start_epoch': 0,
-                'end_epoch': self.warmup_epochs
-            })
-            # Map warmup epochs
-            for e in range(0, self.warmup_epochs):
-                epoch_to_phase[e] = self.equilibrium_phase
-
-        # Cycles
-        epoch = self.warmup_epochs
-        for cycle in range(self.num_cycles):
-            # Equilibrium phase
-            start = epoch
-            end = epoch + self.equilibrium_phase.duration_epochs
-            schedule.append({
-                'name': f'cycle_{cycle}_equilibrium',
-                'phase': self.equilibrium_phase,
-                'start_epoch': start,
-                'end_epoch': end
-            })
-            # Map equilibrium epochs
-            for e in range(start, end):
-                epoch_to_phase[e] = self.equilibrium_phase
-            epoch = end
-
-            # Exploration phase
-            start = epoch
-            end = epoch + self.exploration_phase.duration_epochs
-            schedule.append({
-                'name': f'cycle_{cycle}_exploration',
-                'phase': self.exploration_phase,
-                'start_epoch': start,
-                'end_epoch': end
-            })
-            # Map exploration epochs
-            for e in range(start, end):
-                epoch_to_phase[e] = self.exploration_phase
-            epoch = end
-
-        self.epoch_to_phase = epoch_to_phase
-        return schedule
-
-    def get_phase_config(self, epoch: int) -> PhaseConfig:
-        """
-        Get the phase configuration for a given epoch.
-
-        Args:
-            epoch: Current training epoch
-
-        Returns:
-            PhaseConfig for the current epoch
-        """
-        # O(1) lookup using pre-computed mapping
-        if epoch in self.epoch_to_phase:
-            return self.epoch_to_phase[epoch]
-
-        # Default to exploration phase after all cycles
-        return self.exploration_phase
-
-    def is_exploration_phase(self, epoch: int) -> bool:
-        """Check if current epoch is in exploration phase."""
-        phase = self.get_phase_config(epoch)
-        return phase.name == 'exploration'
-
-    def step(self, epoch: int) -> Dict[str, any]:
-        """
-        Update scheduler state and return current phase info.
-
-        Args:
-            epoch: Current training epoch
-
-        Returns:
-            Dictionary with phase information
-        """
-        self.current_epoch = epoch
-        self.current_phase = self.get_phase_config(epoch)
-
-        return {
-            'phase_name': self.current_phase.name,
-            'lambda_mean': self.current_phase.lambda_mean,
-            'excitation_amplitude': self.current_phase.excitation_amplitude,
-            'is_exploration': self.current_phase.name == 'exploration'
-        }
-
-    def get_total_epochs(self) -> int:
-        """Get total number of epochs in the schedule."""
-        if not self.schedule:
-            return 0
-        return self.schedule[-1]['end_epoch']
-
-    def print_schedule(self):
-        """Print the complete phase schedule."""
-        print("=" * 70)
-        print("EQUILIBRIUM-EXPLORATION CYCLE SCHEDULE")
-        print("=" * 70)
-
-        for entry in self.schedule:
-            phase = entry['phase']
-            print(f"\n{entry['name'].upper()}")
-            print(f"  Epochs: {entry['start_epoch']}-{entry['end_epoch']} "
-                  f"({entry['end_epoch'] - entry['start_epoch']} epochs)")
-            print(f"  Lambda Mean: {phase.lambda_mean:.3f}")
-            print(f"  Excitation Amplitude: {phase.excitation_amplitude:.3f}")
-            print(f"  Phase Type: {phase.name}")
-
-        print(f"\nTotal Training Epochs: {self.get_total_epochs()}")
-        print("=" * 70)
-
-
-class CycleTrainingMixin:
-    """
-    Mixin class to add cycle scheduling to existing trainers.
-
-    Usage:
-        class MyTrainer(CycleTrainingMixin, BaseTrainer):
-            ...
-    """
-
-    def setup_cycle_scheduler(
-        self,
-        equilibrium_config: Dict = None,
-        exploration_config: Dict = None,
-        num_cycles: int = 5,
-        warmup_epochs: int = 10
-    ):
-        """
-        Initialize the phase scheduler.
-
-        Call this in your trainer's __init__ method.
-        """
-        self.cycle_scheduler = EquilibriumExplorationScheduler(
-            equilibrium_config=equilibrium_config,
-            exploration_config=exploration_config,
-            num_cycles=num_cycles,
-            warmup_epochs=warmup_epochs
-        )
-
-        self.cycle_enabled = True
-        self.cycle_scheduler.print_schedule()
-
-    def update_phase(self, epoch: int, model, loss_fn):
-        """
-        Update model and loss function for current phase.
-
-        Args:
-            epoch: Current training epoch
-            model: IntegratorModel
-            loss_fn: IntegratorLoss
-        """
-        if not hasattr(self, 'cycle_enabled') or not self.cycle_enabled:
-            return
-
-        # Get current phase config
-        phase_info = self.cycle_scheduler.step(epoch)
-
-        # Update loss function phase
-        if hasattr(loss_fn, 'set_exploration_phase'):
-            loss_fn.set_exploration_phase(phase_info['is_exploration'])
-
-        # Update lambda_mean in loss function
-        if hasattr(loss_fn, 'lambda_mean'):
-            loss_fn.lambda_mean = phase_info['lambda_mean']
-
-        # Update model excitation amplitude
-        if hasattr(model, 'inl') and hasattr(model.inl, 'excitation_amplitude'):
-            model.inl.excitation_amplitude = phase_info['excitation_amplitude']
-
-        # Update all INL blocks in language model if applicable
-        if hasattr(model, 'blocks'):
-            for block in model.blocks:
-                if hasattr(block, 'inl') and hasattr(block.inl, 'excitation_amplitude'):
-                    block.inl.excitation_amplitude = phase_info['excitation_amplitude']
-
-        return phase_info
-
-
-# Example configuration presets
-CYCLE_PRESETS = {
-    'conservative': {
-        'equilibrium_config': {
-            'lambda_mean': 0.8,
-            'excitation_amplitude': 0.0,
-            'duration_epochs': 15
-        },
-        'exploration_config': {
-            'lambda_mean': 0.1,
-            'excitation_amplitude': 0.02,
-            'duration_epochs': 15
-        },
-        'num_cycles': 4,
-        'warmup_epochs': 20
-    },
-
-    'balanced': {
-        'equilibrium_config': {
-            'lambda_mean': 0.5,
-            'excitation_amplitude': 0.0,
-            'duration_epochs': 10
-        },
-        'exploration_config': {
-            'lambda_mean': 0.05,
-            'excitation_amplitude': 0.05,
-            'duration_epochs': 20
-        },
-        'num_cycles': 5,
-        'warmup_epochs': 10
-    },
-
-    'aggressive': {
-        'equilibrium_config': {
-            'lambda_mean': 0.3,
-            'excitation_amplitude': 0.0,
-            'duration_epochs': 5
-        },
-        'exploration_config': {
-            'lambda_mean': 0.01,
-            'excitation_amplitude': 0.08,
-            'duration_epochs': 25
-        },
-        'num_cycles': 6,
-        'warmup_epochs': 5
-    }
-}
-
-
-def _scale_config_to_epochs(config: dict, total_epochs: int, preset_name: str) -> dict:
-    """
-    Scale a preset configuration to fit a target number of epochs.
-
-    Strategy based on preset ratios:
-    - conservative: 30% warmup, 35% equilibrium, 35% exploration
-    - balanced: 25% warmup, 25% equilibrium, 50% exploration
-    - aggressive: 15% warmup, 15% equilibrium, 70% exploration
-
-    Args:
-        config: Original preset config
-        total_epochs: Target total epochs
-        preset_name: Name of the preset (for ratio selection)
-
-    Returns:
-        Scaled configuration
-    """
-    # Define ratios for each preset
-    ratios = {
-        'conservative': {'warmup': 0.30, 'equilibrium': 0.35, 'exploration': 0.35},
-        'balanced': {'warmup': 0.25, 'equilibrium': 0.25, 'exploration': 0.50},
-        'aggressive': {'warmup': 0.15, 'equilibrium': 0.15, 'exploration': 0.70}
-    }
-
-    ratio = ratios.get(preset_name, ratios['balanced'])
-
-    # Calculate epochs for each phase
-    warmup_epochs = max(1, int(total_epochs * ratio['warmup']))
-    equilibrium_epochs = max(1, int(total_epochs * ratio['equilibrium']))
-    exploration_epochs = max(1, total_epochs - warmup_epochs - equilibrium_epochs)
-
-    # Update config
-    config['warmup_epochs'] = warmup_epochs
-    config['num_cycles'] = 1  # Single cycle for simplicity
-    config['equilibrium_config']['duration_epochs'] = equilibrium_epochs
-    config['exploration_config']['duration_epochs'] = exploration_epochs
-
-    return config
-
-
-def create_cycle_scheduler(preset: str = 'balanced', total_epochs: int = None, **overrides) -> EquilibriumExplorationScheduler:
-    """
-    Create a cycle scheduler from a preset configuration.
-
-    Args:
-        preset: One of 'conservative', 'balanced', 'aggressive'
-        total_epochs: If provided, automatically scales the preset to fit this many epochs
-        **overrides: Override any preset parameters. For nested configs like
-                    equilibrium_config or exploration_config, partial overrides
-                    are merged with preset defaults.
-
-    Returns:
-        Configured EquilibriumExplorationScheduler
-
-    Examples:
-        # Automatic scaling to fit 20 epochs
-        scheduler = create_cycle_scheduler('balanced', total_epochs=20)
-
-        # Override just lambda_mean in equilibrium phase
-        scheduler = create_cycle_scheduler('balanced',
-            equilibrium_config={'lambda_mean': 0.9})
-
-        # Override multiple top-level parameters
-        scheduler = create_cycle_scheduler('aggressive',
-            num_cycles=10,
-            warmup_epochs=15)
-    """
-    if preset not in CYCLE_PRESETS:
-        raise ValueError(f"Unknown preset '{preset}'. Choose from: {list(CYCLE_PRESETS.keys())}")
-
-    # Deep copy to avoid mutating the preset
-    config = deepcopy(CYCLE_PRESETS[preset])
-
-    # AUTO-SCALE: Adapt preset to fit total_epochs
-    if total_epochs is not None:
-        config = _scale_config_to_epochs(config, total_epochs, preset)
-
-    # Merge nested configs intelligently
-    for key, value in overrides.items():
-        if key in ('equilibrium_config', 'exploration_config') and isinstance(value, dict):
-            # Merge nested config instead of replacing it entirely
-            if key in config:
-                config[key].update(value)
-            else:
-                config[key] = value
-        else:
-            # Simple override for non-nested parameters
-            config[key] = value
-
-    return EquilibriumExplorationScheduler(**config)
-
-
-if __name__ == '__main__':
-    # Demo: print different scheduler configurations
-    print("\n" + "=" * 70)
-    print("CYCLE SCHEDULER DEMONSTRATION")
-    print("=" * 70)
-
-    for preset_name in ['conservative', 'balanced', 'aggressive']:
-        print(f"\n\nPRESET: {preset_name.upper()}")
-        scheduler = create_cycle_scheduler(preset_name)
-        scheduler.print_schedule()
-
-        # Show epoch-by-epoch evolution for first 50 epochs
-        print(f"\nFirst 50 epochs evolution:")
-        for epoch in range(min(50, scheduler.get_total_epochs())):
-            if epoch % 10 == 0:
-                phase_info = scheduler.step(epoch)
-                print(f"  Epoch {epoch:3d}: {phase_info['phase_name']:20s} "
-                      f"λ={phase_info['lambda_mean']:.3f} β={phase_info['excitation_amplitude']:.3f}")
+"""
+INL-LLM: Equilibrium-Exploration Cycle Scheduler
+
+Implements rhythmic training phases that alternate between:
+- Equilibrium Phase: Strong stability constraint, low excitation (stabilization)
+- Exploration Phase: Weak stability constraint, high excitation (discovery)
+
+This deterministic cycling encourages structured exploration without randomness.
+"""
+
+from typing import Dict, NamedTuple
+import torch
+from copy import deepcopy
+
+
+class PhaseConfig(NamedTuple):
+    """Configuration for a training phase."""
+    name: str
+    lambda_mean: float
+    excitation_amplitude: float
+    duration_epochs: int
+
+
+class EquilibriumExplorationScheduler:
+    """
+    Manages equilibrium-exploration cycles for v2 training.
+
+    Example cycle:
+    - Equilibrium (10 epochs): lambda_mean=0.5, excitation=0.0
+    - Exploration (20 epochs): lambda_mean=0.05, excitation=0.05
+    - Repeat...
+    """
+
+    def __init__(
+        self,
+        equilibrium_config: Dict = None,
+        exploration_config: Dict = None,
+        num_cycles: int = 5,
+        warmup_epochs: int = 10
+    ):
+        """
+        Args:
+            equilibrium_config: Config for equilibrium phase
+            exploration_config: Config for exploration phase
+            num_cycles: Number of complete cycles to perform
+            warmup_epochs: Initial warmup before cycling starts
+        """
+        # Default equilibrium phase: stabilization
+        if equilibrium_config is None:
+            equilibrium_config = {
+                'lambda_mean': 0.5,
+                'excitation_amplitude': 0.0,
+                'duration_epochs': 10
+            }
+
+        # Default exploration phase: discovery
+        if exploration_config is None:
+            exploration_config = {
+                'lambda_mean': 0.05,
+                'excitation_amplitude': 0.05,
+                'duration_epochs': 20
+            }
+
+        self.equilibrium_phase = PhaseConfig(
+            name='equilibrium',
+            **equilibrium_config
+        )
+        self.exploration_phase = PhaseConfig(
+            name='exploration',
+            **exploration_config
+        )
+
+        self.num_cycles = num_cycles
+        self.warmup_epochs = warmup_epochs
+
+        # Build phase schedule
+        self.schedule = self._build_schedule()
+
+        # Current state
+        self.current_epoch = 0
+        self.current_phase = None
+
+    def _build_schedule(self):
+        """Build the complete phase schedule and epoch-to-phase mapping."""
+        schedule = []
+        epoch_to_phase = {}
+
+        # Warmup: equilibrium phase
+        if self.warmup_epochs > 0:
+            schedule.append({
+                'name': 'warmup',
+                'phase': self.equilibrium_phase,
+                'start_epoch': 0,
+                'end_epoch': self.warmup_epochs
+            })
+            # Map warmup epochs
+            for e in range(0, self.warmup_epochs):
+                epoch_to_phase[e] = self.equilibrium_phase
+
+        # Cycles
+        epoch = self.warmup_epochs
+        for cycle in range(self.num_cycles):
+            # Equilibrium phase
+            start = epoch
+            end = epoch + self.equilibrium_phase.duration_epochs
+            schedule.append({
+                'name': f'cycle_{cycle}_equilibrium',
+                'phase': self.equilibrium_phase,
+                'start_epoch': start,
+                'end_epoch': end
+            })
+            # Map equilibrium epochs
+            for e in range(start, end):
+                epoch_to_phase[e] = self.equilibrium_phase
+            epoch = end
+
+            # Exploration phase
+            start = epoch
+            end = epoch + self.exploration_phase.duration_epochs
+            schedule.append({
+                'name': f'cycle_{cycle}_exploration',
+                'phase': self.exploration_phase,
+                'start_epoch': start,
+                'end_epoch': end
+            })
+            # Map exploration epochs
+            for e in range(start, end):
+                epoch_to_phase[e] = self.exploration_phase
+            epoch = end
+
+        self.epoch_to_phase = epoch_to_phase
+        return schedule
+
+    def get_phase_config(self, epoch: int) -> PhaseConfig:
+        """
+        Get the phase configuration for a given epoch.
+
+        Args:
+            epoch: Current training epoch
+
+        Returns:
+            PhaseConfig for the current epoch
+        """
+        # O(1) lookup using pre-computed mapping
+        if epoch in self.epoch_to_phase:
+            return self.epoch_to_phase[epoch]
+
+        # Default to exploration phase after all cycles
+        return self.exploration_phase
+
+    def is_exploration_phase(self, epoch: int) -> bool:
+        """Check if current epoch is in exploration phase."""
+        phase = self.get_phase_config(epoch)
+        return phase.name == 'exploration'
+
+    def step(self, epoch: int) -> Dict[str, any]:
+        """
+        Update scheduler state and return current phase info.
+
+        Args:
+            epoch: Current training epoch
+
+        Returns:
+            Dictionary with phase information
+        """
+        self.current_epoch = epoch
+        self.current_phase = self.get_phase_config(epoch)
+
+        return {
+            'phase_name': self.current_phase.name,
+            'lambda_mean': self.current_phase.lambda_mean,
+            'excitation_amplitude': self.current_phase.excitation_amplitude,
+            'is_exploration': self.current_phase.name == 'exploration'
+        }
+
+    def get_total_epochs(self) -> int:
+        """Get total number of epochs in the schedule."""
+        if not self.schedule:
+            return 0
+        return self.schedule[-1]['end_epoch']
+
+    def print_schedule(self):
+        """Print the complete phase schedule."""
+        print("=" * 70)
+        print("EQUILIBRIUM-EXPLORATION CYCLE SCHEDULE")
+        print("=" * 70)
+
+        for entry in self.schedule:
+            phase = entry['phase']
+            print(f"\n{entry['name'].upper()}")
+            print(f"  Epochs: {entry['start_epoch']}-{entry['end_epoch']} "
+                  f"({entry['end_epoch'] - entry['start_epoch']} epochs)")
+            print(f"  Lambda Mean: {phase.lambda_mean:.3f}")
+            print(f"  Excitation Amplitude: {phase.excitation_amplitude:.3f}")
+            print(f"  Phase Type: {phase.name}")
+
+        print(f"\nTotal Training Epochs: {self.get_total_epochs()}")
+        print("=" * 70)
+
+
+class CycleTrainingMixin:
+    """
+    Mixin class to add cycle scheduling to existing trainers.
+
+    Usage:
+        class MyTrainer(CycleTrainingMixin, BaseTrainer):
+            ...
+    """
+
+    def setup_cycle_scheduler(
+        self,
+        equilibrium_config: Dict = None,
+        exploration_config: Dict = None,
+        num_cycles: int = 5,
+        warmup_epochs: int = 10
+    ):
+        """
+        Initialize the phase scheduler.
+
+        Call this in your trainer's __init__ method.
+        """
+        self.cycle_scheduler = EquilibriumExplorationScheduler(
+            equilibrium_config=equilibrium_config,
+            exploration_config=exploration_config,
+            num_cycles=num_cycles,
+            warmup_epochs=warmup_epochs
+        )
+
+        self.cycle_enabled = True
+        self.cycle_scheduler.print_schedule()
+
+    def update_phase(self, epoch: int, model, loss_fn):
+        """
+        Update model and loss function for current phase.
+
+        Args:
+            epoch: Current training epoch
+            model: IntegratorModel
+            loss_fn: IntegratorLoss
+        """
+        if not hasattr(self, 'cycle_enabled') or not self.cycle_enabled:
+            return
+
+        # Get current phase config
+        phase_info = self.cycle_scheduler.step(epoch)
+
+        # Update loss function phase
+        if hasattr(loss_fn, 'set_exploration_phase'):
+            loss_fn.set_exploration_phase(phase_info['is_exploration'])
+
+        # Update lambda_mean in loss function
+        if hasattr(loss_fn, 'lambda_mean'):
+            loss_fn.lambda_mean = phase_info['lambda_mean']
+
+        # Update model excitation amplitude
+        if hasattr(model, 'inl') and hasattr(model.inl, 'excitation_amplitude'):
+            model.inl.excitation_amplitude = phase_info['excitation_amplitude']
+
+        # Update all INL blocks in language model if applicable
+        if hasattr(model, 'blocks'):
+            for block in model.blocks:
+                if hasattr(block, 'inl') and hasattr(block.inl, 'excitation_amplitude'):
+                    block.inl.excitation_amplitude = phase_info['excitation_amplitude']
+
+        return phase_info
+
+
+# Example configuration presets
+CYCLE_PRESETS = {
+    'conservative': {
+        'equilibrium_config': {
+            'lambda_mean': 0.8,
+            'excitation_amplitude': 0.0,
+            'duration_epochs': 15
+        },
+        'exploration_config': {
+            'lambda_mean': 0.1,
+            'excitation_amplitude': 0.02,
+            'duration_epochs': 15
+        },
+        'num_cycles': 4,
+        'warmup_epochs': 20
+    },
+
+    'balanced': {
+        'equilibrium_config': {
+            'lambda_mean': 0.5,
+            'excitation_amplitude': 0.0,
+            'duration_epochs': 10
+        },
+        'exploration_config': {
+            'lambda_mean': 0.05,
+            'excitation_amplitude': 0.05,
+            'duration_epochs': 20
+        },
+        'num_cycles': 5,
+        'warmup_epochs': 10
+    },
+
+    'aggressive': {
+        'equilibrium_config': {
+            'lambda_mean': 0.3,
+            'excitation_amplitude': 0.0,
+            'duration_epochs': 5
+        },
+        'exploration_config': {
+            'lambda_mean': 0.01,
+            'excitation_amplitude': 0.08,
+            'duration_epochs': 25
+        },
+        'num_cycles': 6,
+        'warmup_epochs': 5
+    }
+}
+
+
+def _scale_config_to_epochs(config: dict, total_epochs: int, preset_name: str) -> dict:
+    """
+    Scale a preset configuration to fit a target number of epochs.
+
+    Strategy based on preset ratios:
+    - conservative: 30% warmup, 35% equilibrium, 35% exploration
+    - balanced: 25% warmup, 25% equilibrium, 50% exploration
+    - aggressive: 15% warmup, 15% equilibrium, 70% exploration
+
+    Args:
+        config: Original preset config
+        total_epochs: Target total epochs
+        preset_name: Name of the preset (for ratio selection)
+
+    Returns:
+        Scaled configuration
+    """
+    # Define ratios for each preset
+    ratios = {
+        'conservative': {'warmup': 0.30, 'equilibrium': 0.35, 'exploration': 0.35},
+        'balanced': {'warmup': 0.25, 'equilibrium': 0.25, 'exploration': 0.50},
+        'aggressive': {'warmup': 0.15, 'equilibrium': 0.15, 'exploration': 0.70}
+    }
+
+    ratio = ratios.get(preset_name, ratios['balanced'])
+
+    # Calculate epochs for each phase
+    warmup_epochs = max(1, int(total_epochs * ratio['warmup']))
+    equilibrium_epochs = max(1, int(total_epochs * ratio['equilibrium']))
+    exploration_epochs = max(1, total_epochs - warmup_epochs - equilibrium_epochs)
+
+    # Update config
+    config['warmup_epochs'] = warmup_epochs
+    config['num_cycles'] = 1  # Single cycle for simplicity
+    config['equilibrium_config']['duration_epochs'] = equilibrium_epochs
+    config['exploration_config']['duration_epochs'] = exploration_epochs
+
+    return config
+
+
+def create_cycle_scheduler(preset: str = 'balanced', total_epochs: int = None, **overrides) -> EquilibriumExplorationScheduler:
+    """
+    Create a cycle scheduler from a preset configuration.
+
+    Args:
+        preset: One of 'conservative', 'balanced', 'aggressive'
+        total_epochs: If provided, automatically scales the preset to fit this many epochs
+        **overrides: Override any preset parameters. For nested configs like
+                    equilibrium_config or exploration_config, partial overrides
+                    are merged with preset defaults.
+
+    Returns:
+        Configured EquilibriumExplorationScheduler
+
+    Examples:
+        # Automatic scaling to fit 20 epochs
+        scheduler = create_cycle_scheduler('balanced', total_epochs=20)
+
+        # Override just lambda_mean in equilibrium phase
+        scheduler = create_cycle_scheduler('balanced',
+            equilibrium_config={'lambda_mean': 0.9})
+
+        # Override multiple top-level parameters
+        scheduler = create_cycle_scheduler('aggressive',
+            num_cycles=10,
+            warmup_epochs=15)
+    """
+    if preset not in CYCLE_PRESETS:
+        raise ValueError(f"Unknown preset '{preset}'. Choose from: {list(CYCLE_PRESETS.keys())}")
+
+    # Deep copy to avoid mutating the preset
+    config = deepcopy(CYCLE_PRESETS[preset])
+
+    # AUTO-SCALE: Adapt preset to fit total_epochs
+    if total_epochs is not None:
+        config = _scale_config_to_epochs(config, total_epochs, preset)
+
+    # Merge nested configs intelligently
+    for key, value in overrides.items():
+        if key in ('equilibrium_config', 'exploration_config') and isinstance(value, dict):
+            # Merge nested config instead of replacing it entirely
+            if key in config:
+                config[key].update(value)
+            else:
+                config[key] = value
+        else:
+            # Simple override for non-nested parameters
+            config[key] = value
+
+    return EquilibriumExplorationScheduler(**config)
+
+
+if __name__ == '__main__':
+    # Demo: print different scheduler configurations
+    print("\n" + "=" * 70)
+    print("CYCLE SCHEDULER DEMONSTRATION")
+    print("=" * 70)
+
+    for preset_name in ['conservative', 'balanced', 'aggressive']:
+        print(f"\n\nPRESET: {preset_name.upper()}")
+        scheduler = create_cycle_scheduler(preset_name)
+        scheduler.print_schedule()
+
+        # Show epoch-by-epoch evolution for first 50 epochs
+        print(f"\nFirst 50 epochs evolution:")
+        for epoch in range(min(50, scheduler.get_total_epochs())):
+            if epoch % 10 == 0:
+                phase_info = scheduler.step(epoch)
+                print(f"  Epoch {epoch:3d}: {phase_info['phase_name']:20s} "
+                      f"λ={phase_info['lambda_mean']:.3f} β={phase_info['excitation_amplitude']:.3f}")

inl_llm/core/moe_budget_integration.py

@@ -0,0 +1,484 @@
+"""
+Integration: MoE Controller + AdaptiveBudgetAllocator-v2
+
+This module combines the power of:
+1. MoE Controller: Intelligent routing between specialized experts
+2. AdaptiveBudgetAllocator-v2: Smart iteration budget management
+
+The combination enables:
+- Expert specialization per layer + phase
+- Budget allocation adapted to expert choices
+- Loss-component feedback to both MoE and budget allocator
+- Comprehensive monitoring and statistics
+
+Expected Performance:
+- 30-50% compute savings (budget allocator)
+- 2-3x model capacity (MoE)
+- Automatic specialization (emergent behavior)
+- Phase-aware adaptation (equilibrium/exploration)
+
+Author: Boris Peyriguère
+"""
+
+import torch
+import torch.nn as nn
+from typing import Dict, List, Optional, Tuple, Any
+
+from .moe_controller import INLMixtureOfExperts, create_moe_controller
+from .adaptive_budget_allocator import (
+    AdaptiveBudgetAllocator,
+    create_budget_allocator
+)
+
+
+class MoEBudgetAwareINLLayer(nn.Module):
+    """
+    INL Layer with BOTH MoE Controller AND Adaptive Budget Allocation.
+
+    This is the ULTIMATE optimization combining:
+    - MoE: Smart expert routing for capacity
+    - Budget Allocator: Smart iteration management for efficiency
+    - Multi-criteria convergence
+    - Budget redistribution
+    - Phase awareness
+    - Loss-component feedback
+
+    The two systems work synergistically:
+    - MoE provides specialized control strategies
+    - Budget allocator optimizes compute per layer
+    - Both adapt to phase and loss signals
+    """
+
+    def __init__(
+        self,
+        inl_layer: nn.Module,
+        layer_idx: int,
+        d_model: int,
+        num_layers: int,
+        budget_allocator: Optional[AdaptiveBudgetAllocator] = None,
+        moe_controller: Optional[INLMixtureOfExperts] = None,
+        use_moe_for_mu: bool = False
+    ):
+        """
+        Args:
+            inl_layer: Base INL layer (can be None if using MoE for all dynamics)
+            layer_idx: Layer index
+            d_model: Model dimension
+            num_layers: Total number of layers
+            budget_allocator: Budget allocator instance (shared across layers)
+            moe_controller: MoE controller instance (shared across layers)
+            use_moe_for_mu: Use MoE to predict equilibrium mu (experimental)
+        """
+        super().__init__()
+
+        self.inl_layer = inl_layer
+        self.layer_idx = layer_idx
+        self.d_model = d_model
+        self.num_layers = num_layers
+        self.budget_allocator = budget_allocator
+        self.moe_controller = moe_controller
+        self.use_moe_for_mu = use_moe_for_mu
+
+        # Optional: MoE-predicted equilibrium
+        if use_moe_for_mu and moe_controller is not None:
+            self.mu_predictor = nn.Linear(d_model, d_model)
+
+    def forward(
+        self,
+        h: torch.Tensor,
+        x_init: torch.Tensor,
+        v_init: torch.Tensor,
+        default_iterations: int = 5,
+        return_trajectory: bool = False,
+        mu: Optional[torch.Tensor] = None,
+        loss_components: Optional[Dict[str, float]] = None,
+        phase: str = 'equilibrium',
+        attention_weights: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor, Dict]:
+        """
+        Forward pass with MoE control and adaptive budget.
+
+        Args:
+            h: Context embedding [batch, d_model]
+            x_init: Initial state [batch, d_model]
+            v_init: Initial velocity [batch, d_model]
+            default_iterations: Default iterations if no budget allocator
+            return_trajectory: Whether to return full trajectory
+            mu: Learned equilibrium (for error-based convergence)
+            loss_components: Loss components dict (L_speed, L_energy, L_mean)
+            phase: Training phase ('equilibrium' or 'exploration')
+            attention_weights: Attention pattern for MoE routing
+
+        Returns:
+            x_final: Final state
+            v_final: Final velocity
+            info: Dictionary with comprehensive statistics
+        """
+        batch_size = h.size(0)
+        device = h.device
+
+        # Phase 1: Get iteration budget (with redistribution bonus)
+        if self.budget_allocator is not None:
+            bonus = self.budget_allocator.get_redistribution_bonus(self.layer_idx)
+            max_iters = self.budget_allocator.get_layer_budget(
+                self.layer_idx,
+                training=self.training,
+                bonus_budget=bonus
+            )
+        else:
+            max_iters = default_iterations
+
+        # Phase 2: Optional - Predict mu using MoE
+        if self.use_moe_for_mu and self.moe_controller is not None and mu is None:
+            # Use MoE to predict equilibrium target
+            with torch.no_grad():
+                alpha_pred, _, _, _ , _ = self.moe_controller(h, x_init, self.layer_idx, phase)
+                mu = self.mu_predictor(alpha_pred)
+
+        # Phase 3: Run integrator with MoE control
+        x, v = x_init, v_init
+        x_prev = x_init
+
+        if return_trajectory:
+            x_traj = [x.clone()]
+            v_traj = [v.clone()]
+
+        actual_iterations = 0
+        converged = False
+        convergence_metrics = {}
+        moe_info_history = []
+
+        for iteration in range(max_iters):
+            # Get MoE control parameters
+            if self.moe_controller is not None:
+                alpha, beta, gate, v_cand, moe_info = self.moe_controller(
+                    h, x, self.layer_idx, phase, attention_weights
+                )
+                moe_info_history.append(moe_info)
+            else:
+                # Fallback: use base INL layer controller
+                alpha, beta, gate, v_cand = self._get_default_control(h, x)
+                moe_info = {}
+
+            # INL integration step with MoE control
+            x_next, v_next = self._integration_step(
+                h, x, v, alpha, beta, gate, v_cand, mu, iteration
+            )
+
+            # Check convergence (multi-criteria if enabled)
+            if self.budget_allocator is not None and iteration >= self.budget_allocator.warmup_iterations:
+                converged, convergence_metrics = self.budget_allocator.check_convergence(
+                    x_next, x, iteration,
+                    v_current=v_next,
+                    mu=mu
+                )
+                if converged and not self.training:
+                    # Early stop during inference
+                    x, v = x_next, v_next
+                    actual_iterations = iteration + 1
+                    break
+
+            x_prev = x
+            x, v = x_next, v_next
+            actual_iterations = iteration + 1
+
+            if return_trajectory:
+                x_traj.append(x.clone())
+                v_traj.append(v.clone())
+
+        # Phase 4: Update statistics and redistribute budget
+        if self.budget_allocator is not None:
+            # Add unused budget to redistribution pool
+            unused = max_iters - actual_iterations
+            self.budget_allocator.add_to_budget_pool(unused)
+
+            # Update statistics with all metrics
+            if self.training:
+                final_delta = torch.norm(x - x_prev, dim=-1).mean().item()
+                final_velocity = torch.norm(v, dim=-1).mean().item() if v is not None else 0.0
+                final_error = torch.norm(x - mu, dim=-1).mean().item() if mu is not None else 0.0
+
+                # Extract gradient magnitude if possible
+                grad_mag = None
+                if x.requires_grad and x.grad is not None:
+                    grad_mag = torch.norm(x.grad, dim=-1).mean().item()
+
+                self.budget_allocator.update_statistics(
+                    self.layer_idx,
+                    actual_iterations,
+                    final_delta,
+                    budget_allocated=max_iters,
+                    final_velocity=final_velocity,
+                    final_error=final_error,
+                    loss_components=loss_components,
+                    grad_magnitude=grad_mag
+                )
+
+        # Phase 5: Aggregate MoE information
+        moe_summary = self._aggregate_moe_info(moe_info_history)
+
+        # Prepare comprehensive output info
+        info = {
+            # Budget allocator info
+            'iterations_used': actual_iterations,
+            'max_iterations': max_iters,
+            'converged': converged,
+            'layer_idx': self.layer_idx,
+            'convergence_metrics': convergence_metrics,
+
+            # MoE info
+            'moe_summary': moe_summary,
+
+            # Phase info
+            'phase': phase
+        }
+
+        if return_trajectory:
+            info['x_trajectory'] = torch.stack(x_traj, dim=1)
+            info['v_trajectory'] = torch.stack(v_traj, dim=1)
+
+        return x, v, info
+
+    def _integration_step(
+        self,
+        h: torch.Tensor,
+        x: torch.Tensor,
+        v: torch.Tensor,
+        alpha: torch.Tensor,
+        beta: torch.Tensor,
+        gate: torch.Tensor,
+        v_cand: torch.Tensor,
+        mu: Optional[torch.Tensor],
+        step: int
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Single INL integration step with MoE-provided control parameters.
+
+        Implements the core INL dynamics:
+        error = x - mu
+        v_next = alpha * v + (1 - alpha) * v_cand - beta * error
+        x_next = x + gate * v_next
+        """
+        # Compute error term
+        if mu is not None:
+            error = x - mu
+        else:
+            error = torch.zeros_like(x)
+
+        # Velocity update with MoE control
+        v_next = alpha * v + (1 - alpha) * v_cand - beta * error
+
+        # State update with gating
+        x_next = x + gate * v_next
+
+        return x_next, v_next
+
+    def _get_default_control(
+        self,
+        h: torch.Tensor,
+        x: torch.Tensor
+    ) -> Tuple[torch.Tensor, ...]:
+        """Fallback: get control from base INL layer if no MoE."""
+        if hasattr(self.inl_layer, 'controller'):
+            return self.inl_layer.controller(h, x)
+        else:
+            # Simple defaults
+            batch_size = h.size(0)
+            alpha = torch.ones(batch_size, self.d_model, device=h.device) * 0.5
+            beta = torch.ones(batch_size, self.d_model, device=h.device) * 0.1
+            gate = torch.ones(batch_size, self.d_model, device=h.device) * 0.9
+            v_cand = torch.zeros(batch_size, self.d_model, device=h.device)
+            return alpha, beta, gate, v_cand
+
+    def _aggregate_moe_info(self, moe_info_history: List[Dict]) -> Dict:
+        """Aggregate MoE information across iterations."""
+        if not moe_info_history:
+            return {}
+
+        # Average routing weights across iterations
+        all_weights = [info['routing_weights'] for info in moe_info_history if 'routing_weights' in info]
+        if all_weights:
+            avg_routing_weights = torch.stack(all_weights).mean(dim=0)
+        else:
+            avg_routing_weights = None
+
+        # Collect expert usage
+        expert_usage = {}
+        for info in moe_info_history:
+            if 'selected_experts' in info and info['selected_experts'] is not None:
+                for expert_id in info['selected_experts'].flatten().tolist():
+                    expert_usage[expert_id] = expert_usage.get(expert_id, 0) + 1
+
+        # Aggregate auxiliary losses
+        aux_losses = {}
+        if 'aux_losses' in moe_info_history[-1]:
+            for loss_name, loss_value in moe_info_history[-1]['aux_losses'].items():
+                aux_losses[loss_name] = loss_value
+
+        return {
+            'avg_routing_weights': avg_routing_weights,
+            'expert_usage': expert_usage,
+            'aux_losses': aux_losses,
+            'num_iterations': len(moe_info_history)
+        }
+
+
+def create_moe_budget_model(
+    d_model: int,
+    num_layers: int,
+    # Budget allocator params
+    total_budget: int = 125,
+    budget_strategy: str = 'hybrid',
+    # MoE params
+    num_experts: int = 4,
+    top_k: int = 2,
+    # Shared params
+    use_phase_aware: bool = True,
+    use_loss_tracking: bool = True,
+    **kwargs
+) -> Tuple[AdaptiveBudgetAllocator, INLMixtureOfExperts]:
+    """
+    Helper to create both MoE controller and budget allocator.
+
+    Args:
+        d_model: Model dimension
+        num_layers: Number of layers
+        total_budget: Total iteration budget
+        budget_strategy: Budget allocation strategy
+        num_experts: Number of MoE experts
+        top_k: Number of experts to activate
+        use_phase_aware: Enable phase-aware features
+        use_loss_tracking: Enable loss-component tracking
+        **kwargs: Additional arguments
+
+    Returns:
+        budget_allocator: AdaptiveBudgetAllocator instance
+        moe_controller: INLMixtureOfExperts instance
+    """
+    # Create budget allocator
+    budget_allocator = AdaptiveBudgetAllocator(
+        num_layers=num_layers,
+        total_budget=total_budget,
+        strategy=budget_strategy,
+        use_phase_aware=use_phase_aware,
+        use_loss_tracking=use_loss_tracking,
+        **{k: v for k, v in kwargs.items() if k.startswith('use_') or k in [
+            'min_iterations_per_layer', 'max_iterations_per_layer',
+            'convergence_threshold', 'warmup_iterations',
+            'velocity_threshold', 'error_threshold', 'redistribution_window'
+        ]}
+    )
+
+    # Create MoE controller
+    moe_controller = create_moe_controller(
+        d_model=d_model,
+        num_layers=num_layers,
+        num_experts=num_experts,
+        top_k=top_k,
+        **{k: v for k, v in kwargs.items() if k in [
+            'expert_hidden_dim', 'router_hidden_dim',
+            'use_sparse_routing', 'load_balance_weight',
+            'router_z_loss_weight', 'use_attention_features'
+        ]}
+    )
+
+    return budget_allocator, moe_controller
+
+
+if __name__ == '__main__':
+    print("=" * 70)
+    print("MoE + BUDGET ALLOCATOR INTEGRATION - Test")
+    print("=" * 70)
+
+    # Configuration
+    d_model = 1024
+    num_layers = 25
+    batch_size = 16
+    seq_len = 128
+
+    # Create integrated system
+    print("\n🔧 Creating MoE + Budget Allocator...")
+    budget_allocator, moe_controller = create_moe_budget_model(
+        d_model=d_model,
+        num_layers=num_layers,
+        total_budget=125,
+        budget_strategy='hybrid',
+        num_experts=4,
+        top_k=2
+    )
+
+    print(f"\n{budget_allocator}")
+    print(f"\n{moe_controller}")
+
+    # Create test layer (mock INL layer)
+    class MockINLLayer(nn.Module):
+        def __init__(self, d_model):
+            super().__init__()
+            self.d_model = d_model
+
+        def forward(self, h, x, v, step):
+            # Mock forward
+            return x, v, {}
+
+    test_layer = MoEBudgetAwareINLLayer(
+        inl_layer=MockINLLayer(d_model),
+        layer_idx=12,
+        d_model=d_model,
+        num_layers=num_layers,
+        budget_allocator=budget_allocator,
+        moe_controller=moe_controller
+    )
+
+    # Test forward pass
+    print("\n🧪 Testing integrated forward pass...")
+    h = torch.randn(batch_size, d_model)
+    x_init = torch.randn(batch_size, d_model)
+    v_init = torch.randn(batch_size, d_model)
+    mu = torch.randn(batch_size, d_model)
+
+    # Test different phases
+    for phase in ['equilibrium', 'exploration']:
+        print(f"\n  Phase: {phase}")
+        budget_allocator.set_phase(phase)
+
+        x, v, info = test_layer(
+            h, x_init, v_init,
+            phase=phase,
+            mu=mu,
+            loss_components={'L_speed': 0.1, 'L_energy': 0.05, 'L_mean': 0.2}
+        )
+
+        print(f"    Iterations: {info['iterations_used']}/{info['max_iterations']}")
+        print(f"    Converged: {info['converged']}")
+        print(f"    MoE experts used: {info['moe_summary'].get('expert_usage', {})}")
+
+        if 'convergence_metrics' in info:
+            print(f"    Convergence metrics: {info['convergence_metrics']}")
+
+    # Statistics
+    print("\n📊 System Statistics:")
+
+    print("\n  Budget Allocator:")
+    budget_stats = budget_allocator.get_statistics()
+    print(f"    Phase: {budget_stats['current_phase']}")
+    print(f"    Updates: {int(budget_stats['updates'].item())}")
+    print(f"    Budget pool: {budget_stats['current_budget_pool']:.2f}")
+
+    print("\n  MoE Controller:")
+    moe_stats = moe_controller.get_expert_statistics()
+    print(f"    Load balance score: {moe_stats['load_balance_score'].item():.3f}")
+    print(f"    Router calls: {int(moe_stats['router_calls'].item())}")
+    for i, usage in enumerate(moe_stats['expert_usage']):
+        print(f"    Expert {i}: {usage.item():.1%}")
+
+    print("\n" + "=" * 70)
+    print("✅ INTEGRATION TEST COMPLETE!")
+    print("=" * 70)
+    print("\n💡 This system combines:")
+    print("  - MoE routing for intelligent control")
+    print("  - Adaptive budget for compute efficiency")
+    print("  - Multi-criteria convergence")
+    print("  - Phase-aware adaptation")
+    print("  - Budget redistribution")
+    print("  - Loss-component feedback")
+    print("\n🚀 Expected: 30-50% compute savings + 2-3x capacity!")

inl_llm/core/moe_controller.py

@@ -0,0 +1,618 @@
+"""
+Mixture of Experts (MoE) Controller for INL-LLM
+
+Implements intelligent routing between specialized expert controllers:
+- Multiple expert controllers, each learning different control strategies
+- Smart router that selects experts based on (h, x, layer, phase)
+- Sparse activation (top-k) for compute efficiency
+- Load balancing to prevent expert collapse
+- Automatic specialization emergence during training
+
+Key Features:
+✅ 4-8 specialized experts (automatic specialization)
+✅ Sparse routing (top-k): only activate 1-2 experts per forward
+✅ Context-aware routing (layer, phase, attention patterns)
+✅ Load balancing loss (prevent collapse)
+✅ 2-3x model capacity with 50% compute (vs dense)
+✅ Interpretable (can see which expert does what)
+
+Expected Specialization:
+- Expert 0: Fast convergence (early layers, equilibrium)
+- Expert 1: Complex reasoning (middle layers, high abstraction)
+- Expert 2: Stabilization (exploration phase, high noise)
+- Expert 3: Refinement (late layers, precision needed)
+
+Author: Boris Peyriguère
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, List, Optional, Tuple, Literal
+import math
+
+
+class ExpertController(nn.Module):
+    """
+    Single expert controller for INL dynamics.
+
+    Each expert learns specialized control strategies for different situations.
+    The specialization emerges naturally during training via the router.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        hidden_dim: int = 512,
+        expert_id: int = 0,
+        use_layer_norm: bool = True
+    ):
+        """
+        Args:
+            d_model: Model dimension
+            hidden_dim: Hidden layer dimension
+            expert_id: Expert identifier (for logging/debugging)
+            use_layer_norm: Use LayerNorm for stability
+        """
+        super().__init__()
+
+        self.d_model = d_model
+        self.expert_id = expert_id
+
+        # Fused controller MLP
+        self.mlp = nn.Sequential(
+            nn.Linear(2 * d_model, hidden_dim),
+            nn.LayerNorm(hidden_dim) if use_layer_norm else nn.Identity(),
+            nn.GELU(),
+            nn.Dropout(0.1),
+            nn.Linear(hidden_dim, 4 * d_model)
+        )
+
+        # Output heads for INL parameters
+        self.alpha_head = nn.Linear(d_model, d_model)
+        self.beta_head = nn.Linear(d_model, d_model)
+        self.gate_head = nn.Linear(d_model, d_model)
+        self.v_cand_head = nn.Linear(d_model, d_model)
+
+    def forward(self, h: torch.Tensor, x: torch.Tensor) -> Tuple[torch.Tensor, ...]:
+        """
+        Compute INL control parameters.
+
+        Args:
+            h: Context embedding [batch, d_model]
+            x: Current state [batch, d_model]
+
+        Returns:
+            alpha: Integration gain [batch, d_model]
+            beta: Error correction strength [batch, d_model]
+            gate: Velocity gating [batch, d_model]
+            v_cand: Candidate velocity [batch, d_model]
+        """
+        # Fused forward
+        combined = torch.cat([h, x], dim=-1)
+        output = self.mlp(combined)  # [batch, 4*d_model]
+
+        # Split into 4 parameter groups
+        alpha_feat, beta_feat, gate_feat, v_cand_feat = output.chunk(4, dim=-1)
+
+        # Apply output heads with appropriate activations
+        alpha = torch.sigmoid(self.alpha_head(alpha_feat))  # [0, 1] for momentum
+        beta = F.softplus(self.beta_head(beta_feat))        # [0, inf) for correction
+        gate = torch.sigmoid(self.gate_head(gate_feat))     # [0, 1] for gating
+        v_cand = self.v_cand_head(v_cand_feat)              # [-inf, inf] for velocity
+
+        return alpha, beta, gate, v_cand
+
+
+class INLMixtureOfExperts(nn.Module):
+    """
+    Mixture of Experts controller for INL-LLM.
+
+    Routes between multiple expert controllers based on:
+    - Input features (h, x)
+    - Layer depth (early/mid/late)
+    - Training phase (equilibrium/exploration)
+    - Attention patterns (optional)
+
+    Strategies:
+    - Sparse routing (top-k): Activate only k experts per forward
+    - Load balancing: Prevent expert collapse
+    - Context-aware: Router uses rich contextual features
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_layers: int,
+        num_experts: int = 4,
+        expert_hidden_dim: int = 512,
+        router_hidden_dim: int = 256,
+        top_k: int = 2,
+        use_sparse_routing: bool = True,
+        load_balance_weight: float = 0.01,
+        router_z_loss_weight: float = 0.001,
+        use_attention_features: bool = False
+    ):
+        """
+        Args:
+            d_model: Model dimension
+            num_layers: Number of layers in model
+            num_experts: Number of expert controllers (4-8 recommended)
+            expert_hidden_dim: Hidden dim for each expert
+            router_hidden_dim: Hidden dim for router network
+            top_k: Number of experts to activate per forward (1-2 for efficiency)
+            use_sparse_routing: Use top-k sparse routing vs dense
+            load_balance_weight: Weight for load balancing auxiliary loss
+            router_z_loss_weight: Weight for router z-loss (numerical stability)
+            use_attention_features: Use attention patterns in routing (experimental)
+        """
+        super().__init__()
+
+        self.d_model = d_model
+        self.num_layers = num_layers
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.use_sparse_routing = use_sparse_routing
+        self.load_balance_weight = load_balance_weight
+        self.router_z_loss_weight = router_z_loss_weight
+        self.use_attention_features = use_attention_features
+
+        # Expert controllers
+        self.experts = nn.ModuleList([
+            ExpertController(
+                d_model=d_model,
+                hidden_dim=expert_hidden_dim,
+                expert_id=i
+            )
+            for i in range(num_experts)
+        ])
+
+        # Context embeddings for router
+        self.layer_embeddings = nn.Embedding(num_layers, 32)
+        self.phase_embedding = nn.Embedding(2, 32)  # equilibrium=0, exploration=1
+
+        # Router network (chooses which experts to use)
+        router_input_dim = 2 * d_model + 64  # h + x + layer_emb + phase_emb
+        if use_attention_features:
+            router_input_dim += 32  # attention pattern features
+
+        self.router = nn.Sequential(
+            nn.Linear(router_input_dim, router_hidden_dim),
+            nn.LayerNorm(router_hidden_dim),
+            nn.GELU(),
+            nn.Dropout(0.1),
+            nn.Linear(router_hidden_dim, num_experts)
+        )
+
+        # Statistics tracking
+        self.register_buffer('expert_usage_history', torch.zeros(num_experts))
+        self.register_buffer('router_calls', torch.zeros(1))
+
+        # Jitter for load balancing (training only)
+        self.router_jitter_noise = 0.01
+
+    def forward(
+        self,
+        h: torch.Tensor,
+        x: torch.Tensor,
+        layer_idx: int,
+        phase: str = 'equilibrium',
+        attention_weights: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Dict]:
+        """
+        Forward pass through MoE controller.
+
+        Args:
+            h: Context embedding [batch, d_model]
+            x: Current state [batch, d_model]
+            layer_idx: Current layer index
+            phase: Training phase ('equilibrium' or 'exploration')
+            attention_weights: Optional attention pattern [batch, seq_len] for routing
+
+        Returns:
+            alpha: Integration gain [batch, d_model]
+            beta: Error correction strength [batch, d_model]
+            gate: Velocity gating [batch, d_model]
+            v_cand: Candidate velocity [batch, d_model]
+            info: Dictionary with routing statistics
+        """
+        batch_size = h.size(0)
+        device = h.device
+
+        # Prepare router input with contextual features
+        layer_emb = self.layer_embeddings(
+            torch.tensor([layer_idx], device=device)
+        ).expand(batch_size, -1)  # [batch, 32]
+
+        phase_idx = 0 if phase == 'equilibrium' else 1
+        phase_emb = self.phase_embedding(
+            torch.tensor([phase_idx], device=device)
+        ).expand(batch_size, -1)  # [batch, 32]
+
+        router_input = torch.cat([h, x, layer_emb, phase_emb], dim=-1)
+
+        # Optional: add attention pattern features
+        if self.use_attention_features and attention_weights is not None:
+            attn_features = self._extract_attention_features(attention_weights)
+            router_input = torch.cat([router_input, attn_features], dim=-1)
+
+        # Compute routing logits
+        router_logits = self.router(router_input)  # [batch, num_experts]
+
+        # Add jitter during training for load balancing
+        if self.training and self.router_jitter_noise > 0:
+            router_logits = router_logits + torch.randn_like(router_logits) * self.router_jitter_noise
+
+        # Route to experts
+        if self.use_sparse_routing:
+            alpha, beta, gate, v_cand, routing_info = self._sparse_forward(
+                h, x, router_logits
+            )
+        else:
+            alpha, beta, gate, v_cand, routing_info = self._dense_forward(
+                h, x, router_logits
+            )
+
+        # Compute auxiliary losses (training only)
+        aux_losses = {}
+        if self.training:
+            aux_losses['load_balance_loss'] = self._compute_load_balance_loss(
+                router_logits, routing_info['routing_weights']
+            )
+            aux_losses['router_z_loss'] = self._compute_router_z_loss(router_logits)
+
+        # Update statistics
+        self._update_statistics(routing_info['routing_weights'])
+
+        # Prepare info dict
+        info = {
+            **routing_info,
+            'aux_losses': aux_losses,
+            'expert_usage_history': self.expert_usage_history.clone(),
+            'num_experts': self.num_experts,
+            'top_k': self.top_k if self.use_sparse_routing else self.num_experts
+        }
+
+        return alpha, beta, gate, v_cand, info
+
+    def _sparse_forward(
+        self,
+        h: torch.Tensor,
+        x: torch.Tensor,
+        router_logits: torch.Tensor
+    ) -> Tuple[torch.Tensor, ...]:
+        """
+        Sparse forward: activate only top-k experts.
+
+        Compute efficiency: k/num_experts of full compute.
+        """
+        batch_size = h.size(0)
+
+        # Select top-k experts
+        top_k_logits, top_k_indices = torch.topk(
+            router_logits, self.top_k, dim=-1
+        )  # [batch, top_k], [batch, top_k]
+
+        # Normalize routing weights (softmax over selected experts only)
+        routing_weights = F.softmax(top_k_logits, dim=-1)  # [batch, top_k]
+
+        # Gather expert outputs for selected experts
+        # We need to process each sample's selected experts
+        alpha_list, beta_list, gate_list, v_cand_list = [], [], [], []
+
+        for b in range(batch_size):
+            sample_alphas, sample_betas, sample_gates, sample_v_cands = [], [], [], []
+
+            for k_idx in range(self.top_k):
+                expert_idx = top_k_indices[b, k_idx].item()
+                expert = self.experts[expert_idx]
+
+                # Run expert on this sample
+                alpha, beta, gate, v_cand = expert(h[b:b+1], x[b:b+1])
+
+                sample_alphas.append(alpha)
+                sample_betas.append(beta)
+                sample_gates.append(gate)
+                sample_v_cands.append(v_cand)
+
+            # Stack outputs for this sample
+            sample_alphas = torch.stack(sample_alphas, dim=0)  # [top_k, 1, d_model]
+            sample_betas = torch.stack(sample_betas, dim=0)
+            sample_gates = torch.stack(sample_gates, dim=0)
+            sample_v_cands = torch.stack(sample_v_cands, dim=0)
+
+            # Weighted combination for this sample
+            weights = routing_weights[b:b+1, :, None, None]  # [1, top_k, 1, 1]
+
+            alpha_combined = (weights * sample_alphas).sum(dim=1)  # [1, d_model]
+            beta_combined = (weights * sample_betas).sum(dim=1)
+            gate_combined = (weights * sample_gates).sum(dim=1)
+            v_cand_combined = (weights * sample_v_cands).sum(dim=1)
+
+            alpha_list.append(alpha_combined)
+            beta_list.append(beta_combined)
+            gate_list.append(gate_combined)
+            v_cand_list.append(v_cand_combined)
+
+        # Concatenate all samples
+        alpha = torch.cat(alpha_list, dim=0)  # [batch, d_model]
+        beta = torch.cat(beta_list, dim=0)
+        gate = torch.cat(gate_list, dim=0)
+        v_cand = torch.cat(v_cand_list, dim=0)
+
+        routing_info = {
+            'routing_weights': routing_weights,
+            'selected_experts': top_k_indices,
+            'router_logits': router_logits,
+            'routing_type': 'sparse'
+        }
+
+        return alpha, beta, gate, v_cand, routing_info
+
+    def _dense_forward(
+        self,
+        h: torch.Tensor,
+        x: torch.Tensor,
+        router_logits: torch.Tensor
+    ) -> Tuple[torch.Tensor, ...]:
+        """
+        Dense forward: use all experts (weighted combination).
+
+        Higher capacity but more compute.
+        """
+        # Compute routing weights (softmax over all experts)
+        routing_weights = F.softmax(router_logits, dim=-1)  # [batch, num_experts]
+
+        # Get all expert outputs
+        expert_outputs = []
+        for expert in self.experts:
+            alpha, beta, gate, v_cand = expert(h, x)
+            expert_outputs.append(
+                torch.stack([alpha, beta, gate, v_cand], dim=1)
+            )  # [batch, 4, d_model]
+
+        expert_outputs = torch.stack(expert_outputs, dim=1)  # [batch, num_experts, 4, d_model]
+
+        # Weighted combination
+        weights = routing_weights.unsqueeze(-1).unsqueeze(-1)  # [batch, num_experts, 1, 1]
+        combined = (weights * expert_outputs).sum(dim=1)  # [batch, 4, d_model]
+
+        # Split back
+        alpha, beta, gate, v_cand = combined.unbind(dim=1)
+
+        routing_info = {
+            'routing_weights': routing_weights,
+            'selected_experts': None,
+            'router_logits': router_logits,
+            'routing_type': 'dense'
+        }
+
+        return alpha, beta, gate, v_cand, routing_info
+
+    def _extract_attention_features(self, attention_weights: torch.Tensor) -> torch.Tensor:
+        """
+        Extract features from attention patterns for routing.
+
+        Args:
+            attention_weights: [batch, seq_len] or [batch, heads, seq_len]
+
+        Returns:
+            features: [batch, 32] attention pattern features
+        """
+        if attention_weights.dim() == 3:
+            # Average over heads
+            attention_weights = attention_weights.mean(dim=1)
+
+        # Compute attention statistics
+        attn_mean = attention_weights.mean(dim=-1, keepdim=True)
+        attn_max = attention_weights.max(dim=-1, keepdim=True)[0]
+        attn_std = attention_weights.std(dim=-1, keepdim=True)
+        attn_entropy = -(attention_weights * torch.log(attention_weights + 1e-10)).sum(dim=-1, keepdim=True)
+
+        # Simple MLP to project to 32 dims
+        features = torch.cat([attn_mean, attn_max, attn_std, attn_entropy], dim=-1)
+
+        # Expand to 32 dims (simple linear projection)
+        if not hasattr(self, 'attn_projector'):
+            self.attn_projector = nn.Linear(4, 32).to(features.device)
+
+        return self.attn_projector(features)
+
+    def _compute_load_balance_loss(
+        self,
+        router_logits: torch.Tensor,
+        routing_weights: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Auxiliary loss to encourage balanced expert usage.
+
+        Prevents collapse where model uses only 1-2 experts.
+        Based on: https://arxiv.org/abs/2101.03961 (Switch Transformers)
+        """
+        # Compute fraction of tokens routed to each expert
+        if self.use_sparse_routing:
+            # For sparse routing, count how many tokens go to each expert
+            # routing_weights: [batch, top_k]
+            # We need to map back to expert indices
+            batch_size = routing_weights.size(0)
+            expert_counts = torch.zeros(self.num_experts, device=routing_weights.device)
+
+            # This is approximate - just use router logits distribution
+            router_probs = F.softmax(router_logits, dim=-1)  # [batch, num_experts]
+            expert_usage = router_probs.mean(dim=0)  # [num_experts]
+        else:
+            # For dense routing, directly use routing weights
+            expert_usage = routing_weights.mean(dim=0)  # [num_experts]
+
+        # Target: uniform distribution
+        target = 1.0 / self.num_experts
+
+        # Coefficient of variation penalty
+        mean_usage = expert_usage.mean()
+        usage_variance = ((expert_usage - mean_usage) ** 2).mean()
+        cv_loss = usage_variance / (mean_usage + 1e-10)
+
+        return self.load_balance_weight * cv_loss
+
+    def _compute_router_z_loss(self, router_logits: torch.Tensor) -> torch.Tensor:
+        """
+        Router z-loss for numerical stability.
+
+        Penalizes large logits to prevent router from becoming too confident.
+        From: https://arxiv.org/abs/2202.08906
+        """
+        log_z = torch.logsumexp(router_logits, dim=-1)
+        z_loss = (log_z ** 2).mean()
+
+        return self.router_z_loss_weight * z_loss
+
+    def _update_statistics(self, routing_weights: torch.Tensor):
+        """Update running statistics of expert usage."""
+        if self.use_sparse_routing:
+            # Approximate from routing weights
+            # This is not perfect but gives an idea
+            usage = torch.zeros(self.num_experts, device=routing_weights.device)
+            # Just increment by batch size for now (rough approximation)
+            usage += routing_weights.size(0) / self.num_experts
+        else:
+            usage = routing_weights.sum(dim=0)  # [num_experts]
+
+        # Exponential moving average
+        alpha = 0.99
+        self.expert_usage_history = alpha * self.expert_usage_history + (1 - alpha) * usage
+        self.router_calls += 1
+
+    def get_expert_statistics(self) -> Dict[str, torch.Tensor]:
+        """
+        Get statistics about expert usage.
+
+        Returns:
+            Dictionary with expert usage statistics
+        """
+        # Normalize usage history
+        if self.router_calls > 0:
+            normalized_usage = self.expert_usage_history / self.expert_usage_history.sum()
+        else:
+            normalized_usage = torch.ones(self.num_experts) / self.num_experts
+
+        return {
+            'expert_usage': normalized_usage,
+            'expert_usage_raw': self.expert_usage_history,
+            'router_calls': self.router_calls,
+            'load_balance_score': self._compute_load_balance_score(normalized_usage)
+        }
+
+    def _compute_load_balance_score(self, usage: torch.Tensor) -> torch.Tensor:
+        """
+        Compute load balance score (1.0 = perfectly balanced).
+
+        Uses inverse of coefficient of variation.
+        """
+        target = 1.0 / self.num_experts
+        cv = usage.std() / (usage.mean() + 1e-10)
+        balance_score = 1.0 / (1.0 + cv)
+
+        return balance_score
+
+    def __repr__(self) -> str:
+        stats = self.get_expert_statistics()
+        balance_score = stats['load_balance_score'].item()
+
+        return (
+            f"INLMixtureOfExperts(\n"
+            f"  num_experts={self.num_experts},\n"
+            f"  top_k={self.top_k if self.use_sparse_routing else 'all'},\n"
+            f"  routing={'sparse' if self.use_sparse_routing else 'dense'},\n"
+            f"  load_balance_score={balance_score:.3f},\n"
+            f"  router_calls={int(self.router_calls.item())}\n"
+            f")"
+        )
+
+
+def create_moe_controller(
+    d_model: int,
+    num_layers: int,
+    num_experts: int = 4,
+    top_k: int = 2,
+    **kwargs
+) -> INLMixtureOfExperts:
+    """
+    Helper function to create MoE controller with sensible defaults.
+
+    Args:
+        d_model: Model dimension
+        num_layers: Number of layers
+        num_experts: Number of expert controllers (4-8 recommended)
+        top_k: Number of experts to activate (1-2 for efficiency)
+        **kwargs: Additional arguments for INLMixtureOfExperts
+
+    Returns:
+        Configured INLMixtureOfExperts controller
+    """
+    return INLMixtureOfExperts(
+        d_model=d_model,
+        num_layers=num_layers,
+        num_experts=num_experts,
+        top_k=top_k,
+        **kwargs
+    )
+
+
+if __name__ == '__main__':
+    print("=" * 70)
+    print("MIXTURE OF EXPERTS CONTROLLER - Test")
+    print("=" * 70)
+
+    # Configuration
+    d_model = 1024
+    num_layers = 25
+    batch_size = 16
+
+    # Create MoE controller
+    moe = create_moe_controller(
+        d_model=d_model,
+        num_layers=num_layers,
+        num_experts=4,
+        top_k=2,
+        use_sparse_routing=True
+    )
+
+    print(f"\n{moe}")
+
+    # Test forward pass
+    print("\n🧪 Testing forward pass...")
+    h = torch.randn(batch_size, d_model)
+    x = torch.randn(batch_size, d_model)
+
+    # Test different layers and phases
+    test_configs = [
+        (0, 'equilibrium'),
+        (12, 'equilibrium'),
+        (24, 'equilibrium'),
+        (12, 'exploration')
+    ]
+
+    print("\n📊 Routing Analysis:")
+    for layer_idx, phase in test_configs:
+        alpha, beta, gate, v_cand, info = moe(h, x, layer_idx, phase)
+
+        print(f"\n  Layer {layer_idx:2d} ({phase}):")
+        print(f"    Output shapes: alpha={alpha.shape}, beta={beta.shape}")
+        print(f"    Routing type: {info['routing_type']}")
+        print(f"    Selected experts (sample 0): {info['selected_experts'][0].tolist()}")
+        print(f"    Routing weights (sample 0): {info['routing_weights'][0].tolist()}")
+
+        if 'aux_losses' in info:
+            print(f"    Load balance loss: {info['aux_losses']['load_balance_loss']:.6f}")
+            print(f"    Router z-loss: {info['aux_losses']['router_z_loss']:.6f}")
+
+    # Expert usage statistics
+    print("\n📈 Expert Usage Statistics:")
+    stats = moe.get_expert_statistics()
+    for i, usage in enumerate(stats['expert_usage']):
+        print(f"  Expert {i}: {usage.item():.1%}")
+    print(f"  Load Balance Score: {stats['load_balance_score'].item():.3f}")
+
+    print("\n" + "=" * 70)
+    print("✅ MoE CONTROLLER TEST COMPLETE!")
+    print("=" * 70)

inl_llm/models/__init__.py

@@ -1,31 +1,31 @@
-"""
-Complete INL-LLM model with all optimizations (Level 1 + 2).
-
-Single production-ready model with maximum efficiency.
-"""
-
-from .integrator_language_model import (
-    UltraOptimizedIntegratorLanguageModel,
-    create_ultra_optimized_model
-)
-
-# HuggingFace-compatible wrappers (for vLLM support)
-from .modeling_inl_llm import (
-    INLLLMConfig,
-    INLLLMForCausalLM
-)
-
-# Aliases for simpler API
-IntegratorLanguageModel = UltraOptimizedIntegratorLanguageModel
-create_model = create_ultra_optimized_model
-
-__all__ = [
-    'IntegratorLanguageModel',
-    'create_model',
-    # Legacy aliases
-    'UltraOptimizedIntegratorLanguageModel',
-    'create_ultra_optimized_model',
-    # HuggingFace compatibility
-    'INLLLMConfig',
-    'INLLLMForCausalLM'
-]
+"""
+Complete INL-LLM model with all optimizations (Level 1 + 2).
+
+Single production-ready model with maximum efficiency.
+"""
+
+from .integrator_language_model import (
+    UltraOptimizedIntegratorLanguageModel,
+    create_ultra_optimized_model
+)
+
+# HuggingFace-compatible wrappers (for vLLM support)
+from .modeling_inl_llm import (
+    INLLLMConfig,
+    INLLLMForCausalLM
+)
+
+# Aliases for simpler API
+IntegratorLanguageModel = UltraOptimizedIntegratorLanguageModel
+create_model = create_ultra_optimized_model
+
+__all__ = [
+    'IntegratorLanguageModel',
+    'create_model',
+    # Legacy aliases
+    'UltraOptimizedIntegratorLanguageModel',
+    'create_ultra_optimized_model',
+    # HuggingFace compatibility
+    'INLLLMConfig',
+    'INLLLMForCausalLM'
+]

inl_llm/models/inl_diffusion.py

@@ -1,814 +1,814 @@
-"""
-INL-Diffusion: Latent Diffusion Model with Integrator Neuron dynamics
-
-A text-to-image generation model inspired by Stable Diffusion but using
-INL dynamics instead of standard transformers.
-
-Architecture:
-1. VAE: Encode images to latent space (compress 512x512 -> 64x64x4)
-2. Text Encoder: Encode text prompts to embeddings
-3. U-Net with INL blocks: Denoise latent representations conditioned on text
-4. VAE Decoder: Decode latents back to images
-
-Author: Boris Peyriguère
-"""
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from typing import Optional, Tuple, List
-import math
-
-from .inl_vision import SimpleINLDynamics
-
-
-class TimeEmbedding(nn.Module):
-    """
-    Sinusoidal time embedding for diffusion timesteps.
-    """
-    def __init__(self, dim: int):
-        super().__init__()
-        self.dim = dim
-
-    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            timesteps: (B,) tensor of timestep indices
-
-        Returns:
-            embeddings: (B, dim) time embeddings
-        """
-        half_dim = self.dim // 2
-        emb = math.log(10000) / (half_dim - 1)
-        emb = torch.exp(torch.arange(half_dim, device=timesteps.device) * -emb)
-        emb = timesteps[:, None] * emb[None, :]
-        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
-        return emb
-
-
-class ResnetBlock(nn.Module):
-    """
-    Residual block for U-Net with time conditioning.
-    """
-    def __init__(self, in_channels: int, out_channels: int, time_emb_dim: int):
-        super().__init__()
-
-        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
-        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
-
-        self.time_mlp = nn.Sequential(
-            nn.SiLU(),
-            nn.Linear(time_emb_dim, out_channels)
-        )
-
-        self.norm1 = nn.GroupNorm(8, in_channels)
-        self.norm2 = nn.GroupNorm(8, out_channels)
-
-        self.act = nn.SiLU()
-
-        if in_channels != out_channels:
-            self.shortcut = nn.Conv2d(in_channels, out_channels, 1)
-        else:
-            self.shortcut = nn.Identity()
-
-    def forward(self, x: torch.Tensor, time_emb: torch.Tensor) -> torch.Tensor:
-        h = self.norm1(x)
-        h = self.act(h)
-        h = self.conv1(h)
-
-        # Add time conditioning
-        time_cond = self.time_mlp(time_emb)[:, :, None, None]
-        h = h + time_cond
-
-        h = self.norm2(h)
-        h = self.act(h)
-        h = self.conv2(h)
-
-        return h + self.shortcut(x)
-
-
-class INLAttentionBlock(nn.Module):
-    """
-    Attention block using INL dynamics for refinement.
-    """
-    def __init__(self, channels: int, num_heads: int = 8, num_iterations: int = 3):
-        super().__init__()
-
-        self.channels = channels
-        self.num_heads = num_heads
-        self.norm = nn.GroupNorm(8, channels)
-
-        self.qkv = nn.Conv2d(channels, channels * 3, 1)
-        self.proj_out = nn.Conv2d(channels, channels, 1)
-
-        # INL dynamics for iterative refinement
-        self.inl = SimpleINLDynamics(
-            d_model=channels,
-            num_iterations=num_iterations,
-            dt=0.1
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        B, C, H, W = x.shape
-        h = self.norm(x)
-
-        # QKV projection
-        qkv = self.qkv(h)
-        q, k, v = torch.chunk(qkv, 3, dim=1)
-
-        # Reshape for attention
-        q = q.reshape(B, self.num_heads, C // self.num_heads, H * W).transpose(-1, -2)
-        k = k.reshape(B, self.num_heads, C // self.num_heads, H * W).transpose(-1, -2)
-        v = v.reshape(B, self.num_heads, C // self.num_heads, H * W).transpose(-1, -2)
-
-        # Attention
-        scale = (C // self.num_heads) ** -0.5
-        attn = torch.softmax(q @ k.transpose(-1, -2) * scale, dim=-1)
-        h = attn @ v
-
-        # Reshape back
-        h = h.transpose(-1, -2).reshape(B, C, H, W)
-
-        # Apply INL dynamics for refinement
-        h_flat = h.reshape(B, C, H * W).transpose(1, 2)  # (B, H*W, C)
-        h_refined = self.inl(h_flat)
-        h = h_refined.transpose(1, 2).reshape(B, C, H, W)
-
-        h = self.proj_out(h)
-
-        return x + h
-
-
-class INLUNet(nn.Module):
-    """
-    U-Net with INL dynamics for latent diffusion.
-
-    Denoises latent representations conditioned on text embeddings.
-    """
-    def __init__(
-        self,
-        in_channels: int = 4,
-        out_channels: int = 4,
-        model_channels: int = 320,
-        num_res_blocks: int = 2,
-        attention_resolutions: List[int] = [4, 2, 1],
-        channel_mult: List[int] = [1, 2, 4, 4],
-        num_heads: int = 8,
-        context_dim: int = 768  # Text embedding dimension
-    ):
-        super().__init__()
-
-        self.in_channels = in_channels
-        self.model_channels = model_channels
-
-        # Time embedding
-        time_embed_dim = model_channels * 4
-        self.time_embed = nn.Sequential(
-            TimeEmbedding(model_channels),
-            nn.Linear(model_channels, time_embed_dim),
-            nn.SiLU(),
-            nn.Linear(time_embed_dim, time_embed_dim)
-        )
-
-        # Text conditioning projection
-        self.context_proj = nn.Linear(context_dim, model_channels)
-
-        # Input convolution
-        self.input_blocks = nn.ModuleList([
-            nn.Conv2d(in_channels, model_channels, 3, padding=1)
-        ])
-
-        # Encoder (downsampling)
-        ch = model_channels
-        input_block_chans = [ch]
-
-        for level, mult in enumerate(channel_mult):
-            for _ in range(num_res_blocks):
-                layers = [
-                    ResnetBlock(ch, mult * model_channels, time_embed_dim)
-                ]
-
-                ch = mult * model_channels
-
-                # Add attention at specified resolutions
-                if level in attention_resolutions:
-                    layers.append(INLAttentionBlock(ch, num_heads))
-
-                self.input_blocks.append(nn.Sequential(*layers))
-                input_block_chans.append(ch)
-
-            # Downsample
-            if level != len(channel_mult) - 1:
-                self.input_blocks.append(nn.Conv2d(ch, ch, 3, stride=2, padding=1))
-                input_block_chans.append(ch)
-
-        # Middle
-        self.middle_block = nn.Sequential(
-            ResnetBlock(ch, ch, time_embed_dim),
-            INLAttentionBlock(ch, num_heads, num_iterations=5),
-            ResnetBlock(ch, ch, time_embed_dim)
-        )
-
-        # Decoder (upsampling)
-        self.output_blocks = nn.ModuleList([])
-
-        for level, mult in list(enumerate(channel_mult))[::-1]:
-            for i in range(num_res_blocks + 1):
-                ich = input_block_chans.pop()
-                layers = [
-                    ResnetBlock(ch + ich, mult * model_channels, time_embed_dim)
-                ]
-
-                ch = mult * model_channels
-
-                if level in attention_resolutions:
-                    layers.append(INLAttentionBlock(ch, num_heads))
-
-                # Upsample
-                if level != 0 and i == num_res_blocks:
-                    layers.append(nn.Upsample(scale_factor=2, mode='nearest'))
-
-                self.output_blocks.append(nn.Sequential(*layers))
-
-        # Output
-        self.out = nn.Sequential(
-            nn.GroupNorm(8, ch),
-            nn.SiLU(),
-            nn.Conv2d(ch, out_channels, 3, padding=1)
-        )
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        timesteps: torch.Tensor,
-        context: Optional[torch.Tensor] = None
-    ) -> torch.Tensor:
-        """
-        Args:
-            x: Noisy latents (B, 4, H, W)
-            timesteps: Diffusion timesteps (B,)
-            context: Text embeddings (B, seq_len, context_dim)
-
-        Returns:
-            Predicted noise (B, 4, H, W)
-        """
-        # Time embedding
-        t_emb = self.time_embed(timesteps)
-
-        # Text conditioning (average pooling for simplicity)
-        if context is not None:
-            context = context.mean(dim=1)  # (B, context_dim)
-            context_emb = self.context_proj(context)
-            t_emb = t_emb + context_emb
-
-        # Encoder
-        hs = []
-        h = x
-        for module in self.input_blocks:
-            if isinstance(module, nn.Sequential):
-                for layer in module:
-                    if isinstance(layer, ResnetBlock):
-                        h = layer(h, t_emb)
-                    else:
-                        h = layer(h)
-            else:
-                h = module(h)
-            hs.append(h)
-
-        # Middle
-        for layer in self.middle_block:
-            if isinstance(layer, ResnetBlock):
-                h = layer(h, t_emb)
-            else:
-                h = layer(h)
-
-        # Decoder
-        for module in self.output_blocks:
-            h = torch.cat([h, hs.pop()], dim=1)
-            for layer in module:
-                if isinstance(layer, ResnetBlock):
-                    h = layer(h, t_emb)
-                else:
-                    h = layer(h)
-
-        # Output
-        return self.out(h)
-
-
-class VAEResBlock(nn.Module):
-    """
-    Residual block for VAE with GroupNorm.
-    Similar to Stable Diffusion VAE architecture.
-    """
-    def __init__(self, in_channels: int, out_channels: int):
-        super().__init__()
-
-        self.norm1 = nn.GroupNorm(32, in_channels)
-        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
-
-        self.norm2 = nn.GroupNorm(32, out_channels)
-        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
-
-        if in_channels != out_channels:
-            self.shortcut = nn.Conv2d(in_channels, out_channels, 1)
-        else:
-            self.shortcut = nn.Identity()
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        h = self.norm1(x)
-        h = F.silu(h)
-        h = self.conv1(h)
-
-        h = self.norm2(h)
-        h = F.silu(h)
-        h = self.conv2(h)
-
-        return h + self.shortcut(x)
-
-
-class VAEAttentionBlock(nn.Module):
-    """
-    Self-attention block for VAE.
-    """
-    def __init__(self, channels: int):
-        super().__init__()
-
-        self.channels = channels
-        self.norm = nn.GroupNorm(32, channels)
-
-        self.q = nn.Conv2d(channels, channels, 1)
-        self.k = nn.Conv2d(channels, channels, 1)
-        self.v = nn.Conv2d(channels, channels, 1)
-
-        self.proj_out = nn.Conv2d(channels, channels, 1)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        B, C, H, W = x.shape
-        h = self.norm(x)
-
-        q = self.q(h).reshape(B, C, H * W).transpose(1, 2)  # (B, HW, C)
-        k = self.k(h).reshape(B, C, H * W).transpose(1, 2)
-        v = self.v(h).reshape(B, C, H * W).transpose(1, 2)
-
-        # Attention
-        scale = C ** -0.5
-        attn = torch.softmax(q @ k.transpose(-1, -2) * scale, dim=-1)
-        h = attn @ v
-
-        # Reshape back
-        h = h.transpose(1, 2).reshape(B, C, H, W)
-        h = self.proj_out(h)
-
-        return x + h
-
-
-class Downsample(nn.Module):
-    """Downsampling layer."""
-    def __init__(self, channels: int):
-        super().__init__()
-        self.conv = nn.Conv2d(channels, channels, 3, stride=2, padding=0)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # Asymmetric padding to match Stable Diffusion
-        x = F.pad(x, (0, 1, 0, 1), mode='constant', value=0)
-        return self.conv(x)
-
-
-class Upsample(nn.Module):
-    """Upsampling layer."""
-    def __init__(self, channels: int):
-        super().__init__()
-        self.conv = nn.Conv2d(channels, channels, 3, padding=1)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = F.interpolate(x, scale_factor=2.0, mode='nearest')
-        return self.conv(x)
-
-
-class StableDiffusionVAE(nn.Module):
-    """
-    MASSIVE VAE for ultra-high quality latent diffusion.
-
-    Architecture:
-    - Input: 256x256x3 (or 512x512x3)
-    - Latent: 32x32x4 (8x downsampling)
-    - Deep ResNet blocks with GroupNorm
-    - Multi-head attention at multiple resolutions
-    - ~2.15B parameters for SOTA reconstruction quality
-
-    This beast preserves ALL details with near-perfect reconstruction.
-
-    Memory optimization:
-    - Use gradient_checkpointing=True to reduce memory by ~70% (trades 25% speed)
-    - Essential for training on GPUs < 24GB
-    """
-    def __init__(
-        self,
-        in_channels: int = 3,
-        latent_channels: int = 4,
-        base_channels: int = 256,  # 128 -> 256 (doubled!)
-        channel_multipliers: List[int] = [1, 2, 4, 8],  # [1,2,4,4] -> [1,2,4,8] (doubled max!)
-        num_res_blocks: int = 6,  # 2 -> 6 (tripled!)
-        attn_resolutions: List[int] = [128, 64, 32],  # More attention layers!
-        use_gradient_checkpointing: bool = False  # Enable for memory-constrained GPUs
-    ):
-        super().__init__()
-
-        self.latent_channels = latent_channels
-        self.use_gradient_checkpointing = use_gradient_checkpointing
-
-        # ========== ENCODER ==========
-        # Input: 256x256x3 -> 256x256x256
-        self.encoder_conv_in = nn.Conv2d(in_channels, base_channels, 3, padding=1)
-
-        # Downsampling blocks
-        self.encoder_blocks = nn.ModuleList()
-        ch = base_channels
-        resolutions = []
-        current_res = 256  # Assume 256x256 input
-
-        for level, mult in enumerate(channel_multipliers):
-            out_ch = base_channels * mult
-
-            # Add MANY residual blocks (6 per level!)
-            for _ in range(num_res_blocks):
-                self.encoder_blocks.append(VAEResBlock(ch, out_ch))
-                ch = out_ch
-                resolutions.append(current_res)
-
-            # Add attention at specified resolutions
-            if current_res in attn_resolutions:
-                # Add MULTIPLE attention blocks for better quality
-                self.encoder_blocks.append(VAEAttentionBlock(ch))
-                self.encoder_blocks.append(VAEAttentionBlock(ch))
-                resolutions.append(current_res)
-                resolutions.append(current_res)
-
-            # Downsample (except last level)
-            if level != len(channel_multipliers) - 1:
-                self.encoder_blocks.append(Downsample(ch))
-                current_res //= 2
-                resolutions.append(current_res)
-
-        # Middle blocks (at 32x32x2048!) - MASSIVE bottleneck
-        self.encoder_mid_block1 = VAEResBlock(ch, ch)
-        self.encoder_mid_attn1 = VAEAttentionBlock(ch)
-        self.encoder_mid_block2 = VAEResBlock(ch, ch)
-        self.encoder_mid_attn2 = VAEAttentionBlock(ch)
-        self.encoder_mid_block3 = VAEResBlock(ch, ch)
-        self.encoder_mid_attn3 = VAEAttentionBlock(ch)
-        self.encoder_mid_block4 = VAEResBlock(ch, ch)
-
-        # Output: mu and logvar
-        self.encoder_norm_out = nn.GroupNorm(32, ch)
-        self.encoder_conv_out = nn.Conv2d(ch, latent_channels * 2, 3, padding=1)
-
-        # ========== DECODER ==========
-        # Input: 32x32x4 -> 32x32x2048
-        self.decoder_conv_in = nn.Conv2d(latent_channels, ch, 3, padding=1)
-
-        # Middle blocks - MASSIVE processing
-        self.decoder_mid_block1 = VAEResBlock(ch, ch)
-        self.decoder_mid_attn1 = VAEAttentionBlock(ch)
-        self.decoder_mid_block2 = VAEResBlock(ch, ch)
-        self.decoder_mid_attn2 = VAEAttentionBlock(ch)
-        self.decoder_mid_block3 = VAEResBlock(ch, ch)
-        self.decoder_mid_attn3 = VAEAttentionBlock(ch)
-        self.decoder_mid_block4 = VAEResBlock(ch, ch)
-
-        # Upsampling blocks
-        self.decoder_blocks = nn.ModuleList()
-
-        for level, mult in reversed(list(enumerate(channel_multipliers))):
-            out_ch = base_channels * mult
-
-            # MANY residual blocks per level
-            for _ in range(num_res_blocks + 1):
-                self.decoder_blocks.append(VAEResBlock(ch, out_ch))
-                ch = out_ch
-
-            # Add attention at specified resolutions
-            if current_res in attn_resolutions:
-                # Multiple attention blocks
-                self.decoder_blocks.append(VAEAttentionBlock(ch))
-                self.decoder_blocks.append(VAEAttentionBlock(ch))
-
-            # Upsample (except first level, which is last in reversed order)
-            if level != 0:
-                self.decoder_blocks.append(Upsample(ch))
-                current_res *= 2
-
-        # Output: 256x256x3
-        self.decoder_norm_out = nn.GroupNorm(32, ch)
-        self.decoder_conv_out = nn.Conv2d(ch, in_channels, 3, padding=1)
-
-    def encode(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Encode image to latent distribution parameters.
-
-        Args:
-            x: (B, 3, H, W) images in range [-1, 1]
-
-        Returns:
-            mu: (B, latent_channels, H/8, W/8)
-            logvar: (B, latent_channels, H/8, W/8)
-        """
-        # Input conv
-        h = self.encoder_conv_in(x)
-
-        # Encoder blocks
-        for block in self.encoder_blocks:
-            h = block(h)
-
-        # Middle - DEEP processing
-        h = self.encoder_mid_block1(h)
-        h = self.encoder_mid_attn1(h)
-        h = self.encoder_mid_block2(h)
-        h = self.encoder_mid_attn2(h)
-        h = self.encoder_mid_block3(h)
-        h = self.encoder_mid_attn3(h)
-        h = self.encoder_mid_block4(h)
-
-        # Output
-        h = self.encoder_norm_out(h)
-        h = F.silu(h)
-        h = self.encoder_conv_out(h)
-
-        mu, logvar = torch.chunk(h, 2, dim=1)
-        return mu, logvar
-
-    def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
-        """Sample from latent distribution."""
-        std = torch.exp(0.5 * logvar)
-        eps = torch.randn_like(std)
-        return mu + eps * std
-
-    def decode(self, z: torch.Tensor) -> torch.Tensor:
-        """
-        Decode latent to image.
-
-        Args:
-            z: (B, latent_channels, H/8, W/8) latent codes
-
-        Returns:
-            x: (B, 3, H, W) reconstructed images in range [-1, 1]
-        """
-        # Input conv
-        h = self.decoder_conv_in(z)
-
-        # Middle - DEEP processing
-        h = self.decoder_mid_block1(h)
-        h = self.decoder_mid_attn1(h)
-        h = self.decoder_mid_block2(h)
-        h = self.decoder_mid_attn2(h)
-        h = self.decoder_mid_block3(h)
-        h = self.decoder_mid_attn3(h)
-        h = self.decoder_mid_block4(h)
-
-        # Decoder blocks
-        for block in self.decoder_blocks:
-            h = block(h)
-
-        # Output
-        h = self.decoder_norm_out(h)
-        h = F.silu(h)
-        h = self.decoder_conv_out(h)
-
-        return torch.tanh(h)
-
-    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Full forward pass: encode -> sample -> decode."""
-        mu, logvar = self.encode(x)
-        z = self.reparameterize(mu, logvar)
-        recon = self.decode(z)
-        return recon, mu, logvar
-
-
-class SimpleVAE(nn.Module):
-    """
-    DEPRECATED: Simple VAE with only 2M parameters.
-    Use StableDiffusionVAE instead for production.
-
-    Simple VAE for encoding images to latent space.
-    Compress 512x512x3 -> 64x64x4
-    """
-    def __init__(self, in_channels: int = 3, latent_channels: int = 4):
-        super().__init__()
-
-        # Encoder (512 -> 64, 8x downsampling)
-        self.encoder = nn.Sequential(
-            nn.Conv2d(in_channels, 128, 3, padding=1),
-            nn.ReLU(),
-            nn.Conv2d(128, 128, 3, stride=2, padding=1),  # 256
-            nn.ReLU(),
-            nn.Conv2d(128, 256, 3, stride=2, padding=1),  # 128
-            nn.ReLU(),
-            nn.Conv2d(256, 256, 3, stride=2, padding=1),  # 64
-            nn.ReLU(),
-            nn.Conv2d(256, latent_channels * 2, 3, padding=1)  # mu, logvar
-        )
-
-        # Decoder (64 -> 512)
-        self.decoder = nn.Sequential(
-            nn.Conv2d(latent_channels, 256, 3, padding=1),
-            nn.ReLU(),
-            nn.Upsample(scale_factor=2, mode='nearest'),  # 128
-            nn.Conv2d(256, 256, 3, padding=1),
-            nn.ReLU(),
-            nn.Upsample(scale_factor=2, mode='nearest'),  # 256
-            nn.Conv2d(256, 128, 3, padding=1),
-            nn.ReLU(),
-            nn.Upsample(scale_factor=2, mode='nearest'),  # 512
-            nn.Conv2d(128, 128, 3, padding=1),
-            nn.ReLU(),
-            nn.Conv2d(128, in_channels, 3, padding=1),
-            nn.Tanh()
-        )
-
-    def encode(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        h = self.encoder(x)
-        mu, logvar = torch.chunk(h, 2, dim=1)
-        return mu, logvar
-
-    def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
-        std = torch.exp(0.5 * logvar)
-        eps = torch.randn_like(std)
-        return mu + eps * std
-
-    def decode(self, z: torch.Tensor) -> torch.Tensor:
-        return self.decoder(z)
-
-    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        mu, logvar = self.encode(x)
-        z = self.reparameterize(mu, logvar)
-        recon = self.decode(z)
-        return recon, mu, logvar
-
-
-class INLTextEncoder(nn.Module):
-    """
-    Text encoder using pre-trained INL-LLM model.
-
-    Reuses the trained INL-LLM (1.1B) as a powerful text encoder
-    with integrator neuron dynamics.
-    """
-    def __init__(self, inl_llm_model, embed_dim: int = 768):
-        super().__init__()
-
-        # Use pretrained INL-LLM
-        self.inl_llm = inl_llm_model
-
-        # Freeze INL-LLM (use as feature extractor)
-        for param in self.inl_llm.parameters():
-            param.requires_grad = False
-
-        # Project INL-LLM hidden states to diffusion context dimension
-        llm_hidden_dim = self.inl_llm.d_model
-        self.projection = nn.Linear(llm_hidden_dim, embed_dim)
-
-    def forward(self, text_tokens: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            text_tokens: (B, seq_len) token IDs
-
-        Returns:
-            text_embeddings: (B, seq_len, embed_dim)
-        """
-        with torch.no_grad():
-            # Get hidden states from INL-LLM (no generation, just encoding)
-            # Use the model's embedding + transformer blocks
-            x = self.inl_llm.token_embedding(text_tokens)
-            x = self.inl_llm.pos_encoding(x)
-
-            # Pass through INL layers to get contextualized representations
-            for layer in self.inl_llm.layers:
-                x, _ = layer(x)
-
-            x = self.inl_llm.norm(x)
-
-        # Project to context dimension
-        x = self.projection(x)
-
-        return x
-
-
-class INLLatentDiffusion(nn.Module):
-    """
-    Complete Latent Diffusion Model with INL dynamics.
-
-    Text → Image generation pipeline.
-    """
-    def __init__(
-        self,
-        img_size: int = 512,
-        latent_size: int = 64,
-        inl_llm_model = None,  # Pre-trained INL-LLM for text encoding
-        context_dim: int = 768
-    ):
-        super().__init__()
-
-        self.img_size = img_size
-        self.latent_size = latent_size
-
-        # Components - Use MASSIVE 1B+ parameter VAE
-        self.vae = StableDiffusionVAE(
-            in_channels=3,
-            latent_channels=4,
-            base_channels=256,
-            channel_multipliers=[1, 2, 4, 8],
-            num_res_blocks=6,
-            attn_resolutions=[128, 64, 32]
-        )
-        print(f"✅ Using StableDiffusionVAE with {sum(p.numel() for p in self.vae.parameters()):,} parameters")
-
-        # Use INL-LLM as text encoder if provided
-        if inl_llm_model is not None:
-            self.text_encoder = INLTextEncoder(inl_llm_model, embed_dim=context_dim)
-            print("✅ Using pre-trained INL-LLM as text encoder (frozen)")
-        else:
-            # Fallback to simple encoder
-            print("⚠️ No INL-LLM provided, using simple text encoder")
-            from .integrator_language_model import UltraOptimizedIntegratorLanguageModel
-            # Create a small text encoder
-            small_llm = UltraOptimizedIntegratorLanguageModel(
-                vocab_size=50000,
-                d_model=512,
-                num_layers=6,
-                num_heads=8,
-                num_iterations_per_layer=3
-            )
-            self.text_encoder = INLTextEncoder(small_llm, embed_dim=context_dim)
-
-        self.unet = INLUNet(
-            in_channels=4,
-            out_channels=4,
-            model_channels=320,
-            context_dim=context_dim
-        )
-
-        # Diffusion parameters
-        self.num_timesteps = 1000
-        self.register_buffer('betas', self._cosine_beta_schedule())
-        self.register_buffer('alphas', 1.0 - self.betas)
-        self.register_buffer('alphas_cumprod', torch.cumprod(self.alphas, dim=0))
-
-    def _cosine_beta_schedule(self, s: float = 0.008) -> torch.Tensor:
-        """Cosine schedule from Improved DDPM."""
-        steps = self.num_timesteps + 1
-        x = torch.linspace(0, self.num_timesteps, steps)
-        alphas_cumprod = torch.cos(((x / self.num_timesteps) + s) / (1 + s) * math.pi * 0.5) ** 2
-        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
-        betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
-        return torch.clip(betas, 0.0001, 0.9999)
-
-    @torch.no_grad()
-    def generate(
-        self,
-        text_tokens: torch.Tensor,
-        num_inference_steps: int = 50,
-        guidance_scale: float = 7.5
-    ) -> torch.Tensor:
-        """
-        Generate images from text prompts.
-
-        Args:
-            text_tokens: (B, seq_len) text token IDs
-            num_inference_steps: Number of denoising steps
-            guidance_scale: Classifier-free guidance scale
-
-        Returns:
-            generated_images: (B, 3, img_size, img_size)
-        """
-        B = text_tokens.size(0)
-        device = text_tokens.device
-
-        # Encode text
-        context = self.text_encoder(text_tokens)
-
-        # Start from random noise
-        latents = torch.randn(B, 4, self.latent_size, self.latent_size, device=device)
-
-        # Denoising loop (DDPM sampling)
-        timesteps = torch.linspace(self.num_timesteps - 1, 0, num_inference_steps, dtype=torch.long, device=device)
-
-        for t in timesteps:
-            t_batch = t.repeat(B)
-
-            # Predict noise
-            noise_pred = self.unet(latents, t_batch, context)
-
-            # Update latents (simplified DDPM step)
-            alpha = self.alphas_cumprod[t]
-            alpha_prev = self.alphas_cumprod[t - 1] if t > 0 else torch.tensor(1.0, device=device)
-
-            beta_t = 1 - alpha / alpha_prev
-            latents = (latents - beta_t * noise_pred) / torch.sqrt(1 - beta_t)
-
-        # Decode latents to images
-        images = self.vae.decode(latents)
-
-        return images
-
-    def get_num_params(self):
-        """Count total parameters."""
-        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+"""
+INL-Diffusion: Latent Diffusion Model with Integrator Neuron dynamics
+
+A text-to-image generation model inspired by Stable Diffusion but using
+INL dynamics instead of standard transformers.
+
+Architecture:
+1. VAE: Encode images to latent space (compress 512x512 -> 64x64x4)
+2. Text Encoder: Encode text prompts to embeddings
+3. U-Net with INL blocks: Denoise latent representations conditioned on text
+4. VAE Decoder: Decode latents back to images
+
+Author: Boris Peyriguère
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List
+import math
+
+from .inl_vision import SimpleINLDynamics
+
+
+class TimeEmbedding(nn.Module):
+    """
+    Sinusoidal time embedding for diffusion timesteps.
+    """
+    def __init__(self, dim: int):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            timesteps: (B,) tensor of timestep indices
+
+        Returns:
+            embeddings: (B, dim) time embeddings
+        """
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=timesteps.device) * -emb)
+        emb = timesteps[:, None] * emb[None, :]
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+        return emb
+
+
+class ResnetBlock(nn.Module):
+    """
+    Residual block for U-Net with time conditioning.
+    """
+    def __init__(self, in_channels: int, out_channels: int, time_emb_dim: int):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
+
+        self.time_mlp = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(time_emb_dim, out_channels)
+        )
+
+        self.norm1 = nn.GroupNorm(8, in_channels)
+        self.norm2 = nn.GroupNorm(8, out_channels)
+
+        self.act = nn.SiLU()
+
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv2d(in_channels, out_channels, 1)
+        else:
+            self.shortcut = nn.Identity()
+
+    def forward(self, x: torch.Tensor, time_emb: torch.Tensor) -> torch.Tensor:
+        h = self.norm1(x)
+        h = self.act(h)
+        h = self.conv1(h)
+
+        # Add time conditioning
+        time_cond = self.time_mlp(time_emb)[:, :, None, None]
+        h = h + time_cond
+
+        h = self.norm2(h)
+        h = self.act(h)
+        h = self.conv2(h)
+
+        return h + self.shortcut(x)
+
+
+class INLAttentionBlock(nn.Module):
+    """
+    Attention block using INL dynamics for refinement.
+    """
+    def __init__(self, channels: int, num_heads: int = 8, num_iterations: int = 3):
+        super().__init__()
+
+        self.channels = channels
+        self.num_heads = num_heads
+        self.norm = nn.GroupNorm(8, channels)
+
+        self.qkv = nn.Conv2d(channels, channels * 3, 1)
+        self.proj_out = nn.Conv2d(channels, channels, 1)
+
+        # INL dynamics for iterative refinement
+        self.inl = SimpleINLDynamics(
+            d_model=channels,
+            num_iterations=num_iterations,
+            dt=0.1
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, C, H, W = x.shape
+        h = self.norm(x)
+
+        # QKV projection
+        qkv = self.qkv(h)
+        q, k, v = torch.chunk(qkv, 3, dim=1)
+
+        # Reshape for attention
+        q = q.reshape(B, self.num_heads, C // self.num_heads, H * W).transpose(-1, -2)
+        k = k.reshape(B, self.num_heads, C // self.num_heads, H * W).transpose(-1, -2)
+        v = v.reshape(B, self.num_heads, C // self.num_heads, H * W).transpose(-1, -2)
+
+        # Attention
+        scale = (C // self.num_heads) ** -0.5
+        attn = torch.softmax(q @ k.transpose(-1, -2) * scale, dim=-1)
+        h = attn @ v
+
+        # Reshape back
+        h = h.transpose(-1, -2).reshape(B, C, H, W)
+
+        # Apply INL dynamics for refinement
+        h_flat = h.reshape(B, C, H * W).transpose(1, 2)  # (B, H*W, C)
+        h_refined = self.inl(h_flat)
+        h = h_refined.transpose(1, 2).reshape(B, C, H, W)
+
+        h = self.proj_out(h)
+
+        return x + h
+
+
+class INLUNet(nn.Module):
+    """
+    U-Net with INL dynamics for latent diffusion.
+
+    Denoises latent representations conditioned on text embeddings.
+    """
+    def __init__(
+        self,
+        in_channels: int = 4,
+        out_channels: int = 4,
+        model_channels: int = 320,
+        num_res_blocks: int = 2,
+        attention_resolutions: List[int] = [4, 2, 1],
+        channel_mult: List[int] = [1, 2, 4, 4],
+        num_heads: int = 8,
+        context_dim: int = 768  # Text embedding dimension
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+
+        # Time embedding
+        time_embed_dim = model_channels * 4
+        self.time_embed = nn.Sequential(
+            TimeEmbedding(model_channels),
+            nn.Linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            nn.Linear(time_embed_dim, time_embed_dim)
+        )
+
+        # Text conditioning projection
+        self.context_proj = nn.Linear(context_dim, model_channels)
+
+        # Input convolution
+        self.input_blocks = nn.ModuleList([
+            nn.Conv2d(in_channels, model_channels, 3, padding=1)
+        ])
+
+        # Encoder (downsampling)
+        ch = model_channels
+        input_block_chans = [ch]
+
+        for level, mult in enumerate(channel_mult):
+            for _ in range(num_res_blocks):
+                layers = [
+                    ResnetBlock(ch, mult * model_channels, time_embed_dim)
+                ]
+
+                ch = mult * model_channels
+
+                # Add attention at specified resolutions
+                if level in attention_resolutions:
+                    layers.append(INLAttentionBlock(ch, num_heads))
+
+                self.input_blocks.append(nn.Sequential(*layers))
+                input_block_chans.append(ch)
+
+            # Downsample
+            if level != len(channel_mult) - 1:
+                self.input_blocks.append(nn.Conv2d(ch, ch, 3, stride=2, padding=1))
+                input_block_chans.append(ch)
+
+        # Middle
+        self.middle_block = nn.Sequential(
+            ResnetBlock(ch, ch, time_embed_dim),
+            INLAttentionBlock(ch, num_heads, num_iterations=5),
+            ResnetBlock(ch, ch, time_embed_dim)
+        )
+
+        # Decoder (upsampling)
+        self.output_blocks = nn.ModuleList([])
+
+        for level, mult in list(enumerate(channel_mult))[::-1]:
+            for i in range(num_res_blocks + 1):
+                ich = input_block_chans.pop()
+                layers = [
+                    ResnetBlock(ch + ich, mult * model_channels, time_embed_dim)
+                ]
+
+                ch = mult * model_channels
+
+                if level in attention_resolutions:
+                    layers.append(INLAttentionBlock(ch, num_heads))
+
+                # Upsample
+                if level != 0 and i == num_res_blocks:
+                    layers.append(nn.Upsample(scale_factor=2, mode='nearest'))
+
+                self.output_blocks.append(nn.Sequential(*layers))
+
+        # Output
+        self.out = nn.Sequential(
+            nn.GroupNorm(8, ch),
+            nn.SiLU(),
+            nn.Conv2d(ch, out_channels, 3, padding=1)
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        timesteps: torch.Tensor,
+        context: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: Noisy latents (B, 4, H, W)
+            timesteps: Diffusion timesteps (B,)
+            context: Text embeddings (B, seq_len, context_dim)
+
+        Returns:
+            Predicted noise (B, 4, H, W)
+        """
+        # Time embedding
+        t_emb = self.time_embed(timesteps)
+
+        # Text conditioning (average pooling for simplicity)
+        if context is not None:
+            context = context.mean(dim=1)  # (B, context_dim)
+            context_emb = self.context_proj(context)
+            t_emb = t_emb + context_emb
+
+        # Encoder
+        hs = []
+        h = x
+        for module in self.input_blocks:
+            if isinstance(module, nn.Sequential):
+                for layer in module:
+                    if isinstance(layer, ResnetBlock):
+                        h = layer(h, t_emb)
+                    else:
+                        h = layer(h)
+            else:
+                h = module(h)
+            hs.append(h)
+
+        # Middle
+        for layer in self.middle_block:
+            if isinstance(layer, ResnetBlock):
+                h = layer(h, t_emb)
+            else:
+                h = layer(h)
+
+        # Decoder
+        for module in self.output_blocks:
+            h = torch.cat([h, hs.pop()], dim=1)
+            for layer in module:
+                if isinstance(layer, ResnetBlock):
+                    h = layer(h, t_emb)
+                else:
+                    h = layer(h)
+
+        # Output
+        return self.out(h)
+
+
+class VAEResBlock(nn.Module):
+    """
+    Residual block for VAE with GroupNorm.
+    Similar to Stable Diffusion VAE architecture.
+    """
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+
+        self.norm1 = nn.GroupNorm(32, in_channels)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
+
+        self.norm2 = nn.GroupNorm(32, out_channels)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
+
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv2d(in_channels, out_channels, 1)
+        else:
+            self.shortcut = nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.norm1(x)
+        h = F.silu(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = F.silu(h)
+        h = self.conv2(h)
+
+        return h + self.shortcut(x)
+
+
+class VAEAttentionBlock(nn.Module):
+    """
+    Self-attention block for VAE.
+    """
+    def __init__(self, channels: int):
+        super().__init__()
+
+        self.channels = channels
+        self.norm = nn.GroupNorm(32, channels)
+
+        self.q = nn.Conv2d(channels, channels, 1)
+        self.k = nn.Conv2d(channels, channels, 1)
+        self.v = nn.Conv2d(channels, channels, 1)
+
+        self.proj_out = nn.Conv2d(channels, channels, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, C, H, W = x.shape
+        h = self.norm(x)
+
+        q = self.q(h).reshape(B, C, H * W).transpose(1, 2)  # (B, HW, C)
+        k = self.k(h).reshape(B, C, H * W).transpose(1, 2)
+        v = self.v(h).reshape(B, C, H * W).transpose(1, 2)
+
+        # Attention
+        scale = C ** -0.5
+        attn = torch.softmax(q @ k.transpose(-1, -2) * scale, dim=-1)
+        h = attn @ v
+
+        # Reshape back
+        h = h.transpose(1, 2).reshape(B, C, H, W)
+        h = self.proj_out(h)
+
+        return x + h
+
+
+class Downsample(nn.Module):
+    """Downsampling layer."""
+    def __init__(self, channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, 3, stride=2, padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Asymmetric padding to match Stable Diffusion
+        x = F.pad(x, (0, 1, 0, 1), mode='constant', value=0)
+        return self.conv(x)
+
+
+class Upsample(nn.Module):
+    """Upsampling layer."""
+    def __init__(self, channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, 3, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.interpolate(x, scale_factor=2.0, mode='nearest')
+        return self.conv(x)
+
+
+class StableDiffusionVAE(nn.Module):
+    """
+    MASSIVE VAE for ultra-high quality latent diffusion.
+
+    Architecture:
+    - Input: 256x256x3 (or 512x512x3)
+    - Latent: 32x32x4 (8x downsampling)
+    - Deep ResNet blocks with GroupNorm
+    - Multi-head attention at multiple resolutions
+    - ~2.15B parameters for SOTA reconstruction quality
+
+    This beast preserves ALL details with near-perfect reconstruction.
+
+    Memory optimization:
+    - Use gradient_checkpointing=True to reduce memory by ~70% (trades 25% speed)
+    - Essential for training on GPUs < 24GB
+    """
+    def __init__(
+        self,
+        in_channels: int = 3,
+        latent_channels: int = 4,
+        base_channels: int = 256,  # 128 -> 256 (doubled!)
+        channel_multipliers: List[int] = [1, 2, 4, 8],  # [1,2,4,4] -> [1,2,4,8] (doubled max!)
+        num_res_blocks: int = 6,  # 2 -> 6 (tripled!)
+        attn_resolutions: List[int] = [128, 64, 32],  # More attention layers!
+        use_gradient_checkpointing: bool = False  # Enable for memory-constrained GPUs
+    ):
+        super().__init__()
+
+        self.latent_channels = latent_channels
+        self.use_gradient_checkpointing = use_gradient_checkpointing
+
+        # ========== ENCODER ==========
+        # Input: 256x256x3 -> 256x256x256
+        self.encoder_conv_in = nn.Conv2d(in_channels, base_channels, 3, padding=1)
+
+        # Downsampling blocks
+        self.encoder_blocks = nn.ModuleList()
+        ch = base_channels
+        resolutions = []
+        current_res = 256  # Assume 256x256 input
+
+        for level, mult in enumerate(channel_multipliers):
+            out_ch = base_channels * mult
+
+            # Add MANY residual blocks (6 per level!)
+            for _ in range(num_res_blocks):
+                self.encoder_blocks.append(VAEResBlock(ch, out_ch))
+                ch = out_ch
+                resolutions.append(current_res)
+
+            # Add attention at specified resolutions
+            if current_res in attn_resolutions:
+                # Add MULTIPLE attention blocks for better quality
+                self.encoder_blocks.append(VAEAttentionBlock(ch))
+                self.encoder_blocks.append(VAEAttentionBlock(ch))
+                resolutions.append(current_res)
+                resolutions.append(current_res)
+
+            # Downsample (except last level)
+            if level != len(channel_multipliers) - 1:
+                self.encoder_blocks.append(Downsample(ch))
+                current_res //= 2
+                resolutions.append(current_res)
+
+        # Middle blocks (at 32x32x2048!) - MASSIVE bottleneck
+        self.encoder_mid_block1 = VAEResBlock(ch, ch)
+        self.encoder_mid_attn1 = VAEAttentionBlock(ch)
+        self.encoder_mid_block2 = VAEResBlock(ch, ch)
+        self.encoder_mid_attn2 = VAEAttentionBlock(ch)
+        self.encoder_mid_block3 = VAEResBlock(ch, ch)
+        self.encoder_mid_attn3 = VAEAttentionBlock(ch)
+        self.encoder_mid_block4 = VAEResBlock(ch, ch)
+
+        # Output: mu and logvar
+        self.encoder_norm_out = nn.GroupNorm(32, ch)
+        self.encoder_conv_out = nn.Conv2d(ch, latent_channels * 2, 3, padding=1)
+
+        # ========== DECODER ==========
+        # Input: 32x32x4 -> 32x32x2048
+        self.decoder_conv_in = nn.Conv2d(latent_channels, ch, 3, padding=1)
+
+        # Middle blocks - MASSIVE processing
+        self.decoder_mid_block1 = VAEResBlock(ch, ch)
+        self.decoder_mid_attn1 = VAEAttentionBlock(ch)
+        self.decoder_mid_block2 = VAEResBlock(ch, ch)
+        self.decoder_mid_attn2 = VAEAttentionBlock(ch)
+        self.decoder_mid_block3 = VAEResBlock(ch, ch)
+        self.decoder_mid_attn3 = VAEAttentionBlock(ch)
+        self.decoder_mid_block4 = VAEResBlock(ch, ch)
+
+        # Upsampling blocks
+        self.decoder_blocks = nn.ModuleList()
+
+        for level, mult in reversed(list(enumerate(channel_multipliers))):
+            out_ch = base_channels * mult
+
+            # MANY residual blocks per level
+            for _ in range(num_res_blocks + 1):
+                self.decoder_blocks.append(VAEResBlock(ch, out_ch))
+                ch = out_ch
+
+            # Add attention at specified resolutions
+            if current_res in attn_resolutions:
+                # Multiple attention blocks
+                self.decoder_blocks.append(VAEAttentionBlock(ch))
+                self.decoder_blocks.append(VAEAttentionBlock(ch))
+
+            # Upsample (except first level, which is last in reversed order)
+            if level != 0:
+                self.decoder_blocks.append(Upsample(ch))
+                current_res *= 2
+
+        # Output: 256x256x3
+        self.decoder_norm_out = nn.GroupNorm(32, ch)
+        self.decoder_conv_out = nn.Conv2d(ch, in_channels, 3, padding=1)
+
+    def encode(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Encode image to latent distribution parameters.
+
+        Args:
+            x: (B, 3, H, W) images in range [-1, 1]
+
+        Returns:
+            mu: (B, latent_channels, H/8, W/8)
+            logvar: (B, latent_channels, H/8, W/8)
+        """
+        # Input conv
+        h = self.encoder_conv_in(x)
+
+        # Encoder blocks
+        for block in self.encoder_blocks:
+            h = block(h)
+
+        # Middle - DEEP processing
+        h = self.encoder_mid_block1(h)
+        h = self.encoder_mid_attn1(h)
+        h = self.encoder_mid_block2(h)
+        h = self.encoder_mid_attn2(h)
+        h = self.encoder_mid_block3(h)
+        h = self.encoder_mid_attn3(h)
+        h = self.encoder_mid_block4(h)
+
+        # Output
+        h = self.encoder_norm_out(h)
+        h = F.silu(h)
+        h = self.encoder_conv_out(h)
+
+        mu, logvar = torch.chunk(h, 2, dim=1)
+        return mu, logvar
+
+    def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+        """Sample from latent distribution."""
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        """
+        Decode latent to image.
+
+        Args:
+            z: (B, latent_channels, H/8, W/8) latent codes
+
+        Returns:
+            x: (B, 3, H, W) reconstructed images in range [-1, 1]
+        """
+        # Input conv
+        h = self.decoder_conv_in(z)
+
+        # Middle - DEEP processing
+        h = self.decoder_mid_block1(h)
+        h = self.decoder_mid_attn1(h)
+        h = self.decoder_mid_block2(h)
+        h = self.decoder_mid_attn2(h)
+        h = self.decoder_mid_block3(h)
+        h = self.decoder_mid_attn3(h)
+        h = self.decoder_mid_block4(h)
+
+        # Decoder blocks
+        for block in self.decoder_blocks:
+            h = block(h)
+
+        # Output
+        h = self.decoder_norm_out(h)
+        h = F.silu(h)
+        h = self.decoder_conv_out(h)
+
+        return torch.tanh(h)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Full forward pass: encode -> sample -> decode."""
+        mu, logvar = self.encode(x)
+        z = self.reparameterize(mu, logvar)
+        recon = self.decode(z)
+        return recon, mu, logvar
+
+
+class SimpleVAE(nn.Module):
+    """
+    DEPRECATED: Simple VAE with only 2M parameters.
+    Use StableDiffusionVAE instead for production.
+
+    Simple VAE for encoding images to latent space.
+    Compress 512x512x3 -> 64x64x4
+    """
+    def __init__(self, in_channels: int = 3, latent_channels: int = 4):
+        super().__init__()
+
+        # Encoder (512 -> 64, 8x downsampling)
+        self.encoder = nn.Sequential(
+            nn.Conv2d(in_channels, 128, 3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(128, 128, 3, stride=2, padding=1),  # 256
+            nn.ReLU(),
+            nn.Conv2d(128, 256, 3, stride=2, padding=1),  # 128
+            nn.ReLU(),
+            nn.Conv2d(256, 256, 3, stride=2, padding=1),  # 64
+            nn.ReLU(),
+            nn.Conv2d(256, latent_channels * 2, 3, padding=1)  # mu, logvar
+        )
+
+        # Decoder (64 -> 512)
+        self.decoder = nn.Sequential(
+            nn.Conv2d(latent_channels, 256, 3, padding=1),
+            nn.ReLU(),
+            nn.Upsample(scale_factor=2, mode='nearest'),  # 128
+            nn.Conv2d(256, 256, 3, padding=1),
+            nn.ReLU(),
+            nn.Upsample(scale_factor=2, mode='nearest'),  # 256
+            nn.Conv2d(256, 128, 3, padding=1),
+            nn.ReLU(),
+            nn.Upsample(scale_factor=2, mode='nearest'),  # 512
+            nn.Conv2d(128, 128, 3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(128, in_channels, 3, padding=1),
+            nn.Tanh()
+        )
+
+    def encode(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        h = self.encoder(x)
+        mu, logvar = torch.chunk(h, 2, dim=1)
+        return mu, logvar
+
+    def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        return self.decoder(z)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        mu, logvar = self.encode(x)
+        z = self.reparameterize(mu, logvar)
+        recon = self.decode(z)
+        return recon, mu, logvar
+
+
+class INLTextEncoder(nn.Module):
+    """
+    Text encoder using pre-trained INL-LLM model.
+
+    Reuses the trained INL-LLM (1.1B) as a powerful text encoder
+    with integrator neuron dynamics.
+    """
+    def __init__(self, inl_llm_model, embed_dim: int = 768):
+        super().__init__()
+
+        # Use pretrained INL-LLM
+        self.inl_llm = inl_llm_model
+
+        # Freeze INL-LLM (use as feature extractor)
+        for param in self.inl_llm.parameters():
+            param.requires_grad = False
+
+        # Project INL-LLM hidden states to diffusion context dimension
+        llm_hidden_dim = self.inl_llm.d_model
+        self.projection = nn.Linear(llm_hidden_dim, embed_dim)
+
+    def forward(self, text_tokens: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            text_tokens: (B, seq_len) token IDs
+
+        Returns:
+            text_embeddings: (B, seq_len, embed_dim)
+        """
+        with torch.no_grad():
+            # Get hidden states from INL-LLM (no generation, just encoding)
+            # Use the model's embedding + transformer blocks
+            x = self.inl_llm.token_embedding(text_tokens)
+            x = self.inl_llm.pos_encoding(x)
+
+            # Pass through INL layers to get contextualized representations
+            for layer in self.inl_llm.layers:
+                x, _ = layer(x)
+
+            x = self.inl_llm.norm(x)
+
+        # Project to context dimension
+        x = self.projection(x)
+
+        return x
+
+
+class INLLatentDiffusion(nn.Module):
+    """
+    Complete Latent Diffusion Model with INL dynamics.
+
+    Text → Image generation pipeline.
+    """
+    def __init__(
+        self,
+        img_size: int = 512,
+        latent_size: int = 64,
+        inl_llm_model = None,  # Pre-trained INL-LLM for text encoding
+        context_dim: int = 768
+    ):
+        super().__init__()
+
+        self.img_size = img_size
+        self.latent_size = latent_size
+
+        # Components - Use MASSIVE 1B+ parameter VAE
+        self.vae = StableDiffusionVAE(
+            in_channels=3,
+            latent_channels=4,
+            base_channels=256,
+            channel_multipliers=[1, 2, 4, 8],
+            num_res_blocks=6,
+            attn_resolutions=[128, 64, 32]
+        )
+        print(f"✅ Using StableDiffusionVAE with {sum(p.numel() for p in self.vae.parameters()):,} parameters")
+
+        # Use INL-LLM as text encoder if provided
+        if inl_llm_model is not None:
+            self.text_encoder = INLTextEncoder(inl_llm_model, embed_dim=context_dim)
+            print("✅ Using pre-trained INL-LLM as text encoder (frozen)")
+        else:
+            # Fallback to simple encoder
+            print("⚠️ No INL-LLM provided, using simple text encoder")
+            from .integrator_language_model import UltraOptimizedIntegratorLanguageModel
+            # Create a small text encoder
+            small_llm = UltraOptimizedIntegratorLanguageModel(
+                vocab_size=50000,
+                d_model=512,
+                num_layers=6,
+                num_heads=8,
+                num_iterations_per_layer=3
+            )
+            self.text_encoder = INLTextEncoder(small_llm, embed_dim=context_dim)
+
+        self.unet = INLUNet(
+            in_channels=4,
+            out_channels=4,
+            model_channels=320,
+            context_dim=context_dim
+        )
+
+        # Diffusion parameters
+        self.num_timesteps = 1000
+        self.register_buffer('betas', self._cosine_beta_schedule())
+        self.register_buffer('alphas', 1.0 - self.betas)
+        self.register_buffer('alphas_cumprod', torch.cumprod(self.alphas, dim=0))
+
+    def _cosine_beta_schedule(self, s: float = 0.008) -> torch.Tensor:
+        """Cosine schedule from Improved DDPM."""
+        steps = self.num_timesteps + 1
+        x = torch.linspace(0, self.num_timesteps, steps)
+        alphas_cumprod = torch.cos(((x / self.num_timesteps) + s) / (1 + s) * math.pi * 0.5) ** 2
+        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+        betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+        return torch.clip(betas, 0.0001, 0.9999)
+
+    @torch.no_grad()
+    def generate(
+        self,
+        text_tokens: torch.Tensor,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5
+    ) -> torch.Tensor:
+        """
+        Generate images from text prompts.
+
+        Args:
+            text_tokens: (B, seq_len) text token IDs
+            num_inference_steps: Number of denoising steps
+            guidance_scale: Classifier-free guidance scale
+
+        Returns:
+            generated_images: (B, 3, img_size, img_size)
+        """
+        B = text_tokens.size(0)
+        device = text_tokens.device
+
+        # Encode text
+        context = self.text_encoder(text_tokens)
+
+        # Start from random noise
+        latents = torch.randn(B, 4, self.latent_size, self.latent_size, device=device)
+
+        # Denoising loop (DDPM sampling)
+        timesteps = torch.linspace(self.num_timesteps - 1, 0, num_inference_steps, dtype=torch.long, device=device)
+
+        for t in timesteps:
+            t_batch = t.repeat(B)
+
+            # Predict noise
+            noise_pred = self.unet(latents, t_batch, context)
+
+            # Update latents (simplified DDPM step)
+            alpha = self.alphas_cumprod[t]
+            alpha_prev = self.alphas_cumprod[t - 1] if t > 0 else torch.tensor(1.0, device=device)
+
+            beta_t = 1 - alpha / alpha_prev
+            latents = (latents - beta_t * noise_pred) / torch.sqrt(1 - beta_t)
+
+        # Decode latents to images
+        images = self.vae.decode(latents)
+
+        return images
+
+    def get_num_params(self):
+        """Count total parameters."""
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)

inl_llm/models/inl_vision.py

@@ -1,366 +1,366 @@
-"""
-INL-Vision: Image-to-Image model based on Integrator Neuron dynamics
-
-Adapts the INL-LLM architecture for vision tasks by treating image patches
-as tokens and using the same equilibrium-based dynamics.
-
-Author: Boris Peyriguère
-"""
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from typing import Optional, Tuple
-import math
-
-from ..optimizations.optimizations import (
-    LowRankEmbedding,
-    AdaptiveIntegratorNeuronLayer
-)
-from ..core.integrator_neuron_layer import IntegratorNeuronLayer
-
-
-class SimpleINLDynamics(nn.Module):
-    """
-    Simplified Integrator Neuron Layer for vision.
-
-    Uses integrator dynamics without the full complexity of INL:
-    - x_{t+1} = x_t + dt * MLP(x_t)
-    - Iterated num_iterations times for equilibrium
-
-    This gives similar dynamics but simpler implementation.
-    """
-    def __init__(
-        self,
-        d_model: int,
-        num_iterations: int = 5,
-        dt: float = 0.1
-    ):
-        super().__init__()
-
-        self.d_model = d_model
-        self.num_iterations = num_iterations
-        self.dt = dt
-
-        # Simple MLP for dynamics
-        self.dynamics_mlp = nn.Sequential(
-            nn.Linear(d_model, d_model),
-            nn.GELU(),
-            nn.Linear(d_model, d_model)
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            x: Input state (B, seq_len, d_model)
-
-        Returns:
-            Final state after iterations (B, seq_len, d_model)
-        """
-        # Iterate to refine representation
-        state = x
-        for _ in range(self.num_iterations):
-            delta = self.dynamics_mlp(state)
-            state = state + self.dt * delta
-
-        return state
-
-
-class PatchEmbedding(nn.Module):
-    """
-    Split image into patches and embed them.
-    Similar to ViT (Vision Transformer) patch embedding.
-    """
-    def __init__(
-        self,
-        img_size: int = 224,
-        patch_size: int = 16,
-        in_channels: int = 3,
-        embed_dim: int = 768
-    ):
-        super().__init__()
-        self.img_size = img_size
-        self.patch_size = patch_size
-        self.num_patches = (img_size // patch_size) ** 2
-
-        # Convolutional projection
-        self.proj = nn.Conv2d(
-            in_channels,
-            embed_dim,
-            kernel_size=patch_size,
-            stride=patch_size
-        )
-
-    def forward(self, x):
-        # x: (B, C, H, W)
-        B, C, H, W = x.shape
-
-        # Project patches
-        x = self.proj(x)  # (B, embed_dim, H/patch_size, W/patch_size)
-
-        # Flatten spatial dimensions
-        x = x.flatten(2).transpose(1, 2)  # (B, num_patches, embed_dim)
-
-        return x
-
-
-class INLVisionBlock(nn.Module):
-    """
-    Vision block using Integrator Neuron Layer dynamics.
-    Applies equilibrium-based processing to image patch embeddings.
-    """
-    def __init__(
-        self,
-        d_model: int,
-        num_heads: int,
-        num_iterations: int,
-        layer_idx: int,
-        feedforward_dim: int,
-        dropout: float = 0.1,
-        group_size: int = 64,
-        excitation_sparsity: float = 0.1
-    ):
-        super().__init__()
-
-        self.d_model = d_model
-        self.num_iterations = num_iterations
-        self.layer_idx = layer_idx
-
-        # Layer normalization
-        self.norm1 = nn.LayerNorm(d_model)
-        self.norm2 = nn.LayerNorm(d_model)
-        self.norm_attn = nn.LayerNorm(d_model)
-
-        # Multi-head attention (for patch-to-patch interactions)
-        self.attention = nn.MultiheadAttention(
-            embed_dim=d_model,
-            num_heads=num_heads,
-            dropout=dropout,
-            batch_first=True
-        )
-
-        # Feedforward network
-        self.ffn = nn.Sequential(
-            nn.Linear(d_model, feedforward_dim),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(feedforward_dim, d_model),
-            nn.Dropout(dropout)
-        )
-
-        # Use simplified INL dynamics for vision
-        self.inl_layer = SimpleINLDynamics(
-            d_model=d_model,
-            num_iterations=num_iterations,
-            dt=0.1
-        )
-
-    def forward(self, x, return_trajectory=False):
-        """
-        Forward pass with integrator dynamics.
-
-        Args:
-            x: (B, num_patches, d_model)
-            return_trajectory: Return full dynamics trajectory
-        """
-        trajectory = None
-
-        # Self-attention on patches
-        attn_out, _ = self.attention(
-            self.norm_attn(x),
-            self.norm_attn(x),
-            self.norm_attn(x)
-        )
-        x = x + attn_out
-
-        # Apply integrator dynamics to patch embeddings (iterate multiple times)
-        x_normed = self.norm1(x)
-
-        # Run integrator dynamics (wrapper handles iterations internally)
-        inl_out = self.inl_layer(x_normed)
-        x = x + inl_out
-
-        trajectory = None  # Simplified: no trajectory tracking yet
-
-        # Feedforward
-        x = x + self.ffn(self.norm2(x))
-
-        return (x, trajectory) if return_trajectory else x
-
-
-class INLVisionModel(nn.Module):
-    """
-    Complete INL-Vision model for image-to-image tasks.
-
-    Uses integrator neuron dynamics to process image patches iteratively,
-    allowing the model to refine representations through equilibrium-based dynamics.
-    """
-    def __init__(
-        self,
-        img_size: int = 224,
-        patch_size: int = 16,
-        in_channels: int = 3,
-        out_channels: int = 3,
-        d_model: int = 768,
-        num_layers: int = 12,
-        num_heads: int = 12,
-        num_iterations_per_layer: int = 5,
-        feedforward_dim: int = None,
-        dropout: float = 0.1,
-        # Optimizations
-        use_shared_controllers: bool = True,
-        hierarchical_group_size: int = 64,
-        excitation_sparsity: float = 0.1
-    ):
-        super().__init__()
-
-        self.img_size = img_size
-        self.patch_size = patch_size
-        self.d_model = d_model
-        self.num_layers = num_layers
-
-        if feedforward_dim is None:
-            feedforward_dim = 4 * d_model
-
-        # Patch embedding
-        self.patch_embed = PatchEmbedding(
-            img_size=img_size,
-            patch_size=patch_size,
-            in_channels=in_channels,
-            embed_dim=d_model
-        )
-
-        num_patches = self.patch_embed.num_patches
-
-        # Positional encoding for patches
-        self.pos_embedding = nn.Parameter(
-            torch.randn(1, num_patches, d_model) * 0.02
-        )
-
-        # Note: For simplicity in this vision model, we don't use shared controllers
-        # Each block has its own integrator layer
-        self.use_shared_controllers = use_shared_controllers
-        if use_shared_controllers:
-            print(f"ℹ️  Shared controllers disabled for INL-Vision (using per-layer controllers)")
-        self.shared_controller = None
-
-        # Vision blocks with integrator dynamics
-        self.blocks = nn.ModuleList([
-            INLVisionBlock(
-                d_model=d_model,
-                num_heads=num_heads,
-                num_iterations=num_iterations_per_layer,
-                layer_idx=i,
-                feedforward_dim=feedforward_dim,
-                dropout=dropout,
-                group_size=hierarchical_group_size,
-                excitation_sparsity=excitation_sparsity
-            )
-            for i in range(num_layers)
-        ])
-
-        # Final layer norm
-        self.norm = nn.LayerNorm(d_model)
-
-        # Decoder: patches back to image
-        self.decoder = nn.Sequential(
-            nn.Linear(d_model, patch_size * patch_size * out_channels),
-            nn.Tanh()  # Output in [-1, 1]
-        )
-
-        self.out_channels = out_channels
-
-    def forward(self, x, return_aux=False):
-        """
-        Forward pass.
-
-        Args:
-            x: Input image (B, C, H, W)
-            return_aux: Return auxiliary information (trajectories)
-
-        Returns:
-            Output image (B, C, H, W)
-            Optional: trajectories from all layers
-        """
-        B, C, H, W = x.shape
-
-        # Embed patches
-        x = self.patch_embed(x)  # (B, num_patches, d_model)
-
-        # Add positional encoding
-        x = x + self.pos_embedding
-
-        # Apply vision blocks with integrator dynamics
-        trajectories = []
-        for block in self.blocks:
-            if return_aux:
-                x, traj = block(x, return_trajectory=True)
-                trajectories.append(traj)
-            else:
-                x = block(x)
-
-        # Final norm
-        x = self.norm(x)
-
-        # Decode patches back to image
-        x = self.decoder(x)  # (B, num_patches, patch_size^2 * C)
-
-        # Reshape to image
-        num_patches_per_side = self.img_size // self.patch_size
-        x = x.reshape(B, num_patches_per_side, num_patches_per_side,
-                     self.patch_size, self.patch_size, self.out_channels)
-
-        # Rearrange to (B, C, H, W)
-        x = x.permute(0, 5, 1, 3, 2, 4).contiguous()
-        x = x.reshape(B, self.out_channels, self.img_size, self.img_size)
-
-        if return_aux:
-            return x, trajectories
-        return x
-
-    def get_num_params(self):
-        """Count total parameters."""
-        return sum(p.numel() for p in self.parameters() if p.requires_grad)
-
-
-def create_inl_vision_model(size='small', img_size=224, **kwargs):
-    """
-    Factory function to create INL-Vision models of different sizes.
-
-    Args:
-        size: 'tiny', 'small', 'base', 'large'
-        img_size: Input image size
-        **kwargs: Override default parameters
-    """
-    configs = {
-        'tiny': {
-            'd_model': 192,
-            'num_layers': 12,
-            'num_heads': 3,
-            'feedforward_dim': 768
-        },
-        'small': {
-            'd_model': 384,
-            'num_layers': 12,
-            'num_heads': 6,
-            'feedforward_dim': 1536
-        },
-        'base': {
-            'd_model': 768,
-            'num_layers': 12,
-            'num_heads': 12,
-            'feedforward_dim': 3072
-        },
-        'large': {
-            'd_model': 1024,
-            'num_layers': 24,
-            'num_heads': 16,
-            'feedforward_dim': 4096
-        }
-    }
-
-    config = configs.get(size, configs['small'])
-    config.update(kwargs)
-    config['img_size'] = img_size
-
-    return INLVisionModel(**config)
+"""
+INL-Vision: Image-to-Image model based on Integrator Neuron dynamics
+
+Adapts the INL-LLM architecture for vision tasks by treating image patches
+as tokens and using the same equilibrium-based dynamics.
+
+Author: Boris Peyriguère
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple
+import math
+
+from ..optimizations.optimizations import (
+    LowRankEmbedding,
+    AdaptiveIntegratorNeuronLayer
+)
+from ..core.integrator_neuron_layer import IntegratorNeuronLayer
+
+
+class SimpleINLDynamics(nn.Module):
+    """
+    Simplified Integrator Neuron Layer for vision.
+
+    Uses integrator dynamics without the full complexity of INL:
+    - x_{t+1} = x_t + dt * MLP(x_t)
+    - Iterated num_iterations times for equilibrium
+
+    This gives similar dynamics but simpler implementation.
+    """
+    def __init__(
+        self,
+        d_model: int,
+        num_iterations: int = 5,
+        dt: float = 0.1
+    ):
+        super().__init__()
+
+        self.d_model = d_model
+        self.num_iterations = num_iterations
+        self.dt = dt
+
+        # Simple MLP for dynamics
+        self.dynamics_mlp = nn.Sequential(
+            nn.Linear(d_model, d_model),
+            nn.GELU(),
+            nn.Linear(d_model, d_model)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Input state (B, seq_len, d_model)
+
+        Returns:
+            Final state after iterations (B, seq_len, d_model)
+        """
+        # Iterate to refine representation
+        state = x
+        for _ in range(self.num_iterations):
+            delta = self.dynamics_mlp(state)
+            state = state + self.dt * delta
+
+        return state
+
+
+class PatchEmbedding(nn.Module):
+    """
+    Split image into patches and embed them.
+    Similar to ViT (Vision Transformer) patch embedding.
+    """
+    def __init__(
+        self,
+        img_size: int = 224,
+        patch_size: int = 16,
+        in_channels: int = 3,
+        embed_dim: int = 768
+    ):
+        super().__init__()
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = (img_size // patch_size) ** 2
+
+        # Convolutional projection
+        self.proj = nn.Conv2d(
+            in_channels,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size
+        )
+
+    def forward(self, x):
+        # x: (B, C, H, W)
+        B, C, H, W = x.shape
+
+        # Project patches
+        x = self.proj(x)  # (B, embed_dim, H/patch_size, W/patch_size)
+
+        # Flatten spatial dimensions
+        x = x.flatten(2).transpose(1, 2)  # (B, num_patches, embed_dim)
+
+        return x
+
+
+class INLVisionBlock(nn.Module):
+    """
+    Vision block using Integrator Neuron Layer dynamics.
+    Applies equilibrium-based processing to image patch embeddings.
+    """
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        num_iterations: int,
+        layer_idx: int,
+        feedforward_dim: int,
+        dropout: float = 0.1,
+        group_size: int = 64,
+        excitation_sparsity: float = 0.1
+    ):
+        super().__init__()
+
+        self.d_model = d_model
+        self.num_iterations = num_iterations
+        self.layer_idx = layer_idx
+
+        # Layer normalization
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm_attn = nn.LayerNorm(d_model)
+
+        # Multi-head attention (for patch-to-patch interactions)
+        self.attention = nn.MultiheadAttention(
+            embed_dim=d_model,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True
+        )
+
+        # Feedforward network
+        self.ffn = nn.Sequential(
+            nn.Linear(d_model, feedforward_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(feedforward_dim, d_model),
+            nn.Dropout(dropout)
+        )
+
+        # Use simplified INL dynamics for vision
+        self.inl_layer = SimpleINLDynamics(
+            d_model=d_model,
+            num_iterations=num_iterations,
+            dt=0.1
+        )
+
+    def forward(self, x, return_trajectory=False):
+        """
+        Forward pass with integrator dynamics.
+
+        Args:
+            x: (B, num_patches, d_model)
+            return_trajectory: Return full dynamics trajectory
+        """
+        trajectory = None
+
+        # Self-attention on patches
+        attn_out, _ = self.attention(
+            self.norm_attn(x),
+            self.norm_attn(x),
+            self.norm_attn(x)
+        )
+        x = x + attn_out
+
+        # Apply integrator dynamics to patch embeddings (iterate multiple times)
+        x_normed = self.norm1(x)
+
+        # Run integrator dynamics (wrapper handles iterations internally)
+        inl_out = self.inl_layer(x_normed)
+        x = x + inl_out
+
+        trajectory = None  # Simplified: no trajectory tracking yet
+
+        # Feedforward
+        x = x + self.ffn(self.norm2(x))
+
+        return (x, trajectory) if return_trajectory else x
+
+
+class INLVisionModel(nn.Module):
+    """
+    Complete INL-Vision model for image-to-image tasks.
+
+    Uses integrator neuron dynamics to process image patches iteratively,
+    allowing the model to refine representations through equilibrium-based dynamics.
+    """
+    def __init__(
+        self,
+        img_size: int = 224,
+        patch_size: int = 16,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        d_model: int = 768,
+        num_layers: int = 12,
+        num_heads: int = 12,
+        num_iterations_per_layer: int = 5,
+        feedforward_dim: int = None,
+        dropout: float = 0.1,
+        # Optimizations
+        use_shared_controllers: bool = True,
+        hierarchical_group_size: int = 64,
+        excitation_sparsity: float = 0.1
+    ):
+        super().__init__()
+
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.d_model = d_model
+        self.num_layers = num_layers
+
+        if feedforward_dim is None:
+            feedforward_dim = 4 * d_model
+
+        # Patch embedding
+        self.patch_embed = PatchEmbedding(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_channels=in_channels,
+            embed_dim=d_model
+        )
+
+        num_patches = self.patch_embed.num_patches
+
+        # Positional encoding for patches
+        self.pos_embedding = nn.Parameter(
+            torch.randn(1, num_patches, d_model) * 0.02
+        )
+
+        # Note: For simplicity in this vision model, we don't use shared controllers
+        # Each block has its own integrator layer
+        self.use_shared_controllers = use_shared_controllers
+        if use_shared_controllers:
+            print(f"ℹ️  Shared controllers disabled for INL-Vision (using per-layer controllers)")
+        self.shared_controller = None
+
+        # Vision blocks with integrator dynamics
+        self.blocks = nn.ModuleList([
+            INLVisionBlock(
+                d_model=d_model,
+                num_heads=num_heads,
+                num_iterations=num_iterations_per_layer,
+                layer_idx=i,
+                feedforward_dim=feedforward_dim,
+                dropout=dropout,
+                group_size=hierarchical_group_size,
+                excitation_sparsity=excitation_sparsity
+            )
+            for i in range(num_layers)
+        ])
+
+        # Final layer norm
+        self.norm = nn.LayerNorm(d_model)
+
+        # Decoder: patches back to image
+        self.decoder = nn.Sequential(
+            nn.Linear(d_model, patch_size * patch_size * out_channels),
+            nn.Tanh()  # Output in [-1, 1]
+        )
+
+        self.out_channels = out_channels
+
+    def forward(self, x, return_aux=False):
+        """
+        Forward pass.
+
+        Args:
+            x: Input image (B, C, H, W)
+            return_aux: Return auxiliary information (trajectories)
+
+        Returns:
+            Output image (B, C, H, W)
+            Optional: trajectories from all layers
+        """
+        B, C, H, W = x.shape
+
+        # Embed patches
+        x = self.patch_embed(x)  # (B, num_patches, d_model)
+
+        # Add positional encoding
+        x = x + self.pos_embedding
+
+        # Apply vision blocks with integrator dynamics
+        trajectories = []
+        for block in self.blocks:
+            if return_aux:
+                x, traj = block(x, return_trajectory=True)
+                trajectories.append(traj)
+            else:
+                x = block(x)
+
+        # Final norm
+        x = self.norm(x)
+
+        # Decode patches back to image
+        x = self.decoder(x)  # (B, num_patches, patch_size^2 * C)
+
+        # Reshape to image
+        num_patches_per_side = self.img_size // self.patch_size
+        x = x.reshape(B, num_patches_per_side, num_patches_per_side,
+                     self.patch_size, self.patch_size, self.out_channels)
+
+        # Rearrange to (B, C, H, W)
+        x = x.permute(0, 5, 1, 3, 2, 4).contiguous()
+        x = x.reshape(B, self.out_channels, self.img_size, self.img_size)
+
+        if return_aux:
+            return x, trajectories
+        return x
+
+    def get_num_params(self):
+        """Count total parameters."""
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+
+def create_inl_vision_model(size='small', img_size=224, **kwargs):
+    """
+    Factory function to create INL-Vision models of different sizes.
+
+    Args:
+        size: 'tiny', 'small', 'base', 'large'
+        img_size: Input image size
+        **kwargs: Override default parameters
+    """
+    configs = {
+        'tiny': {
+            'd_model': 192,
+            'num_layers': 12,
+            'num_heads': 3,
+            'feedforward_dim': 768
+        },
+        'small': {
+            'd_model': 384,
+            'num_layers': 12,
+            'num_heads': 6,
+            'feedforward_dim': 1536
+        },
+        'base': {
+            'd_model': 768,
+            'num_layers': 12,
+            'num_heads': 12,
+            'feedforward_dim': 3072
+        },
+        'large': {
+            'd_model': 1024,
+            'num_layers': 24,
+            'num_heads': 16,
+            'feedforward_dim': 4096
+        }
+    }
+
+    config = configs.get(size, configs['small'])
+    config.update(kwargs)
+    config['img_size'] = img_size
+
+    return INLVisionModel(**config)

inl_llm/models/integrator_language_model.py

@@ -1,873 +1,990 @@
-"""
-ULTRA-Optimized Integrator Language Model (INL-LLM)
-
-Combines ALL optimizations for maximum efficiency:
-
-LEVEL 1 (Basic):
-- Low-rank embeddings (-70-80% embedding params)
-- Gradient checkpointing (-50-70% memory)
-- Adaptive early stopping (+30-50% inference speed)
-
-LEVEL 2 (Advanced):
-- Shared controllers (-96% controller params)
-- Sparse harmonic excitation (10x less compute)
-- Hierarchical equilibrium (-98% equilibrium params)
-
-RESULT: Can scale to 100B+ parameters with MUCH higher efficiency
-
-Author: Boris Peyriguère
-"""
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from typing import Optional, Tuple, Dict, List
-import math
-
-from ..optimizations.optimizations import (
-    LowRankEmbedding,
-    GradientCheckpointedINL,
-    AdaptiveIntegratorNeuronLayer,
-    AdaptiveHierarchicalINL
-)
-from ..optimizations.advanced_optimizations import (
-    SharedController,
-    SparseHarmonicINL,
-    HierarchicalEquilibriumINL
-)
-
-
-# ============================================================================
-# KV CACHE SUPPORT FOR INL-LLM
-# ============================================================================
-
-class INLCacheLayer:
-    """
-    Cache for a single layer, storing:
-    - Attention K, V (standard transformer cache)
-
-    NOTE: We do NOT cache integrator x, v states because integrator dynamics
-    run WITHIN each layer for each token, not across tokens. Only attention
-    needs to look back at previous tokens' K, V.
-    """
-
-    def __init__(self):
-        self.keys: Optional[torch.Tensor] = None          # [B, num_heads, seq_len, head_dim]
-        self.values: Optional[torch.Tensor] = None        # [B, num_heads, seq_len, head_dim]
-
-    def update_attention(
-        self,
-        new_keys: torch.Tensor,
-        new_values: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Update attention cache with new K, V.
-
-        Args:
-            new_keys: [B, num_heads, new_seq_len, head_dim]
-            new_values: [B, num_heads, new_seq_len, head_dim]
-
-        Returns:
-            Full keys, values (concatenated with past)
-        """
-        if self.keys is None:
-            # First time: initialize cache
-            self.keys = new_keys
-            self.values = new_values
-        else:
-            # Concatenate along sequence dimension
-            self.keys = torch.cat([self.keys, new_keys], dim=2)
-            self.values = torch.cat([self.values, new_values], dim=2)
-
-        return self.keys, self.values
-
-    def get_seq_length(self) -> int:
-        """Get current sequence length in cache."""
-        if self.keys is not None:
-            return self.keys.shape[2]
-        return 0
-
-    def reorder_batch(self, beam_idx: torch.LongTensor):
-        """Reorder cache for beam search."""
-        if self.keys is not None:
-            device = self.keys.device
-            self.keys = self.keys.index_select(0, beam_idx.to(device))
-            self.values = self.values.index_select(0, beam_idx.to(device))
-
-
-class INLCache:
-    """
-    Complete cache for INL-LLM model.
-
-    Stores attention K, V for all layers.
-    Compatible with HuggingFace's past_key_values interface.
-
-    NOTE: Simpler than typical transformers - we only cache attention K, V,
-    not integrator states since those are computed fresh for each token.
-    """
-
-    def __init__(self, num_layers: int):
-        self.num_layers = num_layers
-        self.layers: List[INLCacheLayer] = [INLCacheLayer() for _ in range(num_layers)]
-
-    def __getitem__(self, layer_idx: int) -> INLCacheLayer:
-        """Access cache for specific layer."""
-        return self.layers[layer_idx]
-
-    def __len__(self) -> int:
-        """Number of layers."""
-        return self.num_layers
-
-    def get_seq_length(self, layer_idx: int = 0) -> int:
-        """Get current sequence length (all layers should be same)."""
-        return self.layers[layer_idx].get_seq_length()
-
-    def reorder_cache(self, beam_idx: torch.LongTensor):
-        """Reorder all layers for beam search."""
-        for layer in self.layers:
-            layer.reorder_batch(beam_idx)
-
-    def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:
-        """
-        Convert to tuple format for compatibility.
-
-        Returns:
-            Tuple of (K, V) for each layer
-        """
-        return tuple(
-            (layer.keys, layer.values)
-            for layer in self.layers
-        )
-
-    @staticmethod
-    def from_legacy_cache(
-        past_key_values: Tuple[Tuple[torch.Tensor, torch.Tensor], ...]
-    ) -> 'INLCache':
-        """
-        Create INLCache from legacy tuple format.
-
-        Args:
-            past_key_values: Tuple of (K, V) for each layer
-        """
-        num_layers = len(past_key_values)
-        cache = INLCache(num_layers)
-
-        for layer_idx, (keys, values) in enumerate(past_key_values):
-            cache.layers[layer_idx].keys = keys
-            cache.layers[layer_idx].values = values
-
-        return cache
-
-
-class INLCachedAttention(nn.Module):
-    """
-    Multi-head self-attention with KV cache support.
-
-    Replaces nn.MultiheadAttention with a cache-aware implementation.
-    Compatible with INL-LLM's architecture and optimizations.
-    """
-
-    def __init__(
-        self,
-        embed_dim: int,
-        num_heads: int,
-        dropout: float = 0.0,
-        bias: bool = True
-    ):
-        super().__init__()
-
-        if embed_dim % num_heads != 0:
-            raise ValueError(f"embed_dim ({embed_dim}) must be divisible by num_heads ({num_heads})")
-
-        self.embed_dim = embed_dim
-        self.num_heads = num_heads
-        self.head_dim = embed_dim // num_heads
-        self.dropout = dropout
-
-        # Combined QKV projection (more efficient)
-        self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=bias)
-
-        # Output projection
-        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
-
-        self.attn_dropout = nn.Dropout(dropout)
-        self.resid_dropout = nn.Dropout(dropout)
-
-        # Initialize weights
-        self._reset_parameters()
-
-    def _reset_parameters(self):
-        """Initialize parameters like nn.MultiheadAttention."""
-        nn.init.xavier_uniform_(self.qkv_proj.weight)
-        if self.qkv_proj.bias is not None:
-            nn.init.constant_(self.qkv_proj.bias, 0.0)
-
-        nn.init.xavier_uniform_(self.out_proj.weight)
-        if self.out_proj.bias is not None:
-            nn.init.constant_(self.out_proj.bias, 0.0)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        attn_mask: Optional[torch.Tensor] = None,
-        cache_layer: Optional[INLCacheLayer] = None,
-        use_cache: bool = False
-    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
-        """
-        Forward pass with optional KV caching.
-
-        Args:
-            x: Input tensor [batch_size, seq_len, embed_dim]
-            attn_mask: Attention mask [seq_len, seq_len] or [tgt_len, src_len]
-            cache_layer: Cache layer to update (if using cache)
-            use_cache: Whether to use/update cache
-
-        Returns:
-            attn_output: [batch_size, seq_len, embed_dim]
-            new_cache: Updated (keys, values) if use_cache else None
-        """
-        batch_size, seq_len, embed_dim = x.shape
-
-        # Compute Q, K, V
-        qkv = self.qkv_proj(x)  # [B, S, 3*D]
-        qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)
-        qkv = qkv.permute(2, 0, 3, 1, 4)  # [3, B, num_heads, S, head_dim]
-        q, k, v = qkv[0], qkv[1], qkv[2]
-
-        # Handle cache
-        if use_cache and cache_layer is not None:
-            # Update cache with new K, V
-            k, v = cache_layer.update_attention(k, v)
-
-        # Compute attention scores
-        # q: [B, num_heads, tgt_len, head_dim]
-        # k: [B, num_heads, src_len, head_dim]
-        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
-        # attn_weights: [B, num_heads, tgt_len, src_len]
-
-        # Apply attention mask (causal mask for autoregressive generation)
-        if attn_mask is not None:
-            # attn_mask is [tgt_len, src_len] boolean mask (True = masked position)
-            # Expand for batch and heads
-            attn_mask = attn_mask.unsqueeze(0).unsqueeze(0)  # [1, 1, tgt_len, src_len]
-            attn_weights = attn_weights.masked_fill(attn_mask, float('-inf'))
-
-        # Softmax
-        attn_weights = F.softmax(attn_weights, dim=-1)
-        attn_weights = self.attn_dropout(attn_weights)
-
-        # Apply attention to values
-        # v: [B, num_heads, src_len, head_dim]
-        attn_output = torch.matmul(attn_weights, v)  # [B, num_heads, tgt_len, head_dim]
-
-        # Reshape and project
-        attn_output = attn_output.transpose(1, 2)  # [B, tgt_len, num_heads, head_dim]
-        attn_output = attn_output.reshape(batch_size, seq_len, embed_dim)
-        attn_output = self.out_proj(attn_output)
-        attn_output = self.resid_dropout(attn_output)
-
-        # Return cache if requested
-        cache_output = (k, v) if use_cache else None
-
-        return attn_output, cache_output
-
-
-class PositionalEncoding(nn.Module):
-    """Positional encoding."""
-    def __init__(self, d_model: int, max_len: int = 5000):
-        super().__init__()
-        pe = torch.zeros(max_len, d_model)
-        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
-        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
-        pe[:, 0::2] = torch.sin(position * div_term)
-        pe[:, 1::2] = torch.cos(position * div_term)
-        self.register_buffer('pe', pe.unsqueeze(0))
-
-    def forward(self, x, start_pos: int = 0):
-        """
-        Apply positional encoding.
-
-        Args:
-            x: Input tensor [batch_size, seq_len, d_model]
-            start_pos: Starting position for positional encoding (for KV cache)
-
-        Returns:
-            x with positional encoding added
-        """
-        seq_len = x.size(1)
-        return x + self.pe[:, start_pos:start_pos + seq_len, :]
-
-
-class UltraOptimizedINLBlock(nn.Module):
-    """
-    Ultra-optimized INL block with all optimizations enabled.
-
-    Uses:
-    - Shared controllers (across all blocks in the model)
-    - Hierarchical equilibrium
-    - Sparse harmonic excitation
-    - Adaptive early stopping
-    - Gradient checkpointing
-    """
-
-    def __init__(
-        self,
-        d_model: int,
-        num_heads: int,
-        num_iterations: int,
-        shared_controller: SharedController,
-        layer_idx: int,
-        feedforward_dim: int,
-        dropout: float = 0.1,
-        use_gradient_checkpointing: bool = False,
-        use_adaptive_stopping: bool = True,
-        adaptive_convergence_threshold: float = 0.001,
-        group_size: int = 64,
-        excitation_sparsity: float = 0.1
-    ):
-        super().__init__()
-
-        self.d_model = d_model
-        self.num_iterations = num_iterations
-        self.layer_idx = layer_idx
-        self.shared_controller = shared_controller
-        self.use_adaptive_stopping = use_adaptive_stopping
-
-        # Norms
-        self.norm1 = nn.LayerNorm(d_model)
-        self.norm2 = nn.LayerNorm(d_model)
-        self.norm_attn = nn.LayerNorm(d_model)
-
-        # Attention with KV cache support
-        self.attention = INLCachedAttention(
-            embed_dim=d_model,
-            num_heads=num_heads,
-            dropout=dropout
-        )
-
-        # Ultra-optimized INL
-        # Use hierarchical equilibrium + sparse excitation
-        # Wrap with adaptive stopping for 3× faster inference
-        base_inl = HierarchicalEquilibriumINL(
-            hidden_dim=d_model,
-            output_dim=d_model,
-            group_size=group_size,
-            target_value=0.0,
-            dt=0.1
-        )
-
-        if use_adaptive_stopping:
-            self.inl = AdaptiveHierarchicalINL(
-                inl_layer=base_inl,
-                convergence_threshold=adaptive_convergence_threshold,
-                min_iterations=3,
-                max_iterations=num_iterations,
-                check_interval=1
-            )
-        else:
-            self.inl = base_inl
-
-        # Feedforward
-        self.ff = nn.Sequential(
-            nn.Linear(d_model, feedforward_dim),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(feedforward_dim, d_model),
-            nn.Dropout(dropout)
-        )
-
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        mask: Optional[torch.Tensor] = None,
-        cache_layer: Optional[INLCacheLayer] = None,
-        use_cache: bool = False
-    ) -> Tuple[torch.Tensor, Dict]:
-        batch_size, seq_len, d_model = x.shape
-
-        # Step 1: Attention with KV cache
-        x_norm = self.norm_attn(x)
-
-        # Build causal mask
-        if use_cache and cache_layer is not None:
-            # During generation with cache: mask is for new tokens attending to all previous tokens
-            past_len = cache_layer.get_seq_length()
-            total_len = past_len + seq_len
-            # Create mask [seq_len, total_len] where each new token can attend to all previous + itself
-            attn_mask = torch.zeros(seq_len, total_len, device=x.device, dtype=torch.bool)
-            # Only mask future tokens within the new sequence
-            if seq_len > 1:
-                new_causal_mask = torch.triu(
-                    torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool),
-                    diagonal=1
-                )
-                attn_mask[:, past_len:] = new_causal_mask
-        elif mask is None:
-            # Standard causal mask for full sequence
-            attn_mask = torch.triu(
-                torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool),
-                diagonal=1
-            )
-        else:
-            attn_mask = mask
-
-        attn_output, _ = self.attention(x_norm, attn_mask=attn_mask, cache_layer=cache_layer, use_cache=use_cache)
-        x = x + self.dropout(attn_output)
-        context = attn_output
-
-        # Step 2: INL Dynamics (ultra-optimized with adaptive early stopping)
-        x_norm = self.norm1(x)
-
-        # Initialize integrator states (x, v)
-        # NOTE: We always initialize fresh for each forward pass.
-        # The integrator dynamics run WITHIN each layer, not across tokens.
-        # The cache is ONLY for attention K,V to avoid recomputing attention over past tokens.
-        x_state = x_norm.clone()
-        v_state = torch.zeros_like(x_norm)
-
-        # Flatten for INL processing
-        x_flat_init = x_state.reshape(batch_size * seq_len, d_model)
-        v_flat_init = v_state.reshape(batch_size * seq_len, d_model)
-        ctx_flat = context.reshape(batch_size * seq_len, d_model)
-
-        # Use adaptive forward if available (inference mode with early stopping)
-        if self.use_adaptive_stopping and hasattr(self.inl, 'forward_adaptive') and not self.training:
-            # ✅ Adaptive early stopping (3× faster inference)
-            x_final_flat, v_final_flat, adaptive_result = self.inl.forward_adaptive(
-                ctx_flat,
-                x_flat_init,
-                v_flat_init,
-                num_iterations=self.num_iterations,
-                use_early_stopping=True,
-                return_trajectory=True
-            )
-
-            # Get trajectories from adaptive result
-            if 'x_trajectory' in adaptive_result:
-                x_traj_flat = adaptive_result['x_trajectory']  # [B*S, T+1, D]
-                v_traj_flat = adaptive_result['v_trajectory']  # [B*S, T+1, D]
-            else:
-                # Fallback: single final state
-                x_traj_flat = x_final_flat.unsqueeze(1)
-                v_traj_flat = v_final_flat.unsqueeze(1)
-
-            aux_infos = {
-                'x': x_traj_flat,
-                'v': v_traj_flat,
-                'mu': adaptive_result.get('mu'),
-                'mu_global': adaptive_result.get('mu_global'),
-                'mu_offsets': adaptive_result.get('mu_offsets'),
-                'iterations_used': adaptive_result.get('iterations_used'),
-                'avg_iterations': adaptive_result.get('avg_iterations')
-            }
-
-            output = x_final_flat.reshape(batch_size, seq_len, d_model)
-
-        else:
-            # Standard training mode (all iterations)
-            x_trajectory = [x_flat_init.clone()]
-            v_trajectory = [v_flat_init.clone()]
-
-            x_flat, v_flat = x_flat_init, v_flat_init
-
-            for iteration in range(self.num_iterations):
-                x_next_flat, v_next_flat, aux = self.inl(ctx_flat, x_flat, v_flat, step=iteration)
-
-                x_flat, v_flat = x_next_flat, v_next_flat
-
-                # Save trajectories for loss computation
-                x_trajectory.append(x_flat.clone())
-                v_trajectory.append(v_flat.clone())
-
-            # Stack trajectories: [B*S, T+1, D]
-            x_traj_flat = torch.stack(x_trajectory, dim=1)
-            v_traj_flat = torch.stack(v_trajectory, dim=1)
-
-            aux_infos = {
-                'x': x_traj_flat,
-                'v': v_traj_flat,
-                'mu': aux.get('mu', None),
-                'mu_global': aux.get('mu_global', None),
-                'mu_offsets': aux.get('mu_offsets', None)
-            }
-
-            output = x_flat.reshape(batch_size, seq_len, d_model)
-
-        # NOTE: No need to update integrator cache - we don't cache x, v states
-        # since integrator dynamics are computed fresh for each token.
-
-        # Residual
-        x = x + self.dropout(output)
-
-        # Feedforward
-        x = x + self.ff(self.norm2(x))
-
-        return x, aux_infos
-
-
-class UltraOptimizedIntegratorLanguageModel(nn.Module):
-    """
-    ULTRA-OPTIMIZED INL-LLM
-
-    All optimizations enabled by default:
-    ✅ Low-rank embeddings (87% reduction)
-    ✅ Gradient checkpointing (60% memory save)
-    ✅ Adaptive early stopping (40% faster)
-    ✅ Shared controllers (96% controller reduction)
-    ✅ Hierarchical equilibrium (98% μ reduction)
-    ✅ Sparse excitation (10x less compute)
-
-    Can scale to 100B+ parameters efficiently!
-    """
-
-    def __init__(
-        self,
-        vocab_size: int,
-        d_model: int = 512,
-        num_layers: int = 6,
-        num_heads: int = 8,
-        num_iterations_per_layer: int = 5,
-        feedforward_dim: int = None,
-        max_seq_len: int = 2048,
-        dropout: float = 0.1,
-        # Optimization flags
-        use_lowrank_embeddings: bool = True,
-        lowrank_ratio: float = 0.125,
-        use_gradient_checkpointing: bool = True,
-        use_shared_controllers: bool = True,
-        use_adaptive_stopping: bool = True,
-        adaptive_convergence_threshold: float = 0.001,
-        hierarchical_group_size: int = 64,
-        excitation_sparsity: float = 0.1
-    ):
-        super().__init__()
-
-        self.vocab_size = vocab_size
-        self.d_model = d_model
-        self.num_layers = num_layers
-
-        if feedforward_dim is None:
-            feedforward_dim = 4 * d_model
-
-        # Low-rank embeddings
-        if use_lowrank_embeddings:
-            self.token_embedding = LowRankEmbedding(vocab_size, d_model, rank_ratio=lowrank_ratio)
-            print(f"✅ Low-Rank Embeddings: {self.token_embedding}")
-        else:
-            self.token_embedding = nn.Embedding(vocab_size, d_model)
-
-        # Positional encoding
-        self.pos_encoding = PositionalEncoding(d_model, max_seq_len)
-        self.dropout = nn.Dropout(dropout)
-
-        # Shared controller (ONE for all layers!)
-        if use_shared_controllers:
-            self.shared_controller = SharedController(
-                hidden_dim=d_model,
-                output_dim=d_model,
-                num_layers=num_layers,
-                hidden_controller=64
-            )
-            print(f"✅ Shared Controllers: {self.shared_controller.num_parameters():,} params for {num_layers} layers")
-        else:
-            self.shared_controller = None
-
-        # Layers
-        self.layers = nn.ModuleList([
-            UltraOptimizedINLBlock(
-                d_model=d_model,
-                num_heads=num_heads,
-                num_iterations=num_iterations_per_layer,
-                shared_controller=self.shared_controller,
-                layer_idx=i,
-                feedforward_dim=feedforward_dim,
-                dropout=dropout,
-                use_gradient_checkpointing=use_gradient_checkpointing,
-                use_adaptive_stopping=use_adaptive_stopping,
-                adaptive_convergence_threshold=adaptive_convergence_threshold,
-                group_size=hierarchical_group_size,
-                excitation_sparsity=excitation_sparsity
-            )
-            for i in range(num_layers)
-        ])
-
-        # Final norm
-        self.final_norm = nn.LayerNorm(d_model)
-
-        # LM head
-        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
-
-        # Initialize
-        self._init_weights()
-        self._print_optimization_status()
-
-    def _init_weights(self):
-        """Initialize weights."""
-        if not isinstance(self.token_embedding, LowRankEmbedding):
-            with torch.no_grad():
-                nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
-
-        with torch.no_grad():
-            nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.02)
-
-    def _print_optimization_status(self):
-        """Print optimization summary."""
-        print("\n" + "=" * 70)
-        print("ULTRA-OPTIMIZED INL-LLM")
-        print("=" * 70)
-        print("LEVEL 1 (Basic Optimizations):")
-        print(f"  ✅ Low-rank embeddings")
-        print(f"  ✅ Gradient checkpointing")
-        print(f"  ✅ Adaptive early stopping")
-        print("\nLEVEL 2 (Advanced Optimizations):")
-        print(f"  ✅ Shared controllers (across {self.num_layers} layers)")
-        print(f"  ✅ Hierarchical equilibrium")
-        print(f"  ✅ Sparse harmonic excitation")
-        print(f"\nTotal parameters: {self.get_num_params():,}")
-        print("=" * 70 + "\n")
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        past_key_values: Optional[INLCache] = None,
-        use_cache: bool = False,
-        return_aux: bool = False
-    ) -> Tuple[torch.Tensor, Optional[List], Optional[INLCache]]:
-        """
-        Forward pass with optional KV caching.
-
-        Args:
-            input_ids: Input token IDs [batch_size, seq_len]
-            attention_mask: Attention mask (optional)
-            past_key_values: Previous cache (INLCache object)
-            use_cache: Whether to use/update cache
-            return_aux: Whether to return auxiliary info
-
-        Returns:
-            logits: Output logits [batch_size, seq_len, vocab_size]
-            all_aux: Auxiliary info from each layer (if return_aux=True)
-            new_cache: Updated cache (if use_cache=True)
-        """
-        # Initialize cache if needed
-        if use_cache and past_key_values is None:
-            past_key_values = INLCache(num_layers=self.num_layers)
-
-        # Determine starting position for positional encoding
-        start_pos = 0
-        if use_cache and past_key_values is not None:
-            start_pos = past_key_values.get_seq_length()
-
-        # Embedding with correct positional encoding
-        x = self.token_embedding(input_ids)
-        x = self.pos_encoding(x, start_pos=start_pos)
-        x = self.dropout(x)
-
-        # Layers
-        all_aux = [] if return_aux else None
-
-        for layer_idx, layer in enumerate(self.layers):
-            cache_layer = past_key_values[layer_idx] if use_cache else None
-            x, aux = layer(x, mask=attention_mask, cache_layer=cache_layer, use_cache=use_cache)
-            if return_aux:
-                all_aux.append(aux)
-
-        # Final norm
-        x = self.final_norm(x)
-
-        # LM head
-        logits = self.lm_head(x)
-
-        return logits, all_aux, past_key_values if use_cache else None
-
-    def generate(
-        self,
-        input_ids: torch.Tensor,
-        max_new_tokens: int = 100,
-        temperature: float = 1.0,
-        top_k: Optional[int] = None,
-        top_p: Optional[float] = None,
-        do_sample: bool = True,
-        use_cache: bool = True
-    ) -> torch.Tensor:
-        """
-        Autoregressive generation with optional KV caching.
-
-        Args:
-            input_ids: Input token IDs [batch_size, seq_len]
-            max_new_tokens: Number of tokens to generate
-            temperature: Sampling temperature
-            top_k: Top-k sampling (if provided)
-            top_p: Nucleus sampling threshold (if provided)
-            do_sample: Whether to sample or use greedy decoding
-            use_cache: Whether to use KV caching (default: True, much faster!)
-
-        Returns:
-            Generated token IDs [batch_size, seq_len + max_new_tokens]
-        """
-        self.eval()
-        past_key_values = None
-
-        with torch.no_grad():
-            for step in range(max_new_tokens):
-                # Use cache for all steps after the first
-                if use_cache and step > 0:
-                    # Only pass the last token for cached generation
-                    model_input = input_ids[:, -1:]
-                    logits, _, past_key_values = self.forward(
-                        model_input,
-                        past_key_values=past_key_values,
-                        use_cache=True
-                    )
-                else:
-                    # First step or no cache: process full sequence
-                    logits, _, past_key_values = self.forward(
-                        input_ids,
-                        past_key_values=past_key_values if use_cache else None,
-                        use_cache=use_cache
-                    )
-
-                # Get logits for last token
-                logits = logits[:, -1, :] / temperature
-
-                # Apply top-k filtering
-                if top_k is not None:
-                    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
-                    logits[indices_to_remove] = float('-inf')
-
-                # Apply top-p (nucleus) filtering
-                if top_p is not None:
-                    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
-                    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
-                    sorted_indices_to_remove = cumulative_probs > top_p
-                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-                    sorted_indices_to_remove[..., 0] = 0
-                    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
-                    logits[indices_to_remove] = float('-inf')
-
-                # Sample or select greedily
-                if do_sample:
-                    probs = F.softmax(logits, dim=-1)
-                    next_token = torch.multinomial(probs, num_samples=1)
-                else:
-                    next_token = torch.argmax(logits, dim=-1, keepdim=True)
-
-                # Append to sequence
-                input_ids = torch.cat([input_ids, next_token], dim=1)
-
-        return input_ids
-
-    def get_num_params(self) -> int:
-        """Count parameters."""
-        return sum(p.numel() for p in self.parameters())
-
-    def get_inference_stats(self) -> Dict:
-        """
-        Get model statistics and optimization info.
-
-        Returns dict with model configuration and enabled optimizations.
-        """
-        stats = {
-            'num_params': self.get_num_params(),
-            'num_layers': self.num_layers,
-            'd_model': self.d_model,
-            'optimizations_enabled': {
-                'low_rank_embeddings': True,
-                'shared_controllers': True,
-                'hierarchical_equilibrium': True,
-                'sparse_excitation': True,
-                'gradient_checkpointing': True
-            }
-        }
-        return stats
-
-
-def create_ultra_optimized_model(
-    size: str = 'small',
-    vocab_size: int = 50000
-) -> UltraOptimizedIntegratorLanguageModel:
-    """
-    Create ultra-optimized model.
-
-    Sizes: 'small', 'medium', 'large', 'xlarge', '3B', '7B', '13B', '30B', '70B'
-    """
-    configs = {
-        'small': {'d_model': 512, 'num_layers': 6, 'num_heads': 8, 'iterations': 5, 'ff_dim': 2048},
-        'medium': {'d_model': 768, 'num_layers': 12, 'num_heads': 12, 'iterations': 7, 'ff_dim': 3072},
-        'large': {'d_model': 1024, 'num_layers': 24, 'num_heads': 16, 'iterations': 10, 'ff_dim': 4096},
-        'xlarge': {'d_model': 1536, 'num_layers': 32, 'num_heads': 24, 'iterations': 12, 'ff_dim': 6144},
-        '3B': {'d_model': 2048, 'num_layers': 40, 'num_heads': 32, 'iterations': 15, 'ff_dim': 8192},
-        '7B': {'d_model': 4096, 'num_layers': 32, 'num_heads': 32, 'iterations': 10, 'ff_dim': 16384},
-        '13B': {'d_model': 5120, 'num_layers': 40, 'num_heads': 40, 'iterations': 12, 'ff_dim': 20480},
-        '30B': {'d_model': 6656, 'num_layers': 60, 'num_heads': 52, 'iterations': 12, 'ff_dim': 26624},
-        '70B': {'d_model': 8192, 'num_layers': 80, 'num_heads': 64, 'iterations': 12, 'ff_dim': 32768},
-    }
-
-    if size not in configs:
-        raise ValueError(f"Size must be one of {list(configs.keys())}")
-
-    cfg = configs[size]
-
-    model = UltraOptimizedIntegratorLanguageModel(
-        vocab_size=vocab_size,
-        d_model=cfg['d_model'],
-        num_layers=cfg['num_layers'],
-        num_heads=cfg['num_heads'],
-        num_iterations_per_layer=cfg['iterations'],
-        feedforward_dim=cfg['ff_dim'],
-        max_seq_len=2048,
-        # All optimizations enabled
-        use_lowrank_embeddings=True,
-        lowrank_ratio=0.125,
-        use_gradient_checkpointing=True,
-        use_shared_controllers=True,
-        hierarchical_group_size=64,
-        excitation_sparsity=0.1
-    )
-
-    print(f"\n🚀 ULTRA-OPTIMIZED INL-LLM ({size}): {model.get_num_params():,} parameters")
-    print(f"   Ready to scale to 100B+ with maximum efficiency!\n")
-
-    return model
-
-
-if __name__ == '__main__':
-    # Fix imports for standalone execution
-    import sys
-    import os
-    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
-
-    from inl_llm import create_model
-
-    print("\n" + "=" * 70)
-    print("INL-LLM MODEL - Test")
-    print("=" * 70 + "\n")
-
-    # Create model
-    model = create_model(size='medium', vocab_size=50000)
-
-    # Test forward
-    batch_size = 2
-    seq_len = 10
-    input_ids = torch.randint(0, 50000, (batch_size, seq_len))
-
-    print("Running forward pass...")
-    logits, aux = model(input_ids, return_aux=True)
-
-    print(f"✅ Input shape: {input_ids.shape}")
-    print(f"✅ Output shape: {logits.shape}")
-    print(f"✅ Aux layers: {len(aux)}")
-
-    # Test generation
-    print("\nTesting generation...")
-    prompt = torch.randint(0, 50000, (1, 5))
-    generated = model.generate(prompt, max_new_tokens=20, temperature=0.8)
-
-    print(f"✅ Prompt length: {prompt.shape[1]}")
-    print(f"✅ Generated length: {generated.shape[1]}")
-
-    print("\n" + "=" * 70)
-    print("✅ INL-LLM WORKING PERFECTLY!")
-    print("=" * 70 + "\n")
+"""
+ULTRA-Optimized Integrator Language Model (INL-LLM)
+
+Combines ALL optimizations for maximum efficiency:
+
+LEVEL 1 (Basic):
+- Low-rank embeddings (-70-80% embedding params)
+- Gradient checkpointing (-50-70% memory)
+- Adaptive early stopping (+30-50% inference speed)
+
+LEVEL 2 (Advanced):
+- Shared controllers (-96% controller params)
+- Sparse harmonic excitation (10x less compute)
+- Hierarchical equilibrium (-98% equilibrium params)
+
+RESULT: Can scale to 100B+ parameters with MUCH higher efficiency
+
+Author: Boris Peyriguère
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, Dict, List
+import math
+
+from ..optimizations.optimizations import (
+    LowRankEmbedding,
+    GradientCheckpointedINL,
+    AdaptiveIntegratorNeuronLayer,
+    AdaptiveHierarchicalINL
+)
+from ..optimizations.advanced_optimizations import (
+    SharedController,
+    SparseHarmonicINL,
+    HierarchicalEquilibriumINL
+)
+from ..core.adaptive_budget_allocator import (
+    AdaptiveBudgetAllocator,
+    BudgetAwareINLLayer,
+    create_budget_allocator
+)
+from ..core.moe_controller import (
+    INLMixtureOfExperts,
+    create_moe_controller
+)
+from ..core.moe_budget_integration import (
+    MoEBudgetAwareINLLayer
+)
+
+
+# ============================================================================
+# KV CACHE SUPPORT FOR INL-LLM
+# ============================================================================
+
+class INLCacheLayer:
+    """
+    Cache for a single layer, storing:
+    - Attention K, V (standard transformer cache)
+
+    NOTE: We do NOT cache integrator x, v states because integrator dynamics
+    run WITHIN each layer for each token, not across tokens. Only attention
+    needs to look back at previous tokens' K, V.
+    """
+
+    def __init__(self):
+        self.keys: Optional[torch.Tensor] = None          # [B, num_heads, seq_len, head_dim]
+        self.values: Optional[torch.Tensor] = None        # [B, num_heads, seq_len, head_dim]
+
+    def update_attention(
+        self,
+        new_keys: torch.Tensor,
+        new_values: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Update attention cache with new K, V.
+
+        Args:
+            new_keys: [B, num_heads, new_seq_len, head_dim]
+            new_values: [B, num_heads, new_seq_len, head_dim]
+
+        Returns:
+            Full keys, values (concatenated with past)
+        """
+        if self.keys is None:
+            # First time: initialize cache
+            self.keys = new_keys
+            self.values = new_values
+        else:
+            # Concatenate along sequence dimension
+            self.keys = torch.cat([self.keys, new_keys], dim=2)
+            self.values = torch.cat([self.values, new_values], dim=2)
+
+        return self.keys, self.values
+
+    def get_seq_length(self) -> int:
+        """Get current sequence length in cache."""
+        if self.keys is not None:
+            return self.keys.shape[2]
+        return 0
+
+    def reorder_batch(self, beam_idx: torch.LongTensor):
+        """Reorder cache for beam search."""
+        if self.keys is not None:
+            device = self.keys.device
+            self.keys = self.keys.index_select(0, beam_idx.to(device))
+            self.values = self.values.index_select(0, beam_idx.to(device))
+
+
+class INLCache:
+    """
+    Complete cache for INL-LLM model.
+
+    Stores attention K, V for all layers.
+    Compatible with HuggingFace's past_key_values interface.
+
+    NOTE: Simpler than typical transformers - we only cache attention K, V,
+    not integrator states since those are computed fresh for each token.
+    """
+
+    def __init__(self, num_layers: int):
+        self.num_layers = num_layers
+        self.layers: List[INLCacheLayer] = [INLCacheLayer() for _ in range(num_layers)]
+
+    def __getitem__(self, layer_idx: int) -> INLCacheLayer:
+        """Access cache for specific layer."""
+        return self.layers[layer_idx]
+
+    def __len__(self) -> int:
+        """Number of layers."""
+        return self.num_layers
+
+    def get_seq_length(self, layer_idx: int = 0) -> int:
+        """Get current sequence length (all layers should be same)."""
+        return self.layers[layer_idx].get_seq_length()
+
+    def reorder_cache(self, beam_idx: torch.LongTensor):
+        """Reorder all layers for beam search."""
+        for layer in self.layers:
+            layer.reorder_batch(beam_idx)
+
+    def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:
+        """
+        Convert to tuple format for compatibility.
+
+        Returns:
+            Tuple of (K, V) for each layer
+        """
+        return tuple(
+            (layer.keys, layer.values)
+            for layer in self.layers
+        )
+
+    @staticmethod
+    def from_legacy_cache(
+        past_key_values: Tuple[Tuple[torch.Tensor, torch.Tensor], ...]
+    ) -> 'INLCache':
+        """
+        Create INLCache from legacy tuple format.
+
+        Args:
+            past_key_values: Tuple of (K, V) for each layer
+        """
+        num_layers = len(past_key_values)
+        cache = INLCache(num_layers)
+
+        for layer_idx, (keys, values) in enumerate(past_key_values):
+            cache.layers[layer_idx].keys = keys
+            cache.layers[layer_idx].values = values
+
+        return cache
+
+
+class INLCachedAttention(nn.Module):
+    """
+    Multi-head self-attention with KV cache support.
+
+    Replaces nn.MultiheadAttention with a cache-aware implementation.
+    Compatible with INL-LLM's architecture and optimizations.
+    """
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        bias: bool = True
+    ):
+        super().__init__()
+
+        if embed_dim % num_heads != 0:
+            raise ValueError(f"embed_dim ({embed_dim}) must be divisible by num_heads ({num_heads})")
+
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.dropout = dropout
+
+        # Combined QKV projection (more efficient)
+        self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=bias)
+
+        # Output projection
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        self.attn_dropout = nn.Dropout(dropout)
+        self.resid_dropout = nn.Dropout(dropout)
+
+        # Initialize weights
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        """Initialize parameters like nn.MultiheadAttention."""
+        nn.init.xavier_uniform_(self.qkv_proj.weight)
+        if self.qkv_proj.bias is not None:
+            nn.init.constant_(self.qkv_proj.bias, 0.0)
+
+        nn.init.xavier_uniform_(self.out_proj.weight)
+        if self.out_proj.bias is not None:
+            nn.init.constant_(self.out_proj.bias, 0.0)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        cache_layer: Optional[INLCacheLayer] = None,
+        use_cache: bool = False
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+        """
+        Forward pass with optional KV caching.
+
+        Args:
+            x: Input tensor [batch_size, seq_len, embed_dim]
+            attn_mask: Attention mask [seq_len, seq_len] or [tgt_len, src_len]
+            cache_layer: Cache layer to update (if using cache)
+            use_cache: Whether to use/update cache
+
+        Returns:
+            attn_output: [batch_size, seq_len, embed_dim]
+            new_cache: Updated (keys, values) if use_cache else None
+        """
+        batch_size, seq_len, embed_dim = x.shape
+
+        # Compute Q, K, V
+        qkv = self.qkv_proj(x)  # [B, S, 3*D]
+        qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)
+        qkv = qkv.permute(2, 0, 3, 1, 4)  # [3, B, num_heads, S, head_dim]
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        # Handle cache
+        if use_cache and cache_layer is not None:
+            # Update cache with new K, V
+            k, v = cache_layer.update_attention(k, v)
+
+        # Compute attention scores
+        # q: [B, num_heads, tgt_len, head_dim]
+        # k: [B, num_heads, src_len, head_dim]
+        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        # attn_weights: [B, num_heads, tgt_len, src_len]
+
+        # Apply attention mask (causal mask for autoregressive generation)
+        if attn_mask is not None:
+            # attn_mask is [tgt_len, src_len] boolean mask (True = masked position)
+            # Expand for batch and heads
+            attn_mask = attn_mask.unsqueeze(0).unsqueeze(0)  # [1, 1, tgt_len, src_len]
+            attn_weights = attn_weights.masked_fill(attn_mask, float('-inf'))
+
+        # Softmax
+        attn_weights = F.softmax(attn_weights, dim=-1)
+        attn_weights = self.attn_dropout(attn_weights)
+
+        # Apply attention to values
+        # v: [B, num_heads, src_len, head_dim]
+        attn_output = torch.matmul(attn_weights, v)  # [B, num_heads, tgt_len, head_dim]
+
+        # Reshape and project
+        attn_output = attn_output.transpose(1, 2)  # [B, tgt_len, num_heads, head_dim]
+        attn_output = attn_output.reshape(batch_size, seq_len, embed_dim)
+        attn_output = self.out_proj(attn_output)
+        attn_output = self.resid_dropout(attn_output)
+
+        # Return cache if requested
+        cache_output = (k, v) if use_cache else None
+
+        return attn_output, cache_output
+
+
+class PositionalEncoding(nn.Module):
+    """Positional encoding."""
+    def __init__(self, d_model: int, max_len: int = 5000):
+        super().__init__()
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe.unsqueeze(0))
+
+    def forward(self, x, start_pos: int = 0):
+        """
+        Apply positional encoding.
+
+        Args:
+            x: Input tensor [batch_size, seq_len, d_model]
+            start_pos: Starting position for positional encoding (for KV cache)
+
+        Returns:
+            x with positional encoding added
+        """
+        seq_len = x.size(1)
+        return x + self.pe[:, start_pos:start_pos + seq_len, :]
+
+
+class UltraOptimizedINLBlock(nn.Module):
+    """
+    Ultra-optimized INL block with all optimizations enabled.
+
+    Uses:
+    - Shared controllers (across all blocks in the model)
+    - Hierarchical equilibrium
+    - Sparse harmonic excitation
+    - Adaptive early stopping
+    - Gradient checkpointing
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        num_iterations: int,
+        shared_controller: SharedController,
+        layer_idx: int,
+        feedforward_dim: int,
+        dropout: float = 0.1,
+        use_gradient_checkpointing: bool = False,
+        use_adaptive_stopping: bool = True,
+        adaptive_convergence_threshold: float = 0.001,
+        group_size: int = 64,
+        excitation_sparsity: float = 0.1,
+        budget_allocator: Optional[AdaptiveBudgetAllocator] = None,
+        moe_controller: Optional[INLMixtureOfExperts] = None
+    ):
+        super().__init__()
+
+        self.d_model = d_model
+        self.num_iterations = num_iterations
+        self.layer_idx = layer_idx
+        self.shared_controller = shared_controller
+        self.use_adaptive_stopping = use_adaptive_stopping
+        self.budget_allocator = budget_allocator
+        self.moe_controller = moe_controller
+
+        # Norms
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm_attn = nn.LayerNorm(d_model)
+
+        # Attention with KV cache support
+        self.attention = INLCachedAttention(
+            embed_dim=d_model,
+            num_heads=num_heads,
+            dropout=dropout
+        )
+
+        # Ultra-optimized INL
+        # Use hierarchical equilibrium + sparse excitation
+        # Wrap with adaptive stopping for 3× faster inference
+        base_inl = HierarchicalEquilibriumINL(
+            hidden_dim=d_model,
+            output_dim=d_model,
+            group_size=group_size,
+            target_value=0.0,
+            dt=0.1
+        )
+
+        if use_adaptive_stopping:
+            self.inl = AdaptiveHierarchicalINL(
+                inl_layer=base_inl,
+                convergence_threshold=adaptive_convergence_threshold,
+                min_iterations=3,
+                max_iterations=num_iterations,
+                check_interval=1
+            )
+        else:
+            self.inl = base_inl
+
+        # Feedforward
+        self.ff = nn.Sequential(
+            nn.Linear(d_model, feedforward_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(feedforward_dim, d_model),
+            nn.Dropout(dropout)
+        )
+
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        cache_layer: Optional[INLCacheLayer] = None,
+        use_cache: bool = False
+    ) -> Tuple[torch.Tensor, Dict]:
+        batch_size, seq_len, d_model = x.shape
+
+        # Step 1: Attention with KV cache
+        x_norm = self.norm_attn(x)
+
+        # Build causal mask
+        if use_cache and cache_layer is not None:
+            # During generation with cache: mask is for new tokens attending to all previous tokens
+            past_len = cache_layer.get_seq_length()
+            total_len = past_len + seq_len
+            # Create mask [seq_len, total_len] where each new token can attend to all previous + itself
+            attn_mask = torch.zeros(seq_len, total_len, device=x.device, dtype=torch.bool)
+            # Only mask future tokens within the new sequence
+            if seq_len > 1:
+                new_causal_mask = torch.triu(
+                    torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool),
+                    diagonal=1
+                )
+                attn_mask[:, past_len:] = new_causal_mask
+        elif mask is None:
+            # Standard causal mask for full sequence
+            attn_mask = torch.triu(
+                torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool),
+                diagonal=1
+            )
+        else:
+            attn_mask = mask
+
+        attn_output, _ = self.attention(x_norm, attn_mask=attn_mask, cache_layer=cache_layer, use_cache=use_cache)
+        x = x + self.dropout(attn_output)
+        context = attn_output
+
+        # Step 2: INL Dynamics (ultra-optimized with adaptive early stopping)
+        x_norm = self.norm1(x)
+
+        # Initialize integrator states (x, v)
+        # NOTE: We always initialize fresh for each forward pass.
+        # The integrator dynamics run WITHIN each layer, not across tokens.
+        # The cache is ONLY for attention K,V to avoid recomputing attention over past tokens.
+        x_state = x_norm.clone()
+        v_state = torch.zeros_like(x_norm)
+
+        # Flatten for INL processing
+        x_flat_init = x_state.reshape(batch_size * seq_len, d_model)
+        v_flat_init = v_state.reshape(batch_size * seq_len, d_model)
+        ctx_flat = context.reshape(batch_size * seq_len, d_model)
+
+        # Get iteration budget (if budget allocator available)
+        if self.budget_allocator is not None:
+            max_iters = self.budget_allocator.get_layer_budget(self.layer_idx, self.training)
+        else:
+            max_iters = self.num_iterations
+
+        # Use adaptive forward if available (inference mode with early stopping)
+        if self.use_adaptive_stopping and hasattr(self.inl, 'forward_adaptive') and not self.training:
+            # ✅ Adaptive early stopping (3× faster inference)
+            x_final_flat, v_final_flat, adaptive_result = self.inl.forward_adaptive(
+                ctx_flat,
+                x_flat_init,
+                v_flat_init,
+                num_iterations=max_iters,
+                use_early_stopping=True,
+                return_trajectory=True
+            )
+
+            # Get trajectories from adaptive result
+            if 'x_trajectory' in adaptive_result:
+                x_traj_flat = adaptive_result['x_trajectory']  # [B*S, T+1, D]
+                v_traj_flat = adaptive_result['v_trajectory']  # [B*S, T+1, D]
+            else:
+                # Fallback: single final state
+                x_traj_flat = x_final_flat.unsqueeze(1)
+                v_traj_flat = v_final_flat.unsqueeze(1)
+
+            aux_infos = {
+                'x': x_traj_flat,
+                'v': v_traj_flat,
+                'mu': adaptive_result.get('mu'),
+                'mu_global': adaptive_result.get('mu_global'),
+                'mu_offsets': adaptive_result.get('mu_offsets'),
+                'iterations_used': adaptive_result.get('iterations_used'),
+                'avg_iterations': adaptive_result.get('avg_iterations'),
+                'max_iterations': max_iters,
+                'layer_idx': self.layer_idx
+            }
+
+            output = x_final_flat.reshape(batch_size, seq_len, d_model)
+
+        else:
+            # Budget-aware training mode
+            x_trajectory = [x_flat_init.clone()]
+            v_trajectory = [v_flat_init.clone()]
+
+            x_flat, v_flat = x_flat_init, v_flat_init
+            x_prev = x_flat_init
+            actual_iterations = 0
+
+            for iteration in range(max_iters):
+                x_next_flat, v_next_flat, aux = self.inl(ctx_flat, x_flat, v_flat, step=iteration)
+
+                # Check for early stopping (if budget allocator with convergence checking)
+                if (self.budget_allocator is not None and
+                    iteration >= self.budget_allocator.warmup_iterations and
+                    not self.training):
+
+                    converged = self.budget_allocator.check_convergence(x_next_flat, x_flat, iteration)
+                    if converged:
+                        x_flat, v_flat = x_next_flat, v_next_flat
+                        actual_iterations = iteration + 1
+                        x_trajectory.append(x_flat.clone())
+                        v_trajectory.append(v_flat.clone())
+                        break
+
+                x_prev = x_flat
+                x_flat, v_flat = x_next_flat, v_next_flat
+                actual_iterations = iteration + 1
+
+                # Save trajectories for loss computation
+                x_trajectory.append(x_flat.clone())
+                v_trajectory.append(v_flat.clone())
+
+            # Update budget statistics (during training)
+            if self.training and self.budget_allocator is not None:
+                final_delta = torch.norm(x_flat - x_prev, dim=-1).mean().item()
+                self.budget_allocator.update_statistics(
+                    self.layer_idx,
+                    actual_iterations,
+                    final_delta
+                )
+
+            # Stack trajectories: [B*S, T+1, D]
+            x_traj_flat = torch.stack(x_trajectory, dim=1)
+            v_traj_flat = torch.stack(v_trajectory, dim=1)
+
+            aux_infos = {
+                'x': x_traj_flat,
+                'v': v_traj_flat,
+                'mu': aux.get('mu', None),
+                'mu_global': aux.get('mu_global', None),
+                'mu_offsets': aux.get('mu_offsets', None),
+                'iterations_used': actual_iterations,
+                'max_iterations': max_iters,
+                'layer_idx': self.layer_idx
+            }
+
+            output = x_flat.reshape(batch_size, seq_len, d_model)
+
+        # NOTE: No need to update integrator cache - we don't cache x, v states
+        # since integrator dynamics are computed fresh for each token.
+
+        # Residual
+        x = x + self.dropout(output)
+
+        # Feedforward
+        x = x + self.ff(self.norm2(x))
+
+        return x, aux_infos
+
+
+class UltraOptimizedIntegratorLanguageModel(nn.Module):
+    """
+    ULTRA-OPTIMIZED INL-LLM
+
+    All optimizations enabled by default:
+    ✅ Low-rank embeddings (87% reduction)
+    ✅ Gradient checkpointing (60% memory save)
+    ✅ Adaptive early stopping (40% faster)
+    ✅ Shared controllers (96% controller reduction)
+    ✅ Hierarchical equilibrium (98% μ reduction)
+    ✅ Sparse excitation (10x less compute)
+
+    Can scale to 100B+ parameters efficiently!
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        d_model: int = 512,
+        num_layers: int = 6,
+        num_heads: int = 8,
+        num_iterations_per_layer: int = 5,
+        feedforward_dim: int = None,
+        max_seq_len: int = 2048,
+        dropout: float = 0.1,
+        # Optimization flags
+        use_lowrank_embeddings: bool = True,
+        lowrank_ratio: float = 0.125,
+        use_gradient_checkpointing: bool = True,
+        use_shared_controllers: bool = True,
+        use_adaptive_stopping: bool = True,
+        adaptive_convergence_threshold: float = 0.001,
+        hierarchical_group_size: int = 64,
+        excitation_sparsity: float = 0.1,
+        # Adaptive budget allocation
+        use_adaptive_budget: bool = True,
+        budget_strategy: str = 'hybrid',
+        budget_convergence_threshold: float = 0.001,
+        # Mixture of Experts (MoE)
+        use_moe: bool = False,
+        num_experts: int = 4,
+        moe_top_k: int = 2,
+        moe_load_balance_weight: float = 0.01
+    ):
+        super().__init__()
+
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.num_layers = num_layers
+        self.use_adaptive_budget = use_adaptive_budget
+        self.use_moe = use_moe
+
+        if feedforward_dim is None:
+            feedforward_dim = 4 * d_model
+
+        # Low-rank embeddings
+        if use_lowrank_embeddings:
+            self.token_embedding = LowRankEmbedding(vocab_size, d_model, rank_ratio=lowrank_ratio)
+            print(f"✅ Low-Rank Embeddings: {self.token_embedding}")
+        else:
+            self.token_embedding = nn.Embedding(vocab_size, d_model)
+
+        # Positional encoding
+        self.pos_encoding = PositionalEncoding(d_model, max_seq_len)
+        self.dropout = nn.Dropout(dropout)
+
+        # Shared controller (ONE for all layers!)
+        if use_shared_controllers:
+            self.shared_controller = SharedController(
+                hidden_dim=d_model,
+                output_dim=d_model,
+                num_layers=num_layers,
+                hidden_controller=64
+            )
+            print(f"✅ Shared Controllers: {self.shared_controller.num_parameters():,} params for {num_layers} layers")
+        else:
+            self.shared_controller = None
+
+        # Adaptive budget allocator (LEVEL 3!)
+        if use_adaptive_budget:
+            self.budget_allocator = create_budget_allocator(
+                num_layers=num_layers,
+                avg_iterations_per_layer=num_iterations_per_layer,
+                strategy=budget_strategy,
+                convergence_threshold=budget_convergence_threshold,
+                min_iterations_per_layer=max(2, num_iterations_per_layer // 2),
+                max_iterations_per_layer=num_iterations_per_layer * 2
+            )
+            print(f"✅ Adaptive Budget: {self.budget_allocator.total_budget} total iterations, strategy='{budget_strategy}'")
+        else:
+            self.budget_allocator = None
+
+        # Mixture of Experts Controller (LEVEL 4!)
+        if use_moe:
+            self.moe_controller = create_moe_controller(
+                d_model=d_model,
+                num_layers=num_layers,
+                num_experts=num_experts,
+                top_k=moe_top_k,
+                load_balance_weight=moe_load_balance_weight
+            )
+            print(f"✅ MoE Controller: {num_experts} experts, top-{moe_top_k} routing")
+        else:
+            self.moe_controller = None
+
+        # Layers
+        self.layers = nn.ModuleList([
+            UltraOptimizedINLBlock(
+                d_model=d_model,
+                num_heads=num_heads,
+                num_iterations=num_iterations_per_layer,
+                shared_controller=self.shared_controller,
+                layer_idx=i,
+                feedforward_dim=feedforward_dim,
+                dropout=dropout,
+                use_gradient_checkpointing=use_gradient_checkpointing,
+                use_adaptive_stopping=use_adaptive_stopping,
+                adaptive_convergence_threshold=adaptive_convergence_threshold,
+                group_size=hierarchical_group_size,
+                excitation_sparsity=excitation_sparsity,
+                budget_allocator=self.budget_allocator,
+                moe_controller=self.moe_controller
+            )
+            for i in range(num_layers)
+        ])
+
+        # Final norm
+        self.final_norm = nn.LayerNorm(d_model)
+
+        # LM head
+        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
+
+        # Initialize
+        self._init_weights()
+        self._print_optimization_status()
+
+    def _init_weights(self):
+        """Initialize weights."""
+        if not isinstance(self.token_embedding, LowRankEmbedding):
+            with torch.no_grad():
+                nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
+
+        with torch.no_grad():
+            nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.02)
+
+    def _print_optimization_status(self):
+        """Print optimization summary."""
+        print("\n" + "=" * 70)
+        print("ULTRA-OPTIMIZED INL-LLM")
+        print("=" * 70)
+        print("LEVEL 1 (Basic Optimizations):")
+        print(f"  ✅ Low-rank embeddings")
+        print(f"  ✅ Gradient checkpointing")
+        print(f"  ✅ Adaptive early stopping")
+        print("\nLEVEL 2 (Advanced Optimizations):")
+        print(f"  ✅ Shared controllers (across {self.num_layers} layers)")
+        print(f"  ✅ Hierarchical equilibrium")
+        print(f"  ✅ Sparse harmonic excitation")
+        if self.use_adaptive_budget:
+            print("\nLEVEL 3 (Bio-inspired Compute Allocation):")
+            print(f"  ✅ Adaptive budget allocation (strategy: {self.budget_allocator.strategy})")
+            budgets = self.budget_allocator.get_all_budgets(training=False)
+            print(f"  ✅ Dynamic iterations per layer: {min(budgets)}-{max(budgets)} (avg: {sum(budgets)/len(budgets):.1f})")
+            print(f"  ✅ Total compute budget: {self.budget_allocator.total_budget} iterations")
+        if self.use_moe:
+            print("\nLEVEL 4 (Mixture of Experts):")
+            print(f"  ✅ MoE Controller: {self.moe_controller.num_experts} specialized experts")
+            print(f"  ✅ Sparse routing: top-{self.moe_controller.top_k} experts per forward")
+            print(f"  ✅ Load balancing: {self.moe_controller.load_balance_weight} weight")
+            print(f"  ✅ Capacity increase: ~{self.moe_controller.num_experts / self.moe_controller.top_k:.1f}x with same compute")
+        print(f"\nTotal parameters: {self.get_num_params():,}")
+        print("=" * 70 + "\n")
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[INLCache] = None,
+        use_cache: bool = False,
+        return_aux: bool = False
+    ) -> Tuple[torch.Tensor, Optional[List], Optional[INLCache]]:
+        """
+        Forward pass with optional KV caching.
+
+        Args:
+            input_ids: Input token IDs [batch_size, seq_len]
+            attention_mask: Attention mask (optional)
+            past_key_values: Previous cache (INLCache object)
+            use_cache: Whether to use/update cache
+            return_aux: Whether to return auxiliary info
+
+        Returns:
+            logits: Output logits [batch_size, seq_len, vocab_size]
+            all_aux: Auxiliary info from each layer (if return_aux=True)
+            new_cache: Updated cache (if use_cache=True)
+        """
+        # Initialize cache if needed
+        if use_cache and past_key_values is None:
+            past_key_values = INLCache(num_layers=self.num_layers)
+
+        # Determine starting position for positional encoding
+        start_pos = 0
+        if use_cache and past_key_values is not None:
+            start_pos = past_key_values.get_seq_length()
+
+        # Embedding with correct positional encoding
+        x = self.token_embedding(input_ids)
+        x = self.pos_encoding(x, start_pos=start_pos)
+        x = self.dropout(x)
+
+        # Layers
+        all_aux = [] if return_aux else None
+
+        for layer_idx, layer in enumerate(self.layers):
+            cache_layer = past_key_values[layer_idx] if use_cache else None
+            x, aux = layer(x, mask=attention_mask, cache_layer=cache_layer, use_cache=use_cache)
+            if return_aux:
+                all_aux.append(aux)
+
+        # Final norm
+        x = self.final_norm(x)
+
+        # LM head
+        logits = self.lm_head(x)
+
+        return logits, all_aux, past_key_values if use_cache else None
+
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        max_new_tokens: int = 100,
+        temperature: float = 1.0,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        do_sample: bool = True,
+        use_cache: bool = True
+    ) -> torch.Tensor:
+        """
+        Autoregressive generation with optional KV caching.
+
+        Args:
+            input_ids: Input token IDs [batch_size, seq_len]
+            max_new_tokens: Number of tokens to generate
+            temperature: Sampling temperature
+            top_k: Top-k sampling (if provided)
+            top_p: Nucleus sampling threshold (if provided)
+            do_sample: Whether to sample or use greedy decoding
+            use_cache: Whether to use KV caching (default: True, much faster!)
+
+        Returns:
+            Generated token IDs [batch_size, seq_len + max_new_tokens]
+        """
+        self.eval()
+        past_key_values = None
+
+        with torch.no_grad():
+            for step in range(max_new_tokens):
+                # Use cache for all steps after the first
+                if use_cache and step > 0:
+                    # Only pass the last token for cached generation
+                    model_input = input_ids[:, -1:]
+                    logits, _, past_key_values = self.forward(
+                        model_input,
+                        past_key_values=past_key_values,
+                        use_cache=True
+                    )
+                else:
+                    # First step or no cache: process full sequence
+                    logits, _, past_key_values = self.forward(
+                        input_ids,
+                        past_key_values=past_key_values if use_cache else None,
+                        use_cache=use_cache
+                    )
+
+                # Get logits for last token
+                logits = logits[:, -1, :] / temperature
+
+                # Apply top-k filtering
+                if top_k is not None:
+                    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+                    logits[indices_to_remove] = float('-inf')
+
+                # Apply top-p (nucleus) filtering
+                if top_p is not None:
+                    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                    sorted_indices_to_remove = cumulative_probs > top_p
+                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                    sorted_indices_to_remove[..., 0] = 0
+                    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+                    logits[indices_to_remove] = float('-inf')
+
+                # Sample or select greedily
+                if do_sample:
+                    probs = F.softmax(logits, dim=-1)
+                    next_token = torch.multinomial(probs, num_samples=1)
+                else:
+                    next_token = torch.argmax(logits, dim=-1, keepdim=True)
+
+                # Append to sequence
+                input_ids = torch.cat([input_ids, next_token], dim=1)
+
+        return input_ids
+
+    def get_num_params(self) -> int:
+        """Count parameters."""
+        return sum(p.numel() for p in self.parameters())
+
+    def get_inference_stats(self) -> Dict:
+        """
+        Get model statistics and optimization info.
+
+        Returns dict with model configuration and enabled optimizations.
+        """
+        stats = {
+            'num_params': self.get_num_params(),
+            'num_layers': self.num_layers,
+            'd_model': self.d_model,
+            'optimizations_enabled': {
+                'low_rank_embeddings': True,
+                'shared_controllers': True,
+                'hierarchical_equilibrium': True,
+                'sparse_excitation': True,
+                'gradient_checkpointing': True,
+                'adaptive_budget': self.use_adaptive_budget
+            }
+        }
+
+        # Add budget statistics if available
+        if self.use_adaptive_budget and self.budget_allocator is not None:
+            budget_stats = self.budget_allocator.get_statistics()
+            stats['budget_allocation'] = {
+                'layer_budgets': budget_stats['layer_budgets'].tolist(),
+                'total_budget': budget_stats['total_budget'].item(),
+                'avg_iterations_history': budget_stats['layer_iterations_history'].tolist(),
+                'convergence_speeds': budget_stats['layer_convergence_speed'].tolist()
+            }
+
+        return stats
+
+
+def create_ultra_optimized_model(
+    size: str = 'small',
+    vocab_size: int = 50000
+) -> UltraOptimizedIntegratorLanguageModel:
+    """
+    Create ultra-optimized model.
+
+    Sizes: 'small', 'medium', 'large', 'xlarge', '3B', '7B', '13B', '30B', '70B'
+    """
+    configs = {
+        'small': {'d_model': 512, 'num_layers': 6, 'num_heads': 8, 'iterations': 5, 'ff_dim': 2048},
+        'medium': {'d_model': 768, 'num_layers': 12, 'num_heads': 12, 'iterations': 7, 'ff_dim': 3072},
+        'large': {'d_model': 1024, 'num_layers': 24, 'num_heads': 16, 'iterations': 10, 'ff_dim': 4096},
+        'xlarge': {'d_model': 1536, 'num_layers': 32, 'num_heads': 24, 'iterations': 12, 'ff_dim': 6144},
+        '3B': {'d_model': 2048, 'num_layers': 40, 'num_heads': 32, 'iterations': 15, 'ff_dim': 8192},
+        '7B': {'d_model': 4096, 'num_layers': 32, 'num_heads': 32, 'iterations': 10, 'ff_dim': 16384},
+        '13B': {'d_model': 5120, 'num_layers': 40, 'num_heads': 40, 'iterations': 12, 'ff_dim': 20480},
+        '30B': {'d_model': 6656, 'num_layers': 60, 'num_heads': 52, 'iterations': 12, 'ff_dim': 26624},
+        '70B': {'d_model': 8192, 'num_layers': 80, 'num_heads': 64, 'iterations': 12, 'ff_dim': 32768},
+    }
+
+    if size not in configs:
+        raise ValueError(f"Size must be one of {list(configs.keys())}")
+
+    cfg = configs[size]
+
+    model = UltraOptimizedIntegratorLanguageModel(
+        vocab_size=vocab_size,
+        d_model=cfg['d_model'],
+        num_layers=cfg['num_layers'],
+        num_heads=cfg['num_heads'],
+        num_iterations_per_layer=cfg['iterations'],
+        feedforward_dim=cfg['ff_dim'],
+        max_seq_len=2048,
+        # All optimizations enabled
+        use_lowrank_embeddings=True,
+        lowrank_ratio=0.125,
+        use_gradient_checkpointing=True,
+        use_shared_controllers=True,
+        hierarchical_group_size=64,
+        excitation_sparsity=0.1
+    )
+
+    print(f"\n🚀 ULTRA-OPTIMIZED INL-LLM ({size}): {model.get_num_params():,} parameters")
+    print(f"   Ready to scale to 100B+ with maximum efficiency!\n")
+
+    return model
+
+
+if __name__ == '__main__':
+    # Fix imports for standalone execution
+    import sys
+    import os
+    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+
+    from inl_llm import create_model
+
+    print("\n" + "=" * 70)
+    print("INL-LLM MODEL - Test")
+    print("=" * 70 + "\n")
+
+    # Create model
+    model = create_model(size='medium', vocab_size=50000)
+
+    # Test forward
+    batch_size = 2
+    seq_len = 10
+    input_ids = torch.randint(0, 50000, (batch_size, seq_len))
+
+    print("Running forward pass...")
+    logits, aux = model(input_ids, return_aux=True)
+
+    print(f"✅ Input shape: {input_ids.shape}")
+    print(f"✅ Output shape: {logits.shape}")
+    print(f"✅ Aux layers: {len(aux)}")
+
+    # Test generation
+    print("\nTesting generation...")
+    prompt = torch.randint(0, 50000, (1, 5))
+    generated = model.generate(prompt, max_new_tokens=20, temperature=0.8)
+
+    print(f"✅ Prompt length: {prompt.shape[1]}")
+    print(f"✅ Generated length: {generated.shape[1]}")
+
+    print("\n" + "=" * 70)
+    print("✅ INL-LLM WORKING PERFECTLY!")
+    print("=" * 70 + "\n")

inl_llm/models/modeling_inl_llm.py

@@ -1,226 +1,226 @@
-"""
-HuggingFace-compatible wrapper for INL-LLM to enable vLLM support.
-
-This module registers the UltraOptimizedIntegratorLanguageModel with HuggingFace's
-AutoModel system, making it compatible with vLLM and other HF-based serving frameworks.
-
-Author: Boris Peyriguère
-"""
-
-import torch
-import torch.nn as nn
-from typing import Optional, Tuple, Union
-from transformers import PreTrainedModel, PretrainedConfig
-from transformers.modeling_outputs import CausalLMOutputWithPast
-
-from .integrator_language_model import UltraOptimizedIntegratorLanguageModel
-
-
-class INLLLMConfig(PretrainedConfig):
-    """
-    Configuration class for INL-LLM models.
-
-    This is required for HuggingFace AutoModel integration and vLLM compatibility.
-    """
-
-    model_type = "inl-llm"
-
-    def __init__(
-        self,
-        vocab_size: int = 50261,
-        d_model: int = 1728,
-        num_layers: int = 25,
-        num_heads: int = 32,
-        num_iterations_per_layer: int = 5,
-        feedforward_dim: int = 6912,
-        max_seq_len: int = 2048,
-        dropout: float = 0.1,
-        # Optimization settings
-        use_lowrank_embeddings: bool = True,
-        lowrank_ratio: float = 0.125,
-        use_gradient_checkpointing: bool = True,
-        use_shared_controllers: bool = True,
-        use_adaptive_stopping: bool = True,
-        adaptive_convergence_threshold: float = 0.001,
-        hierarchical_group_size: int = 64,
-        excitation_sparsity: float = 0.1,
-        # Token IDs
-        bos_token_id: int = 50256,
-        eos_token_id: int = 50256,
-        pad_token_id: int = 50256,
-        # Integration metadata
-        integrator_type: str = "ultra_optimized",
-        controller_type: str = "shared",
-        equilibrium_type: str = "hierarchical",
-        excitation_type: str = "sparse_harmonic",
-        **kwargs
-    ):
-        super().__init__(
-            bos_token_id=bos_token_id,
-            eos_token_id=eos_token_id,
-            pad_token_id=pad_token_id,
-            **kwargs
-        )
-
-        self.vocab_size = vocab_size
-        self.d_model = d_model
-        self.num_layers = num_layers
-        self.num_heads = num_heads
-        self.num_iterations_per_layer = num_iterations_per_layer
-        self.feedforward_dim = feedforward_dim
-        self.max_seq_len = max_seq_len
-        self.dropout = dropout
-
-        # Optimizations
-        self.use_lowrank_embeddings = use_lowrank_embeddings
-        self.lowrank_ratio = lowrank_ratio
-        self.use_gradient_checkpointing = use_gradient_checkpointing
-        self.use_shared_controllers = use_shared_controllers
-        self.use_adaptive_stopping = use_adaptive_stopping
-        self.adaptive_convergence_threshold = adaptive_convergence_threshold
-        self.hierarchical_group_size = hierarchical_group_size
-        self.excitation_sparsity = excitation_sparsity
-
-        # Metadata
-        self.integrator_type = integrator_type
-        self.controller_type = controller_type
-        self.equilibrium_type = equilibrium_type
-        self.excitation_type = excitation_type
-
-
-class INLLLMForCausalLM(PreTrainedModel):
-    """
-    HuggingFace-compatible wrapper for UltraOptimizedIntegratorLanguageModel.
-
-    This wrapper enables:
-    - vLLM support
-    - HuggingFace AutoModel.from_pretrained()
-    - Compatibility with HF ecosystem (pipelines, etc.)
-    """
-
-    config_class = INLLLMConfig
-    base_model_prefix = "inl_llm"
-    supports_gradient_checkpointing = True
-    _no_split_modules = ["UltraOptimizedINLBlock"]
-
-    def __init__(self, config: INLLLMConfig):
-        super().__init__(config)
-
-        # Create the underlying INL-LLM model
-        self.model = UltraOptimizedIntegratorLanguageModel(
-            vocab_size=config.vocab_size,
-            d_model=config.d_model,
-            num_layers=config.num_layers,
-            num_heads=config.num_heads,
-            num_iterations_per_layer=config.num_iterations_per_layer,
-            feedforward_dim=config.feedforward_dim,
-            max_seq_len=config.max_seq_len,
-            dropout=config.dropout,
-            use_lowrank_embeddings=config.use_lowrank_embeddings,
-            lowrank_ratio=config.lowrank_ratio,
-            use_gradient_checkpointing=config.use_gradient_checkpointing,
-            use_shared_controllers=config.use_shared_controllers,
-            use_adaptive_stopping=config.use_adaptive_stopping,
-            adaptive_convergence_threshold=config.adaptive_convergence_threshold,
-            hierarchical_group_size=config.hierarchical_group_size,
-            excitation_sparsity=config.excitation_sparsity
-        )
-
-        # Language model head (already part of UltraOptimizedIntegratorLanguageModel)
-        # No need to add another lm_head
-
-        # Initialize weights
-        self.post_init()
-
-    def get_input_embeddings(self):
-        """Required for HuggingFace compatibility."""
-        return self.model.token_embedding
-
-    def set_input_embeddings(self, value):
-        """Required for HuggingFace compatibility."""
-        self.model.token_embedding = value
-
-    def get_output_embeddings(self):
-        """Required for HuggingFace compatibility."""
-        return self.model.lm_head
-
-    def set_output_embeddings(self, new_embeddings):
-        """Required for HuggingFace compatibility."""
-        self.model.lm_head = new_embeddings
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        return_dict: Optional[bool] = None,
-        **kwargs
-    ) -> Union[Tuple, CausalLMOutputWithPast]:
-        """
-        Forward pass compatible with HuggingFace's CausalLM interface.
-
-        Args:
-            input_ids: Input token IDs [batch_size, seq_len]
-            attention_mask: Attention mask (currently not used by INL-LLM)
-            labels: Labels for language modeling loss
-            return_dict: Whether to return a ModelOutput object
-
-        Returns:
-            CausalLMOutputWithPast or tuple of (loss, logits)
-        """
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        # Forward through INL-LLM
-        logits = self.model(input_ids)
-
-        # Compute loss if labels provided
-        loss = None
-        if labels is not None:
-            # Shift so that tokens < n predict n
-            shift_logits = logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous()
-
-            # Flatten the tokens
-            loss_fct = nn.CrossEntropyLoss()
-            loss = loss_fct(
-                shift_logits.view(-1, shift_logits.size(-1)),
-                shift_labels.view(-1)
-            )
-
-        if not return_dict:
-            output = (logits,)
-            return ((loss,) + output) if loss is not None else output
-
-        return CausalLMOutputWithPast(
-            loss=loss,
-            logits=logits,
-            past_key_values=None,  # INL-LLM doesn't use KV cache
-            hidden_states=None,
-            attentions=None
-        )
-
-    def prepare_inputs_for_generation(
-        self,
-        input_ids: torch.LongTensor,
-        **kwargs
-    ):
-        """Prepare inputs for generation (required for .generate())."""
-        return {
-            "input_ids": input_ids,
-        }
-
-    @staticmethod
-    def _reorder_cache(past, beam_idx):
-        """Required for beam search (INL-LLM doesn't use cache)."""
-        return past
-
-    def get_num_params(self) -> int:
-        """Get total number of parameters."""
-        return self.model.get_num_params()
-
-
-# Register the model with HuggingFace AutoModel
-from transformers import AutoConfig, AutoModelForCausalLM
-
-AutoConfig.register("inl-llm", INLLLMConfig)
-AutoModelForCausalLM.register(INLLLMConfig, INLLLMForCausalLM)
+"""
+HuggingFace-compatible wrapper for INL-LLM to enable vLLM support.
+
+This module registers the UltraOptimizedIntegratorLanguageModel with HuggingFace's
+AutoModel system, making it compatible with vLLM and other HF-based serving frameworks.
+
+Author: Boris Peyriguère
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional, Tuple, Union
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .integrator_language_model import UltraOptimizedIntegratorLanguageModel
+
+
+class INLLLMConfig(PretrainedConfig):
+    """
+    Configuration class for INL-LLM models.
+
+    This is required for HuggingFace AutoModel integration and vLLM compatibility.
+    """
+
+    model_type = "inl-llm"
+
+    def __init__(
+        self,
+        vocab_size: int = 50261,
+        d_model: int = 1728,
+        num_layers: int = 25,
+        num_heads: int = 32,
+        num_iterations_per_layer: int = 5,
+        feedforward_dim: int = 6912,
+        max_seq_len: int = 2048,
+        dropout: float = 0.1,
+        # Optimization settings
+        use_lowrank_embeddings: bool = True,
+        lowrank_ratio: float = 0.125,
+        use_gradient_checkpointing: bool = True,
+        use_shared_controllers: bool = True,
+        use_adaptive_stopping: bool = True,
+        adaptive_convergence_threshold: float = 0.001,
+        hierarchical_group_size: int = 64,
+        excitation_sparsity: float = 0.1,
+        # Token IDs
+        bos_token_id: int = 50256,
+        eos_token_id: int = 50256,
+        pad_token_id: int = 50256,
+        # Integration metadata
+        integrator_type: str = "ultra_optimized",
+        controller_type: str = "shared",
+        equilibrium_type: str = "hierarchical",
+        excitation_type: str = "sparse_harmonic",
+        **kwargs
+    ):
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            **kwargs
+        )
+
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.num_iterations_per_layer = num_iterations_per_layer
+        self.feedforward_dim = feedforward_dim
+        self.max_seq_len = max_seq_len
+        self.dropout = dropout
+
+        # Optimizations
+        self.use_lowrank_embeddings = use_lowrank_embeddings
+        self.lowrank_ratio = lowrank_ratio
+        self.use_gradient_checkpointing = use_gradient_checkpointing
+        self.use_shared_controllers = use_shared_controllers
+        self.use_adaptive_stopping = use_adaptive_stopping
+        self.adaptive_convergence_threshold = adaptive_convergence_threshold
+        self.hierarchical_group_size = hierarchical_group_size
+        self.excitation_sparsity = excitation_sparsity
+
+        # Metadata
+        self.integrator_type = integrator_type
+        self.controller_type = controller_type
+        self.equilibrium_type = equilibrium_type
+        self.excitation_type = excitation_type
+
+
+class INLLLMForCausalLM(PreTrainedModel):
+    """
+    HuggingFace-compatible wrapper for UltraOptimizedIntegratorLanguageModel.
+
+    This wrapper enables:
+    - vLLM support
+    - HuggingFace AutoModel.from_pretrained()
+    - Compatibility with HF ecosystem (pipelines, etc.)
+    """
+
+    config_class = INLLLMConfig
+    base_model_prefix = "inl_llm"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["UltraOptimizedINLBlock"]
+
+    def __init__(self, config: INLLLMConfig):
+        super().__init__(config)
+
+        # Create the underlying INL-LLM model
+        self.model = UltraOptimizedIntegratorLanguageModel(
+            vocab_size=config.vocab_size,
+            d_model=config.d_model,
+            num_layers=config.num_layers,
+            num_heads=config.num_heads,
+            num_iterations_per_layer=config.num_iterations_per_layer,
+            feedforward_dim=config.feedforward_dim,
+            max_seq_len=config.max_seq_len,
+            dropout=config.dropout,
+            use_lowrank_embeddings=config.use_lowrank_embeddings,
+            lowrank_ratio=config.lowrank_ratio,
+            use_gradient_checkpointing=config.use_gradient_checkpointing,
+            use_shared_controllers=config.use_shared_controllers,
+            use_adaptive_stopping=config.use_adaptive_stopping,
+            adaptive_convergence_threshold=config.adaptive_convergence_threshold,
+            hierarchical_group_size=config.hierarchical_group_size,
+            excitation_sparsity=config.excitation_sparsity
+        )
+
+        # Language model head (already part of UltraOptimizedIntegratorLanguageModel)
+        # No need to add another lm_head
+
+        # Initialize weights
+        self.post_init()
+
+    def get_input_embeddings(self):
+        """Required for HuggingFace compatibility."""
+        return self.model.token_embedding
+
+    def set_input_embeddings(self, value):
+        """Required for HuggingFace compatibility."""
+        self.model.token_embedding = value
+
+    def get_output_embeddings(self):
+        """Required for HuggingFace compatibility."""
+        return self.model.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        """Required for HuggingFace compatibility."""
+        self.model.lm_head = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        """
+        Forward pass compatible with HuggingFace's CausalLM interface.
+
+        Args:
+            input_ids: Input token IDs [batch_size, seq_len]
+            attention_mask: Attention mask (currently not used by INL-LLM)
+            labels: Labels for language modeling loss
+            return_dict: Whether to return a ModelOutput object
+
+        Returns:
+            CausalLMOutputWithPast or tuple of (loss, logits)
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Forward through INL-LLM
+        logits = self.model(input_ids)
+
+        # Compute loss if labels provided
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            # Flatten the tokens
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1)
+            )
+
+        if not return_dict:
+            output = (logits,)
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,  # INL-LLM doesn't use KV cache
+            hidden_states=None,
+            attentions=None
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        **kwargs
+    ):
+        """Prepare inputs for generation (required for .generate())."""
+        return {
+            "input_ids": input_ids,
+        }
+
+    @staticmethod
+    def _reorder_cache(past, beam_idx):
+        """Required for beam search (INL-LLM doesn't use cache)."""
+        return past
+
+    def get_num_params(self) -> int:
+        """Get total number of parameters."""
+        return self.model.get_num_params()
+
+
+# Register the model with HuggingFace AutoModel
+from transformers import AutoConfig, AutoModelForCausalLM
+
+AutoConfig.register("inl-llm", INLLLMConfig)
+AutoModelForCausalLM.register(INLLLMConfig, INLLLMForCausalLM)

inl_llm/optimizations/__init__.py

@@ -1,49 +1,49 @@
-"""
-Optimization modules for INL-LLM.
-
-Level 1 (Production-ready):
-- LowRankEmbedding: Reduces embedding parameters by 70-80%
-- GradientCheckpointedINL: Reduces training memory by 50-70%
-- AdaptiveIntegratorNeuronLayer: Speeds up inference by 30-50%
-
-Level 2 (Research/Experimental):
-- SharedController: Shares controllers across layers (-96% params)
-- SparseHarmonicINL: Sparse excitation (10x less compute)
-- HierarchicalEquilibriumINL: Hierarchical equilibrium learning (-98% params)
-- MixtureOfIntegrators: Conditional computation (MoE-style)
-"""
-
-# Level 1 optimizations
-from .optimizations import (
-    LowRankEmbedding,
-    AdaptiveIntegratorNeuronLayer,
-    AdaptiveHierarchicalINL,
-    GradientCheckpointedINL,
-    compute_parameter_reduction,
-    print_optimization_summary
-)
-
-# Level 2 optimizations
-from .advanced_optimizations import (
-    SharedController,
-    SparseHarmonicINL,
-    HierarchicalEquilibriumINL,
-    MixtureOfIntegrators,
-    compute_advanced_optimization_gains
-)
-
-__all__ = [
-    # Level 1
-    'LowRankEmbedding',
-    'AdaptiveIntegratorNeuronLayer',
-    'AdaptiveHierarchicalINL',
-    'GradientCheckpointedINL',
-    'compute_parameter_reduction',
-    'print_optimization_summary',
-    # Level 2
-    'SharedController',
-    'SparseHarmonicINL',
-    'HierarchicalEquilibriumINL',
-    'MixtureOfIntegrators',
-    'compute_advanced_optimization_gains'
-]
+"""
+Optimization modules for INL-LLM.
+
+Level 1 (Production-ready):
+- LowRankEmbedding: Reduces embedding parameters by 70-80%
+- GradientCheckpointedINL: Reduces training memory by 50-70%
+- AdaptiveIntegratorNeuronLayer: Speeds up inference by 30-50%
+
+Level 2 (Research/Experimental):
+- SharedController: Shares controllers across layers (-96% params)
+- SparseHarmonicINL: Sparse excitation (10x less compute)
+- HierarchicalEquilibriumINL: Hierarchical equilibrium learning (-98% params)
+- MixtureOfIntegrators: Conditional computation (MoE-style)
+"""
+
+# Level 1 optimizations
+from .optimizations import (
+    LowRankEmbedding,
+    AdaptiveIntegratorNeuronLayer,
+    AdaptiveHierarchicalINL,
+    GradientCheckpointedINL,
+    compute_parameter_reduction,
+    print_optimization_summary
+)
+
+# Level 2 optimizations
+from .advanced_optimizations import (
+    SharedController,
+    SparseHarmonicINL,
+    HierarchicalEquilibriumINL,
+    MixtureOfIntegrators,
+    compute_advanced_optimization_gains
+)
+
+__all__ = [
+    # Level 1
+    'LowRankEmbedding',
+    'AdaptiveIntegratorNeuronLayer',
+    'AdaptiveHierarchicalINL',
+    'GradientCheckpointedINL',
+    'compute_parameter_reduction',
+    'print_optimization_summary',
+    # Level 2
+    'SharedController',
+    'SparseHarmonicINL',
+    'HierarchicalEquilibriumINL',
+    'MixtureOfIntegrators',
+    'compute_advanced_optimization_gains'
+]

inl_llm/optimizations/advanced_optimizations.py

@@ -1,619 +1,619 @@
-"""
-Advanced Optimizations for INL-LLM
-
-Implements additional efficiency techniques:
-1. Shared Controllers: Share control MLPs across layers (-15-20% params)
-2. Sparse Harmonic Excitation: Only excite subset of dimensions (-10x compute)
-3. Mixture of Integrators (MoI): Conditional computation like MoE
-4. Hierarchical Equilibrium: Global + local offsets for μ
-
-Author: Boris Peyriguère
-"""
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from typing import Optional, Tuple, Dict, List
-import math
-
-
-class SharedController(nn.Module):
-    """
-    Shared controller MLP across multiple INL layers.
-
-    Instead of each layer having its own controller (α, β, g, v_cand),
-    we use ONE shared controller + small layer-specific modulation.
-
-    Benefit: 15-20% parameter reduction on controller networks
-    """
-
-    def __init__(
-        self,
-        hidden_dim: int,
-        output_dim: int,
-        num_layers: int,
-        hidden_controller: int = 64
-    ):
-        """
-        Args:
-            hidden_dim: Context dimension
-            output_dim: State dimension
-            num_layers: Number of layers sharing this controller
-            hidden_controller: Hidden size for controller MLP
-        """
-        super().__init__()
-
-        self.hidden_dim = hidden_dim
-        self.output_dim = output_dim
-        self.num_layers = num_layers
-
-        # Single shared controller (used by all layers)
-        self.controller_h = nn.Linear(hidden_dim, hidden_controller)
-        self.controller_x = nn.Linear(output_dim, hidden_controller)
-        self.controller_v = nn.Linear(output_dim, hidden_controller)
-        self.controller_mlp = nn.Sequential(
-            nn.ReLU(),
-            nn.Linear(hidden_controller, 4 * output_dim)
-        )
-
-        # Layer-specific modulation (tiny parameters)
-        # Each layer gets 4 scalar multipliers (α, β, g, v_cand)
-        self.layer_scalers = nn.Parameter(torch.ones(num_layers, 4))
-        self.layer_biases = nn.Parameter(torch.zeros(num_layers, 4))
-
-        # Initialize
-        self._init_weights()
-
-    def _init_weights(self):
-        """Initialize controller weights."""
-        with torch.no_grad():
-            nn.init.xavier_uniform_(self.controller_h.weight)
-            nn.init.xavier_uniform_(self.controller_x.weight)
-            nn.init.xavier_uniform_(self.controller_v.weight)
-            self.controller_h.bias.zero_()
-            self.controller_x.bias.zero_()
-            self.controller_v.bias.zero_()
-            self.controller_mlp[-1].weight.normal_(0.0, 0.01)
-
-    def forward(
-        self,
-        h: torch.Tensor,
-        x: torch.Tensor,
-        v: torch.Tensor,
-        layer_idx: int
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Compute controller parameters for specific layer.
-
-        Args:
-            h: Context [batch, hidden_dim]
-            x: State [batch, output_dim]
-            v: Velocity [batch, output_dim]
-            layer_idx: Which layer is requesting control
-
-        Returns:
-            alpha, beta, gate, v_cand (all [batch, output_dim])
-        """
-        # Shared computation
-        controller_hidden = self.controller_h(h) + self.controller_x(x) + self.controller_v(v)
-        controller_output = self.controller_mlp(controller_hidden)
-
-        # Split into components
-        alpha_base, beta_base, gate_base, v_cand_base = torch.split(
-            controller_output, self.output_dim, dim=1
-        )
-
-        # Layer-specific modulation
-        scaler = self.layer_scalers[layer_idx]  # [4]
-        bias = self.layer_biases[layer_idx]      # [4]
-
-        alpha = torch.sigmoid(alpha_base * scaler[0] + bias[0])
-        beta = F.softplus(beta_base * scaler[1] + bias[1])
-        gate = torch.sigmoid(gate_base * scaler[2] + bias[2])
-        v_cand = v_cand_base * scaler[3] + bias[3]
-
-        return alpha, beta, gate, v_cand
-
-    def num_parameters(self) -> int:
-        """Count parameters."""
-        shared = sum(p.numel() for p in [
-            self.controller_h.weight, self.controller_h.bias,
-            self.controller_x.weight, self.controller_x.bias,
-            self.controller_v.weight, self.controller_v.bias
-        ]) + sum(p.numel() for p in self.controller_mlp.parameters())
-
-        layer_specific = self.layer_scalers.numel() + self.layer_biases.numel()
-
-        return shared + layer_specific
-
-
-class SparseHarmonicINL(nn.Module):
-    """
-    INL with Sparse Harmonic Excitation.
-
-    Only applies harmonic noise to a subset of dimensions (e.g., 10%).
-    Reduces compute by 10x while maintaining exploration.
-    """
-
-    def __init__(
-        self,
-        hidden_dim: int,
-        output_dim: int,
-        sparsity: float = 0.1,
-        target_value: float = 5.0,
-        dt: float = 0.1,
-        excitation_amplitude: float = 0.03
-    ):
-        """
-        Args:
-            hidden_dim: Context dimension
-            output_dim: State dimension
-            sparsity: Fraction of dimensions to excite (0.1 = 10%)
-            target_value: Initial equilibrium
-            dt: Time step
-            excitation_amplitude: Amplitude of excitation
-        """
-        super().__init__()
-
-        self.hidden_dim = hidden_dim
-        self.output_dim = output_dim
-        self.sparsity = sparsity
-        self.dt = dt
-
-        # Learnable μ
-        self.mu = nn.Parameter(torch.full((output_dim,), target_value))
-
-        # Excitation parameters (only for sparse subset)
-        self.num_excited = max(1, int(output_dim * sparsity))
-
-        # Fixed sparse indices (deterministic)
-        indices = torch.linspace(0, output_dim - 1, self.num_excited).long()
-        self.register_buffer('excited_indices', indices)
-
-        # Learnable excitation params (only for excited dims)
-        self.register_buffer('excitation_amplitude', torch.tensor(excitation_amplitude))
-        self.excitation_gamma = nn.Parameter(torch.ones(self.num_excited))
-        self.excitation_phi = nn.Parameter(torch.zeros(self.num_excited))
-
-        # Simple controller (for demo - would use shared in practice)
-        self.controller = nn.Sequential(
-            nn.Linear(hidden_dim + 2 * output_dim, 64),
-            nn.ReLU(),
-            nn.Linear(64, 3 * output_dim)  # α, β, g
-        )
-
-    def forward(
-        self,
-        h: torch.Tensor,
-        x: torch.Tensor,
-        v: torch.Tensor,
-        step: int = 0
-    ) -> Tuple[torch.Tensor, torch.Tensor, Dict]:
-        """Forward with sparse excitation."""
-        batch_size = x.shape[0]
-
-        # Compute controllers
-        ctx = torch.cat([h, x, v], dim=-1)
-        controller_out = self.controller(ctx)
-        alpha_raw, beta_raw, gate_raw = torch.split(controller_out, self.output_dim, dim=1)
-
-        alpha = torch.sigmoid(alpha_raw)
-        beta = F.softplus(beta_raw)
-        gate = torch.sigmoid(gate_raw)
-
-        # Velocity update
-        error = x - self.mu
-        v_next = alpha * v - beta * error
-
-        # Sparse harmonic excitation (only on subset of dims)
-        if self.excitation_amplitude.item() > 0 and self.training:
-            t = float(step)
-            # Compute noise only for excited dimensions
-            noise_sparse = self.excitation_amplitude * torch.sin(
-                self.excitation_gamma * t + self.excitation_phi
-            )  # [num_excited]
-
-            # Apply to specific indices (sparse operation)
-            v_next[:, self.excited_indices] += noise_sparse.unsqueeze(0)
-
-        # State update
-        x_next = x + self.dt * gate * v_next
-
-        aux = {'alpha': alpha, 'beta': beta, 'gate': gate}
-        return x_next, v_next, aux
-
-    def init_state(self, batch_size: int, device: torch.device):
-        """Initialize state."""
-        x0 = self.mu.unsqueeze(0).expand(batch_size, -1).to(device)
-        v0 = torch.zeros(batch_size, self.output_dim, device=device)
-        return x0, v0
-
-
-class MixtureOfIntegrators(nn.Module):
-    """
-    Mixture of Integrators (MoI) - like Mixture of Experts for INL.
-
-    Routes each token to top-k integrator experts.
-    Enables sparse, conditional computation.
-
-    Benefit: Can scale capacity without scaling compute linearly
-    """
-
-    def __init__(
-        self,
-        hidden_dim: int,
-        output_dim: int,
-        num_experts: int = 8,
-        top_k: int = 2,
-        target_value: float = 5.0,
-        dt: float = 0.1
-    ):
-        """
-        Args:
-            hidden_dim: Context dimension
-            output_dim: State dimension
-            num_experts: Number of INL experts
-            top_k: Use top-k experts per token
-            target_value: Initial equilibrium
-            dt: Time step
-        """
-        super().__init__()
-
-        self.hidden_dim = hidden_dim
-        self.output_dim = output_dim
-        self.num_experts = num_experts
-        self.top_k = top_k
-        self.dt = dt
-
-        # Shared equilibrium (all experts share same μ)
-        self.mu = nn.Parameter(torch.full((output_dim,), target_value))
-
-        # Router: decides which expert(s) to use
-        self.router = nn.Linear(hidden_dim, num_experts)
-
-        # Expert-specific controllers
-        self.expert_controllers = nn.ModuleList([
-            nn.Sequential(
-                nn.Linear(hidden_dim + 2 * output_dim, 64),
-                nn.ReLU(),
-                nn.Linear(64, 3 * output_dim)  # α, β, g
-            )
-            for _ in range(num_experts)
-        ])
-
-    def forward(
-        self,
-        h: torch.Tensor,
-        x: torch.Tensor,
-        v: torch.Tensor,
-        step: int = 0
-    ) -> Tuple[torch.Tensor, torch.Tensor, Dict]:
-        """
-        Forward with expert routing.
-
-        Args:
-            h: Context [batch, hidden_dim]
-            x: State [batch, output_dim]
-            v: Velocity [batch, output_dim]
-            step: Integration step
-
-        Returns:
-            x_next, v_next, aux_info
-        """
-        batch_size = x.shape[0]
-
-        # Route: which experts to use?
-        router_logits = self.router(h)  # [batch, num_experts]
-        router_probs = F.softmax(router_logits, dim=-1)
-
-        # Select top-k experts
-        top_k_probs, top_k_indices = torch.topk(router_probs, self.top_k, dim=-1)
-        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)  # Renormalize
-
-        # Compute outputs from selected experts
-        x_next_combined = torch.zeros_like(x)
-        v_next_combined = torch.zeros_like(v)
-
-        for k in range(self.top_k):
-            expert_idx = top_k_indices[:, k]  # [batch]
-            weight = top_k_probs[:, k].unsqueeze(-1)  # [batch, 1]
-
-            # Process each sample with its selected expert
-            for i in range(batch_size):
-                exp_id = expert_idx[i].item()
-
-                # Get controller output from this expert
-                ctx_i = torch.cat([h[i:i+1], x[i:i+1], v[i:i+1]], dim=-1)
-                ctrl_out = self.expert_controllers[exp_id](ctx_i)
-                alpha_raw, beta_raw, gate_raw = torch.split(ctrl_out, self.output_dim, dim=1)
-
-                alpha = torch.sigmoid(alpha_raw)
-                beta = F.softplus(beta_raw)
-                gate = torch.sigmoid(gate_raw)
-
-                # INL dynamics
-                error = x[i:i+1] - self.mu
-                v_next_i = alpha * v[i:i+1] - beta * error
-                x_next_i = x[i:i+1] + self.dt * gate * v_next_i
-
-                # Accumulate weighted contribution
-                x_next_combined[i:i+1] += weight[i:i+1] * x_next_i
-                v_next_combined[i:i+1] += weight[i:i+1] * v_next_i
-
-        aux = {
-            'router_probs': router_probs,
-            'top_k_experts': top_k_indices,
-            'expert_weights': top_k_probs
-        }
-
-        return x_next_combined, v_next_combined, aux
-
-    def init_state(self, batch_size: int, device: torch.device):
-        """Initialize state."""
-        x0 = self.mu.unsqueeze(0).expand(batch_size, -1).to(device)
-        v0 = torch.zeros(batch_size, self.output_dim, device=device)
-        return x0, v0
-
-
-class HierarchicalEquilibriumINL(nn.Module):
-    """
-    Hierarchical Equilibrium Learning.
-
-    Instead of learning μ per dimension independently:
-    - Learn global μ_global (1 parameter)
-    - Learn local offsets per group (d_model // group_size parameters)
-
-    Benefit: Fewer parameters, better generalization
-    """
-
-    def __init__(
-        self,
-        hidden_dim: int,
-        output_dim: int,
-        group_size: int = 64,
-        target_value: float = 5.0,
-        dt: float = 0.1
-    ):
-        """
-        Args:
-            hidden_dim: Context dimension
-            output_dim: State dimension
-            group_size: Size of each group sharing offset
-            target_value: Initial global equilibrium
-            dt: Time step
-        """
-        super().__init__()
-
-        self.hidden_dim = hidden_dim
-        self.output_dim = output_dim
-        self.group_size = group_size
-        self.dt = dt
-
-        # Global equilibrium (shared by all)
-        self.mu_global = nn.Parameter(torch.tensor(target_value))
-
-        # Local offsets per group
-        num_groups = (output_dim + group_size - 1) // group_size
-        self.mu_local_offsets = nn.Parameter(torch.zeros(num_groups))
-
-        # Simple controller
-        self.controller = nn.Sequential(
-            nn.Linear(hidden_dim + 2 * output_dim, 64),
-            nn.ReLU(),
-            nn.Linear(64, 3 * output_dim)
-        )
-
-    def get_mu(self) -> torch.Tensor:
-        """
-        Compute full μ from hierarchical representation.
-
-        Returns:
-            mu: [output_dim]
-        """
-        # Repeat each group offset
-        mu_local = self.mu_local_offsets.repeat_interleave(self.group_size)
-
-        # Trim to exact size
-        mu_local = mu_local[:self.output_dim]
-
-        # Combine global + local
-        mu = self.mu_global + mu_local
-        return mu
-
-    def forward(
-        self,
-        h: torch.Tensor,
-        x: torch.Tensor,
-        v: torch.Tensor,
-        step: int = 0
-    ) -> Tuple[torch.Tensor, torch.Tensor, Dict]:
-        """Forward with hierarchical equilibrium."""
-        mu = self.get_mu()
-
-        # Controller
-        ctx = torch.cat([h, x, v], dim=-1)
-        ctrl_out = self.controller(ctx)
-        alpha_raw, beta_raw, gate_raw = torch.split(ctrl_out, self.output_dim, dim=1)
-
-        alpha = torch.sigmoid(alpha_raw)
-        beta = F.softplus(beta_raw)
-        gate = torch.sigmoid(gate_raw)
-
-        # Dynamics
-        error = x - mu
-        v_next = alpha * v - beta * error
-        x_next = x + self.dt * gate * v_next
-
-        aux = {'mu': mu, 'mu_global': self.mu_global, 'mu_offsets': self.mu_local_offsets}
-        return x_next, v_next, aux
-
-    def init_state(self, batch_size: int, device: torch.device):
-        """Initialize state."""
-        mu = self.get_mu()
-        x0 = mu.unsqueeze(0).expand(batch_size, -1).to(device)
-        v0 = torch.zeros(batch_size, self.output_dim, device=device)
-        return x0, v0
-
-    def num_mu_parameters(self) -> int:
-        """Count parameters used for μ."""
-        return 1 + self.mu_local_offsets.numel()
-
-
-def compute_advanced_optimization_gains(
-    d_model: int = 2048,
-    num_layers: int = 24,
-    hidden_controller: int = 64
-):
-    """
-    Compute parameter savings from advanced optimizations.
-
-    Args:
-        d_model: Model dimension
-        num_layers: Number of layers
-        hidden_controller: Controller hidden size
-    """
-    print("=" * 70)
-    print("ADVANCED OPTIMIZATION ANALYSIS")
-    print("=" * 70)
-
-    # 1. Shared Controllers
-    print("\n1. SHARED CONTROLLERS")
-    print("-" * 70)
-
-    # Standard: each layer has own controller
-    params_per_controller = (
-        d_model * hidden_controller +  # h projection
-        d_model * hidden_controller +  # x projection
-        d_model * hidden_controller +  # v projection
-        hidden_controller * (4 * d_model)  # output
-    )
-    standard_total = params_per_controller * num_layers
-
-    # Shared: one controller + layer modulation
-    shared_base = params_per_controller
-    layer_modulation = num_layers * 8  # 4 scalers + 4 biases per layer
-    shared_total = shared_base + layer_modulation
-
-    reduction_pct = (1 - shared_total / standard_total) * 100
-
-    print(f"  Standard (independent): {standard_total:,} params")
-    print(f"  Shared + modulation:    {shared_total:,} params")
-    print(f"  💾 REDUCTION: {reduction_pct:.1f}%")
-
-    # 2. Sparse Harmonic
-    print("\n2. SPARSE HARMONIC EXCITATION")
-    print("-" * 70)
-    sparsity = 0.1
-    compute_reduction = 1 / sparsity
-    print(f"  Sparsity: {sparsity*100:.0f}% of dimensions excited")
-    print(f"  ⚡ COMPUTE REDUCTION: {compute_reduction:.0f}x less operations")
-
-    # 3. Hierarchical μ
-    print("\n3. HIERARCHICAL EQUILIBRIUM")
-    print("-" * 70)
-    group_size = 64
-    num_groups = (d_model + group_size - 1) // group_size
-
-    standard_mu = d_model
-    hierarchical_mu = 1 + num_groups
-    mu_reduction = (1 - hierarchical_mu / standard_mu) * 100
-
-    print(f"  Standard μ:       {standard_mu:,} params")
-    print(f"  Hierarchical μ:   {hierarchical_mu:,} params (global + {num_groups} groups)")
-    print(f"  💾 REDUCTION: {mu_reduction:.1f}%")
-
-    # 4. Combined impact
-    print("\n4. COMBINED IMPACT")
-    print("-" * 70)
-    print(f"  Controller params saved:  {standard_total - shared_total:,}")
-    print(f"  Harmonic compute:         {compute_reduction:.0f}x faster")
-    print(f"  Equilibrium params saved: {standard_mu - hierarchical_mu}")
-    print(f"  Overall controller reduction: {reduction_pct:.1f}%")
-
-    print("\n" + "=" * 70)
-
-
-if __name__ == '__main__':
-    print("\n")
-
-    # Test 1: Shared Controllers
-    print("=" * 70)
-    print("TEST 1: Shared Controllers")
-    print("=" * 70)
-
-    shared_ctrl = SharedController(
-        hidden_dim=512,
-        output_dim=512,
-        num_layers=12
-    )
-
-    h = torch.randn(2, 512)
-    x = torch.randn(2, 512)
-    v = torch.randn(2, 512)
-
-    for layer_idx in range(3):
-        alpha, beta, gate, v_cand = shared_ctrl(h, x, v, layer_idx)
-        print(f"Layer {layer_idx}: alpha={alpha.mean().item():.3f}, beta={beta.mean().item():.3f}")
-
-    print(f"✅ Shared controller parameters: {shared_ctrl.num_parameters():,}")
-
-    # Test 2: Sparse Harmonic
-    print("\n" + "=" * 70)
-    print("TEST 2: Sparse Harmonic Excitation")
-    print("=" * 70)
-
-    sparse_inl = SparseHarmonicINL(
-        hidden_dim=512,
-        output_dim=512,
-        sparsity=0.1
-    )
-
-    x0, v0 = sparse_inl.init_state(2, 'cpu')
-    x_next, v_next, aux = sparse_inl(h, x0, v0, step=0)
-
-    print(f"✅ Sparse excitation: {sparse_inl.num_excited}/{sparse_inl.output_dim} dims excited")
-    print(f"   Sparsity: {sparse_inl.sparsity*100:.0f}%")
-
-    # Test 3: Mixture of Integrators
-    print("\n" + "=" * 70)
-    print("TEST 3: Mixture of Integrators")
-    print("=" * 70)
-
-    moi = MixtureOfIntegrators(
-        hidden_dim=512,
-        output_dim=512,
-        num_experts=8,
-        top_k=2
-    )
-
-    x0, v0 = moi.init_state(2, 'cpu')
-    x_next, v_next, aux = moi(h, x0, v0, step=0)
-
-    print(f"✅ MoI: {moi.num_experts} experts, top-{moi.top_k} routing")
-    print(f"   Expert distribution: {aux['top_k_experts']}")
-
-    # Test 4: Hierarchical Equilibrium
-    print("\n" + "=" * 70)
-    print("TEST 4: Hierarchical Equilibrium")
-    print("=" * 70)
-
-    hier_inl = HierarchicalEquilibriumINL(
-        hidden_dim=512,
-        output_dim=512,
-        group_size=64
-    )
-
-    x0, v0 = hier_inl.init_state(2, 'cpu')
-    x_next, v_next, aux = hier_inl(h, x0, v0)
-
-    print(f"✅ Hierarchical μ: {hier_inl.num_mu_parameters()} params (vs 512 standard)")
-    print(f"   Global μ: {aux['mu_global'].item():.3f}")
-    print(f"   Local offsets: {aux['mu_offsets'][:3].tolist()}")
-
-    # Analysis
-    print("\n")
-    compute_advanced_optimization_gains(
-        d_model=2048,
-        num_layers=24,
-        hidden_controller=64
-    )
+"""
+Advanced Optimizations for INL-LLM
+
+Implements additional efficiency techniques:
+1. Shared Controllers: Share control MLPs across layers (-15-20% params)
+2. Sparse Harmonic Excitation: Only excite subset of dimensions (-10x compute)
+3. Mixture of Integrators (MoI): Conditional computation like MoE
+4. Hierarchical Equilibrium: Global + local offsets for μ
+
+Author: Boris Peyriguère
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, Dict, List
+import math
+
+
+class SharedController(nn.Module):
+    """
+    Shared controller MLP across multiple INL layers.
+
+    Instead of each layer having its own controller (α, β, g, v_cand),
+    we use ONE shared controller + small layer-specific modulation.
+
+    Benefit: 15-20% parameter reduction on controller networks
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        output_dim: int,
+        num_layers: int,
+        hidden_controller: int = 64
+    ):
+        """
+        Args:
+            hidden_dim: Context dimension
+            output_dim: State dimension
+            num_layers: Number of layers sharing this controller
+            hidden_controller: Hidden size for controller MLP
+        """
+        super().__init__()
+
+        self.hidden_dim = hidden_dim
+        self.output_dim = output_dim
+        self.num_layers = num_layers
+
+        # Single shared controller (used by all layers)
+        self.controller_h = nn.Linear(hidden_dim, hidden_controller)
+        self.controller_x = nn.Linear(output_dim, hidden_controller)
+        self.controller_v = nn.Linear(output_dim, hidden_controller)
+        self.controller_mlp = nn.Sequential(
+            nn.ReLU(),
+            nn.Linear(hidden_controller, 4 * output_dim)
+        )
+
+        # Layer-specific modulation (tiny parameters)
+        # Each layer gets 4 scalar multipliers (α, β, g, v_cand)
+        self.layer_scalers = nn.Parameter(torch.ones(num_layers, 4))
+        self.layer_biases = nn.Parameter(torch.zeros(num_layers, 4))
+
+        # Initialize
+        self._init_weights()
+
+    def _init_weights(self):
+        """Initialize controller weights."""
+        with torch.no_grad():
+            nn.init.xavier_uniform_(self.controller_h.weight)
+            nn.init.xavier_uniform_(self.controller_x.weight)
+            nn.init.xavier_uniform_(self.controller_v.weight)
+            self.controller_h.bias.zero_()
+            self.controller_x.bias.zero_()
+            self.controller_v.bias.zero_()
+            self.controller_mlp[-1].weight.normal_(0.0, 0.01)
+
+    def forward(
+        self,
+        h: torch.Tensor,
+        x: torch.Tensor,
+        v: torch.Tensor,
+        layer_idx: int
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Compute controller parameters for specific layer.
+
+        Args:
+            h: Context [batch, hidden_dim]
+            x: State [batch, output_dim]
+            v: Velocity [batch, output_dim]
+            layer_idx: Which layer is requesting control
+
+        Returns:
+            alpha, beta, gate, v_cand (all [batch, output_dim])
+        """
+        # Shared computation
+        controller_hidden = self.controller_h(h) + self.controller_x(x) + self.controller_v(v)
+        controller_output = self.controller_mlp(controller_hidden)
+
+        # Split into components
+        alpha_base, beta_base, gate_base, v_cand_base = torch.split(
+            controller_output, self.output_dim, dim=1
+        )
+
+        # Layer-specific modulation
+        scaler = self.layer_scalers[layer_idx]  # [4]
+        bias = self.layer_biases[layer_idx]      # [4]
+
+        alpha = torch.sigmoid(alpha_base * scaler[0] + bias[0])
+        beta = F.softplus(beta_base * scaler[1] + bias[1])
+        gate = torch.sigmoid(gate_base * scaler[2] + bias[2])
+        v_cand = v_cand_base * scaler[3] + bias[3]
+
+        return alpha, beta, gate, v_cand
+
+    def num_parameters(self) -> int:
+        """Count parameters."""
+        shared = sum(p.numel() for p in [
+            self.controller_h.weight, self.controller_h.bias,
+            self.controller_x.weight, self.controller_x.bias,
+            self.controller_v.weight, self.controller_v.bias
+        ]) + sum(p.numel() for p in self.controller_mlp.parameters())
+
+        layer_specific = self.layer_scalers.numel() + self.layer_biases.numel()
+
+        return shared + layer_specific
+
+
+class SparseHarmonicINL(nn.Module):
+    """
+    INL with Sparse Harmonic Excitation.
+
+    Only applies harmonic noise to a subset of dimensions (e.g., 10%).
+    Reduces compute by 10x while maintaining exploration.
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        output_dim: int,
+        sparsity: float = 0.1,
+        target_value: float = 5.0,
+        dt: float = 0.1,
+        excitation_amplitude: float = 0.03
+    ):
+        """
+        Args:
+            hidden_dim: Context dimension
+            output_dim: State dimension
+            sparsity: Fraction of dimensions to excite (0.1 = 10%)
+            target_value: Initial equilibrium
+            dt: Time step
+            excitation_amplitude: Amplitude of excitation
+        """
+        super().__init__()
+
+        self.hidden_dim = hidden_dim
+        self.output_dim = output_dim
+        self.sparsity = sparsity
+        self.dt = dt
+
+        # Learnable μ
+        self.mu = nn.Parameter(torch.full((output_dim,), target_value))
+
+        # Excitation parameters (only for sparse subset)
+        self.num_excited = max(1, int(output_dim * sparsity))
+
+        # Fixed sparse indices (deterministic)
+        indices = torch.linspace(0, output_dim - 1, self.num_excited).long()
+        self.register_buffer('excited_indices', indices)
+
+        # Learnable excitation params (only for excited dims)
+        self.register_buffer('excitation_amplitude', torch.tensor(excitation_amplitude))
+        self.excitation_gamma = nn.Parameter(torch.ones(self.num_excited))
+        self.excitation_phi = nn.Parameter(torch.zeros(self.num_excited))
+
+        # Simple controller (for demo - would use shared in practice)
+        self.controller = nn.Sequential(
+            nn.Linear(hidden_dim + 2 * output_dim, 64),
+            nn.ReLU(),
+            nn.Linear(64, 3 * output_dim)  # α, β, g
+        )
+
+    def forward(
+        self,
+        h: torch.Tensor,
+        x: torch.Tensor,
+        v: torch.Tensor,
+        step: int = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor, Dict]:
+        """Forward with sparse excitation."""
+        batch_size = x.shape[0]
+
+        # Compute controllers
+        ctx = torch.cat([h, x, v], dim=-1)
+        controller_out = self.controller(ctx)
+        alpha_raw, beta_raw, gate_raw = torch.split(controller_out, self.output_dim, dim=1)
+
+        alpha = torch.sigmoid(alpha_raw)
+        beta = F.softplus(beta_raw)
+        gate = torch.sigmoid(gate_raw)
+
+        # Velocity update
+        error = x - self.mu
+        v_next = alpha * v - beta * error
+
+        # Sparse harmonic excitation (only on subset of dims)
+        if self.excitation_amplitude.item() > 0 and self.training:
+            t = float(step)
+            # Compute noise only for excited dimensions
+            noise_sparse = self.excitation_amplitude * torch.sin(
+                self.excitation_gamma * t + self.excitation_phi
+            )  # [num_excited]
+
+            # Apply to specific indices (sparse operation)
+            v_next[:, self.excited_indices] += noise_sparse.unsqueeze(0)
+
+        # State update
+        x_next = x + self.dt * gate * v_next
+
+        aux = {'alpha': alpha, 'beta': beta, 'gate': gate}
+        return x_next, v_next, aux
+
+    def init_state(self, batch_size: int, device: torch.device):
+        """Initialize state."""
+        x0 = self.mu.unsqueeze(0).expand(batch_size, -1).to(device)
+        v0 = torch.zeros(batch_size, self.output_dim, device=device)
+        return x0, v0
+
+
+class MixtureOfIntegrators(nn.Module):
+    """
+    Mixture of Integrators (MoI) - like Mixture of Experts for INL.
+
+    Routes each token to top-k integrator experts.
+    Enables sparse, conditional computation.
+
+    Benefit: Can scale capacity without scaling compute linearly
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        output_dim: int,
+        num_experts: int = 8,
+        top_k: int = 2,
+        target_value: float = 5.0,
+        dt: float = 0.1
+    ):
+        """
+        Args:
+            hidden_dim: Context dimension
+            output_dim: State dimension
+            num_experts: Number of INL experts
+            top_k: Use top-k experts per token
+            target_value: Initial equilibrium
+            dt: Time step
+        """
+        super().__init__()
+
+        self.hidden_dim = hidden_dim
+        self.output_dim = output_dim
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.dt = dt
+
+        # Shared equilibrium (all experts share same μ)
+        self.mu = nn.Parameter(torch.full((output_dim,), target_value))
+
+        # Router: decides which expert(s) to use
+        self.router = nn.Linear(hidden_dim, num_experts)
+
+        # Expert-specific controllers
+        self.expert_controllers = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(hidden_dim + 2 * output_dim, 64),
+                nn.ReLU(),
+                nn.Linear(64, 3 * output_dim)  # α, β, g
+            )
+            for _ in range(num_experts)
+        ])
+
+    def forward(
+        self,
+        h: torch.Tensor,
+        x: torch.Tensor,
+        v: torch.Tensor,
+        step: int = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor, Dict]:
+        """
+        Forward with expert routing.
+
+        Args:
+            h: Context [batch, hidden_dim]
+            x: State [batch, output_dim]
+            v: Velocity [batch, output_dim]
+            step: Integration step
+
+        Returns:
+            x_next, v_next, aux_info
+        """
+        batch_size = x.shape[0]
+
+        # Route: which experts to use?
+        router_logits = self.router(h)  # [batch, num_experts]
+        router_probs = F.softmax(router_logits, dim=-1)
+
+        # Select top-k experts
+        top_k_probs, top_k_indices = torch.topk(router_probs, self.top_k, dim=-1)
+        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)  # Renormalize
+
+        # Compute outputs from selected experts
+        x_next_combined = torch.zeros_like(x)
+        v_next_combined = torch.zeros_like(v)
+
+        for k in range(self.top_k):
+            expert_idx = top_k_indices[:, k]  # [batch]
+            weight = top_k_probs[:, k].unsqueeze(-1)  # [batch, 1]
+
+            # Process each sample with its selected expert
+            for i in range(batch_size):
+                exp_id = expert_idx[i].item()
+
+                # Get controller output from this expert
+                ctx_i = torch.cat([h[i:i+1], x[i:i+1], v[i:i+1]], dim=-1)
+                ctrl_out = self.expert_controllers[exp_id](ctx_i)
+                alpha_raw, beta_raw, gate_raw = torch.split(ctrl_out, self.output_dim, dim=1)
+
+                alpha = torch.sigmoid(alpha_raw)
+                beta = F.softplus(beta_raw)
+                gate = torch.sigmoid(gate_raw)
+
+                # INL dynamics
+                error = x[i:i+1] - self.mu
+                v_next_i = alpha * v[i:i+1] - beta * error
+                x_next_i = x[i:i+1] + self.dt * gate * v_next_i
+
+                # Accumulate weighted contribution
+                x_next_combined[i:i+1] += weight[i:i+1] * x_next_i
+                v_next_combined[i:i+1] += weight[i:i+1] * v_next_i
+
+        aux = {
+            'router_probs': router_probs,
+            'top_k_experts': top_k_indices,
+            'expert_weights': top_k_probs
+        }
+
+        return x_next_combined, v_next_combined, aux
+
+    def init_state(self, batch_size: int, device: torch.device):
+        """Initialize state."""
+        x0 = self.mu.unsqueeze(0).expand(batch_size, -1).to(device)
+        v0 = torch.zeros(batch_size, self.output_dim, device=device)
+        return x0, v0
+
+
+class HierarchicalEquilibriumINL(nn.Module):
+    """
+    Hierarchical Equilibrium Learning.
+
+    Instead of learning μ per dimension independently:
+    - Learn global μ_global (1 parameter)
+    - Learn local offsets per group (d_model // group_size parameters)
+
+    Benefit: Fewer parameters, better generalization
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        output_dim: int,
+        group_size: int = 64,
+        target_value: float = 5.0,
+        dt: float = 0.1
+    ):
+        """
+        Args:
+            hidden_dim: Context dimension
+            output_dim: State dimension
+            group_size: Size of each group sharing offset
+            target_value: Initial global equilibrium
+            dt: Time step
+        """
+        super().__init__()
+
+        self.hidden_dim = hidden_dim
+        self.output_dim = output_dim
+        self.group_size = group_size
+        self.dt = dt
+
+        # Global equilibrium (shared by all)
+        self.mu_global = nn.Parameter(torch.tensor(target_value))
+
+        # Local offsets per group
+        num_groups = (output_dim + group_size - 1) // group_size
+        self.mu_local_offsets = nn.Parameter(torch.zeros(num_groups))
+
+        # Simple controller
+        self.controller = nn.Sequential(
+            nn.Linear(hidden_dim + 2 * output_dim, 64),
+            nn.ReLU(),
+            nn.Linear(64, 3 * output_dim)
+        )
+
+    def get_mu(self) -> torch.Tensor:
+        """
+        Compute full μ from hierarchical representation.
+
+        Returns:
+            mu: [output_dim]
+        """
+        # Repeat each group offset
+        mu_local = self.mu_local_offsets.repeat_interleave(self.group_size)
+
+        # Trim to exact size
+        mu_local = mu_local[:self.output_dim]
+
+        # Combine global + local
+        mu = self.mu_global + mu_local
+        return mu
+
+    def forward(
+        self,
+        h: torch.Tensor,
+        x: torch.Tensor,
+        v: torch.Tensor,
+        step: int = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor, Dict]:
+        """Forward with hierarchical equilibrium."""
+        mu = self.get_mu()
+
+        # Controller
+        ctx = torch.cat([h, x, v], dim=-1)
+        ctrl_out = self.controller(ctx)
+        alpha_raw, beta_raw, gate_raw = torch.split(ctrl_out, self.output_dim, dim=1)
+
+        alpha = torch.sigmoid(alpha_raw)
+        beta = F.softplus(beta_raw)
+        gate = torch.sigmoid(gate_raw)
+
+        # Dynamics
+        error = x - mu
+        v_next = alpha * v - beta * error
+        x_next = x + self.dt * gate * v_next
+
+        aux = {'mu': mu, 'mu_global': self.mu_global, 'mu_offsets': self.mu_local_offsets}
+        return x_next, v_next, aux
+
+    def init_state(self, batch_size: int, device: torch.device):
+        """Initialize state."""
+        mu = self.get_mu()
+        x0 = mu.unsqueeze(0).expand(batch_size, -1).to(device)
+        v0 = torch.zeros(batch_size, self.output_dim, device=device)
+        return x0, v0
+
+    def num_mu_parameters(self) -> int:
+        """Count parameters used for μ."""
+        return 1 + self.mu_local_offsets.numel()
+
+
+def compute_advanced_optimization_gains(
+    d_model: int = 2048,
+    num_layers: int = 24,
+    hidden_controller: int = 64
+):
+    """
+    Compute parameter savings from advanced optimizations.
+
+    Args:
+        d_model: Model dimension
+        num_layers: Number of layers
+        hidden_controller: Controller hidden size
+    """
+    print("=" * 70)
+    print("ADVANCED OPTIMIZATION ANALYSIS")
+    print("=" * 70)
+
+    # 1. Shared Controllers
+    print("\n1. SHARED CONTROLLERS")
+    print("-" * 70)
+
+    # Standard: each layer has own controller
+    params_per_controller = (
+        d_model * hidden_controller +  # h projection
+        d_model * hidden_controller +  # x projection
+        d_model * hidden_controller +  # v projection
+        hidden_controller * (4 * d_model)  # output
+    )
+    standard_total = params_per_controller * num_layers
+
+    # Shared: one controller + layer modulation
+    shared_base = params_per_controller
+    layer_modulation = num_layers * 8  # 4 scalers + 4 biases per layer
+    shared_total = shared_base + layer_modulation
+
+    reduction_pct = (1 - shared_total / standard_total) * 100
+
+    print(f"  Standard (independent): {standard_total:,} params")
+    print(f"  Shared + modulation:    {shared_total:,} params")
+    print(f"  💾 REDUCTION: {reduction_pct:.1f}%")
+
+    # 2. Sparse Harmonic
+    print("\n2. SPARSE HARMONIC EXCITATION")
+    print("-" * 70)
+    sparsity = 0.1
+    compute_reduction = 1 / sparsity
+    print(f"  Sparsity: {sparsity*100:.0f}% of dimensions excited")
+    print(f"  ⚡ COMPUTE REDUCTION: {compute_reduction:.0f}x less operations")
+
+    # 3. Hierarchical μ
+    print("\n3. HIERARCHICAL EQUILIBRIUM")
+    print("-" * 70)
+    group_size = 64
+    num_groups = (d_model + group_size - 1) // group_size
+
+    standard_mu = d_model
+    hierarchical_mu = 1 + num_groups
+    mu_reduction = (1 - hierarchical_mu / standard_mu) * 100
+
+    print(f"  Standard μ:       {standard_mu:,} params")
+    print(f"  Hierarchical μ:   {hierarchical_mu:,} params (global + {num_groups} groups)")
+    print(f"  💾 REDUCTION: {mu_reduction:.1f}%")
+
+    # 4. Combined impact
+    print("\n4. COMBINED IMPACT")
+    print("-" * 70)
+    print(f"  Controller params saved:  {standard_total - shared_total:,}")
+    print(f"  Harmonic compute:         {compute_reduction:.0f}x faster")
+    print(f"  Equilibrium params saved: {standard_mu - hierarchical_mu}")
+    print(f"  Overall controller reduction: {reduction_pct:.1f}%")
+
+    print("\n" + "=" * 70)
+
+
+if __name__ == '__main__':
+    print("\n")
+
+    # Test 1: Shared Controllers
+    print("=" * 70)
+    print("TEST 1: Shared Controllers")
+    print("=" * 70)
+
+    shared_ctrl = SharedController(
+        hidden_dim=512,
+        output_dim=512,
+        num_layers=12
+    )
+
+    h = torch.randn(2, 512)
+    x = torch.randn(2, 512)
+    v = torch.randn(2, 512)
+
+    for layer_idx in range(3):
+        alpha, beta, gate, v_cand = shared_ctrl(h, x, v, layer_idx)
+        print(f"Layer {layer_idx}: alpha={alpha.mean().item():.3f}, beta={beta.mean().item():.3f}")
+
+    print(f"✅ Shared controller parameters: {shared_ctrl.num_parameters():,}")
+
+    # Test 2: Sparse Harmonic
+    print("\n" + "=" * 70)
+    print("TEST 2: Sparse Harmonic Excitation")
+    print("=" * 70)
+
+    sparse_inl = SparseHarmonicINL(
+        hidden_dim=512,
+        output_dim=512,
+        sparsity=0.1
+    )
+
+    x0, v0 = sparse_inl.init_state(2, 'cpu')
+    x_next, v_next, aux = sparse_inl(h, x0, v0, step=0)
+
+    print(f"✅ Sparse excitation: {sparse_inl.num_excited}/{sparse_inl.output_dim} dims excited")
+    print(f"   Sparsity: {sparse_inl.sparsity*100:.0f}%")
+
+    # Test 3: Mixture of Integrators
+    print("\n" + "=" * 70)
+    print("TEST 3: Mixture of Integrators")
+    print("=" * 70)
+
+    moi = MixtureOfIntegrators(
+        hidden_dim=512,
+        output_dim=512,
+        num_experts=8,
+        top_k=2
+    )
+
+    x0, v0 = moi.init_state(2, 'cpu')
+    x_next, v_next, aux = moi(h, x0, v0, step=0)
+
+    print(f"✅ MoI: {moi.num_experts} experts, top-{moi.top_k} routing")
+    print(f"   Expert distribution: {aux['top_k_experts']}")
+
+    # Test 4: Hierarchical Equilibrium
+    print("\n" + "=" * 70)
+    print("TEST 4: Hierarchical Equilibrium")
+    print("=" * 70)
+
+    hier_inl = HierarchicalEquilibriumINL(
+        hidden_dim=512,
+        output_dim=512,
+        group_size=64
+    )
+
+    x0, v0 = hier_inl.init_state(2, 'cpu')
+    x_next, v_next, aux = hier_inl(h, x0, v0)
+
+    print(f"✅ Hierarchical μ: {hier_inl.num_mu_parameters()} params (vs 512 standard)")
+    print(f"   Global μ: {aux['mu_global'].item():.3f}")
+    print(f"   Local offsets: {aux['mu_offsets'][:3].tolist()}")
+
+    # Analysis
+    print("\n")
+    compute_advanced_optimization_gains(
+        d_model=2048,
+        num_layers=24,
+        hidden_controller=64
+    )

inl_llm/optimizations/optimizations.py

@@ -1,564 +1,564 @@
-"""
-Optimizations for INL-LLM Architecture
-
-This module implements key optimizations to maximize efficiency:
-1. Low-Rank Embeddings: Reduce embedding parameters by 70-80%
-2. Adaptive Early Stopping: 2x speedup in inference
-3. Gradient Checkpointing: Enable scaling to 100B+ parameters
-
-Author: Boris Peyriguère
-"""
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from typing import Optional, Tuple, Dict
-import math
-
-
-class LowRankEmbedding(nn.Module):
-    """
-    Low-rank factorized embedding layer.
-
-    Replaces standard embedding (vocab_size × d_model) with:
-    - Low-rank embedding (vocab_size × rank)
-    - Projection matrix (rank × d_model)
-
-    Memory savings example:
-    - Standard: 50k × 2048 = 102M parameters
-    - Low-rank: 50k × 256 + 256 × 2048 = 13.3M parameters
-    - Savings: 87% reduction!
-    """
-
-    def __init__(
-        self,
-        vocab_size: int,
-        d_model: int,
-        rank: Optional[int] = None,
-        rank_ratio: float = 0.125
-    ):
-        """
-        Args:
-            vocab_size: Size of vocabulary
-            d_model: Model dimension
-            rank: Explicit rank (if None, computed as d_model * rank_ratio)
-            rank_ratio: Ratio of rank to d_model (default: 0.125 = 1/8)
-        """
-        super().__init__()
-
-        if rank is None:
-            rank = max(64, int(d_model * rank_ratio))  # At least 64
-
-        self.vocab_size = vocab_size
-        self.d_model = d_model
-        self.rank = rank
-
-        # Low-rank factorization
-        self.embed_low = nn.Embedding(vocab_size, rank)
-        self.project_up = nn.Linear(rank, d_model, bias=False)
-
-        # Initialize
-        nn.init.normal_(self.embed_low.weight, mean=0.0, std=0.02)
-        nn.init.normal_(self.project_up.weight, mean=0.0, std=0.02)
-
-    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            input_ids: [batch_size, seq_len]
-
-        Returns:
-            embeddings: [batch_size, seq_len, d_model]
-        """
-        low_rank_embed = self.embed_low(input_ids)  # [B, S, rank]
-        full_embed = self.project_up(low_rank_embed)  # [B, S, d_model]
-        return full_embed
-
-    def num_parameters(self) -> int:
-        """Count parameters in this layer."""
-        return self.vocab_size * self.rank + self.rank * self.d_model
-
-    def __repr__(self) -> str:
-        std_params = self.vocab_size * self.d_model
-        our_params = self.num_parameters()
-        reduction = (1 - our_params / std_params) * 100
-
-        return (
-            f"{self.__class__.__name__}(\n"
-            f"  vocab_size={self.vocab_size}, d_model={self.d_model}, rank={self.rank}\n"
-            f"  parameters: {our_params:,} (vs {std_params:,} standard)\n"
-            f"  reduction: {reduction:.1f}%\n"
-            f")"
-        )
-
-
-class AdaptiveIntegratorNeuronLayer(nn.Module):
-    """
-    Integrator Neuron Layer with Adaptive Early Stopping.
-
-    Dynamically adjusts number of integration steps based on convergence.
-    When error is small enough, stops iterating early.
-
-    Benefits:
-    - 30-50% faster inference (fewer iterations needed)
-    - Same training dynamics (max iterations used)
-    - Automatic adaptation per sample
-    """
-
-    def __init__(
-        self,
-        inl_layer: nn.Module,
-        convergence_threshold: float = 0.01,
-        min_iterations: int = 3,
-        max_iterations: int = 10,
-        check_interval: int = 1
-    ):
-        """
-        Args:
-            inl_layer: Base IntegratorNeuronLayer to wrap
-            convergence_threshold: L2 norm threshold for early stopping
-            min_iterations: Minimum iterations before checking convergence
-            max_iterations: Maximum iterations (used during training)
-            check_interval: Check convergence every N iterations
-        """
-        super().__init__()
-
-        self.inl = inl_layer
-        self.convergence_threshold = convergence_threshold
-        self.min_iterations = min_iterations
-        self.max_iterations = max_iterations
-        self.check_interval = check_interval
-
-        # Statistics tracking
-        self.register_buffer('avg_iterations', torch.tensor(0.0))
-        self.register_buffer('num_forwards', torch.tensor(0))
-
-    def forward(
-        self,
-        h: torch.Tensor,
-        num_iterations: Optional[int] = None,
-        use_early_stopping: bool = None,
-        return_trajectory: bool = False
-    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Dict]]:
-        """
-        Forward with adaptive early stopping.
-
-        Args:
-            h: Context embedding [batch_size, hidden_dim]
-            num_iterations: Override max iterations (if None, use self.max_iterations)
-            use_early_stopping: Enable early stopping (default: not training)
-            return_trajectory: Return full trajectory
-
-        Returns:
-            x_final: Final state [batch_size, output_dim]
-            v_final: Final velocity [batch_size, output_dim]
-            info: Dict with 'iterations_used', 'converged', optional 'trajectory'
-        """
-        batch_size = h.shape[0]
-        device = h.device
-
-        if num_iterations is None:
-            num_iterations = self.max_iterations
-
-        if use_early_stopping is None:
-            use_early_stopping = not self.training
-
-        # Initialize state and velocity
-        x, v = self.inl.init_state(batch_size, device)
-
-        # Track trajectory if needed
-        if return_trajectory:
-            x_traj = [x.detach().cpu()]
-            v_traj = [v.detach().cpu()]
-
-        converged = torch.zeros(batch_size, dtype=torch.bool, device=device)
-        iterations_used = torch.zeros(batch_size, dtype=torch.long, device=device)
-
-        for t in range(num_iterations):
-            # Run integration step
-            x_next, v_next, aux = self.inl(h, x, v, step=t, return_aux=True)
-
-            # Update iterations counter for non-converged samples
-            iterations_used[~converged] += 1
-
-            # Check convergence (after min_iterations)
-            if use_early_stopping and t >= self.min_iterations and t % self.check_interval == 0:
-                # Compute error norm per sample
-                error = aux['error']  # [batch_size, output_dim]
-                error_norm = torch.norm(error, dim=-1)  # [batch_size]
-
-                # Mark newly converged samples
-                newly_converged = (error_norm < self.convergence_threshold) & (~converged)
-                converged = converged | newly_converged
-
-                # If all samples converged, stop early
-                if converged.all():
-                    x, v = x_next, v_next
-                    if return_trajectory:
-                        x_traj.append(x.detach().cpu())
-                        v_traj.append(v.detach().cpu())
-                    break
-
-            x, v = x_next, v_next
-
-            if return_trajectory:
-                x_traj.append(x.detach().cpu())
-                v_traj.append(v.detach().cpu())
-
-        # Update statistics (exponential moving average)
-        if not self.training:
-            avg_iters = iterations_used.float().mean()
-            self.num_forwards += 1
-            alpha = 0.99
-            self.avg_iterations = alpha * self.avg_iterations + (1 - alpha) * avg_iters
-
-        info = {
-            'iterations_used': iterations_used,
-            'converged': converged,
-            'avg_iterations': self.avg_iterations.item()
-        }
-
-        if return_trajectory:
-            info['trajectory'] = {
-                'x': torch.stack(x_traj, dim=1),  # [B, T+1, D]
-                'v': torch.stack(v_traj, dim=1)
-            }
-
-        return x, v, info
-
-    def reset_statistics(self):
-        """Reset tracking statistics."""
-        self.avg_iterations.zero_()
-        self.num_forwards.zero_()
-
-
-class AdaptiveHierarchicalINL(nn.Module):
-    """
-    Adaptive wrapper for HierarchicalEquilibriumINL with early stopping.
-
-    Specifically designed for INL blocks in language models.
-    Monitors velocity (rate of change) instead of error for convergence detection.
-    """
-
-    def __init__(
-        self,
-        inl_layer: nn.Module,
-        convergence_threshold: float = 0.001,
-        min_iterations: int = 3,
-        max_iterations: int = 12,
-        check_interval: int = 1
-    ):
-        """
-        Args:
-            inl_layer: HierarchicalEquilibriumINL to wrap
-            convergence_threshold: Velocity threshold for early stopping
-            min_iterations: Minimum iterations before checking convergence
-            max_iterations: Maximum iterations (used during training)
-            check_interval: Check convergence every N iterations
-        """
-        super().__init__()
-
-        self.inl = inl_layer
-        self.convergence_threshold = convergence_threshold
-        self.min_iterations = min_iterations
-        self.max_iterations = max_iterations
-        self.check_interval = check_interval
-
-        # Statistics tracking
-        self.register_buffer('avg_iterations', torch.tensor(0.0))
-        self.register_buffer('num_forwards', torch.tensor(0))
-
-    def forward(
-        self,
-        h: torch.Tensor,
-        x: torch.Tensor,
-        v: torch.Tensor,
-        step: int = 0
-    ):
-        """
-        Forward pass - compatible with INL block usage.
-
-        Note: For use in INL blocks, this is called per-iteration.
-        Early stopping is handled at the block level.
-        """
-        return self.inl(h, x, v, step)
-
-    def forward_adaptive(
-        self,
-        h: torch.Tensor,
-        initial_x: torch.Tensor,
-        initial_v: torch.Tensor,
-        num_iterations: Optional[int] = None,
-        use_early_stopping: bool = None,
-        return_trajectory: bool = False
-    ):
-        """
-        Full adaptive forward with early stopping control.
-
-        Use this method when you want full control over iterations.
-        """
-        batch_size = h.shape[0]
-        device = h.device
-
-        if num_iterations is None:
-            num_iterations = self.max_iterations
-
-        if use_early_stopping is None:
-            use_early_stopping = not self.training
-
-        x, v = initial_x, initial_v
-
-        # Track trajectory if needed
-        x_traj = [x.clone()] if return_trajectory else None
-        v_traj = [v.clone()] if return_trajectory else None
-
-        converged = torch.zeros(batch_size, dtype=torch.bool, device=device)
-        iterations_used = torch.zeros(batch_size, dtype=torch.long, device=device)
-
-        for t in range(num_iterations):
-            x_prev = x.clone()
-
-            # Run integration step
-            x_next, v_next, aux = self.inl(h, x, v, step=t)
-
-            # Update iterations counter for non-converged samples
-            iterations_used[~converged] += 1
-
-            # Check convergence based on velocity (rate of change)
-            if use_early_stopping and t >= self.min_iterations and t % self.check_interval == 0:
-                # Compute change in state
-                delta_x = torch.norm(x_next - x_prev, dim=-1)  # [batch_size]
-
-                # Mark newly converged samples
-                newly_converged = (delta_x < self.convergence_threshold) & (~converged)
-                converged = converged | newly_converged
-
-                # If all samples converged, stop early
-                if converged.all():
-                    x, v = x_next, v_next
-                    if return_trajectory:
-                        x_traj.append(x.clone())
-                        v_traj.append(v.clone())
-                    break
-
-            x, v = x_next, v_next
-
-            if return_trajectory:
-                x_traj.append(x.clone())
-                v_traj.append(v.clone())
-
-        # Update statistics (exponential moving average)
-        if not self.training:
-            avg_iters = iterations_used.float().mean()
-            self.num_forwards += 1
-            alpha = 0.99
-            self.avg_iterations = alpha * self.avg_iterations + (1 - alpha) * avg_iters
-
-        result = {
-            'x': x,
-            'v': v,
-            'iterations_used': iterations_used,
-            'converged': converged,
-            'avg_iterations': self.avg_iterations.item(),
-            'mu': aux.get('mu'),
-            'mu_global': aux.get('mu_global'),
-            'mu_offsets': aux.get('mu_offsets')
-        }
-
-        if return_trajectory:
-            result['x_trajectory'] = torch.stack(x_traj, dim=1)  # [B, T+1, D]
-            result['v_trajectory'] = torch.stack(v_traj, dim=1)
-
-        return x, v, result
-
-    def reset_statistics(self):
-        """Reset tracking statistics."""
-        self.avg_iterations.zero_()
-        self.num_forwards.zero_()
-
-
-class GradientCheckpointedINL(nn.Module):
-    """
-    Wrapper for IntegratorNeuronLayer with gradient checkpointing.
-
-    Trades compute for memory:
-    - Forward: Normal computation
-    - Backward: Recompute forward instead of storing activations
-
-    Memory savings: 50-70% during training
-    Cost: ~30% slower backward pass (but worth it for large models!)
-    """
-
-    def __init__(self, inl_layer: nn.Module):
-        """
-        Args:
-            inl_layer: IntegratorNeuronLayer to wrap
-        """
-        super().__init__()
-        self.inl = inl_layer
-
-    def forward(
-        self,
-        h: torch.Tensor,
-        x: torch.Tensor,
-        v: torch.Tensor,
-        step: int = 0,
-        return_aux: bool = True
-    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Dict]]:
-        """
-        Forward with gradient checkpointing.
-
-        Uses torch.utils.checkpoint to save memory during backward pass.
-        """
-        if self.training:
-            # Use checkpointing during training
-            return torch.utils.checkpoint.checkpoint(
-                self._forward_impl,
-                h, x, v, step, return_aux,
-                use_reentrant=False
-            )
-        else:
-            # No checkpointing during inference
-            return self._forward_impl(h, x, v, step, return_aux)
-
-    def _forward_impl(
-        self,
-        h: torch.Tensor,
-        x: torch.Tensor,
-        v: torch.Tensor,
-        step: int,
-        return_aux: bool
-    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Dict]]:
-        """Actual forward implementation."""
-        return self.inl(h, x, v, step, return_aux)
-
-    def init_state(self, batch_size: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Delegate to wrapped layer."""
-        return self.inl.init_state(batch_size, device)
-
-    def __getattr__(self, name: str):
-        """Delegate attribute access to wrapped layer."""
-        try:
-            return super().__getattr__(name)
-        except AttributeError:
-            return getattr(self.inl, name)
-
-
-def compute_parameter_reduction(
-    vocab_size: int,
-    d_model: int,
-    rank_ratio: float = 0.125
-) -> Dict[str, float]:
-    """
-    Compute parameter reduction from using low-rank embeddings.
-
-    Args:
-        vocab_size: Vocabulary size
-        d_model: Model dimension
-        rank_ratio: Rank ratio for low-rank embedding
-
-    Returns:
-        Dictionary with parameter counts and reduction percentage
-    """
-    rank = max(64, int(d_model * rank_ratio))
-
-    standard_params = vocab_size * d_model
-    lowrank_params = vocab_size * rank + rank * d_model
-
-    reduction_pct = (1 - lowrank_params / standard_params) * 100
-
-    return {
-        'standard_params': standard_params,
-        'lowrank_params': lowrank_params,
-        'reduction_percent': reduction_pct,
-        'rank': rank,
-        'memory_mb_standard': standard_params * 4 / 1e6,  # FP32
-        'memory_mb_lowrank': lowrank_params * 4 / 1e6
-    }
-
-
-def print_optimization_summary(
-    vocab_size: int,
-    d_model: int,
-    num_layers: int,
-    rank_ratio: float = 0.125
-):
-    """
-    Print summary of optimization benefits.
-
-    Args:
-        vocab_size: Vocabulary size
-        d_model: Model dimension
-        num_layers: Number of layers
-        rank_ratio: Low-rank embedding ratio
-    """
-    print("=" * 70)
-    print("INL-LLM OPTIMIZATION SUMMARY")
-    print("=" * 70)
-
-    # Low-rank embedding savings
-    embed_stats = compute_parameter_reduction(vocab_size, d_model, rank_ratio)
-
-    print("\n1. LOW-RANK EMBEDDINGS")
-    print("-" * 70)
-    print(f"  Standard embedding:  {embed_stats['standard_params']:>12,} params "
-          f"({embed_stats['memory_mb_standard']:>6.1f} MB)")
-    print(f"  Low-rank embedding:  {embed_stats['lowrank_params']:>12,} params "
-          f"({embed_stats['memory_mb_lowrank']:>6.1f} MB)")
-    print(f"  Rank: {embed_stats['rank']}")
-    print(f"  💾 REDUCTION: {embed_stats['reduction_percent']:.1f}%")
-
-    print("\n2. ADAPTIVE EARLY STOPPING")
-    print("-" * 70)
-    print("  Training:   Uses max iterations (no change)")
-    print("  Inference:  Adaptive iterations based on convergence")
-    print("  ⚡ SPEEDUP: 30-50% faster inference")
-    print("  Typical iterations: 5-7 (vs 10 max)")
-
-    print("\n3. GRADIENT CHECKPOINTING")
-    print("-" * 70)
-    print("  Memory reduction: ~50-70% during training")
-    print("  Compute overhead: ~30% slower backward")
-    print("  Enables scaling to: 2-3x larger models")
-    print("  🚀 BENEFIT: Train 100B+ models on consumer GPUs")
-
-    print("\n4. COMBINED IMPACT")
-    print("-" * 70)
-    saved_params = embed_stats['standard_params'] - embed_stats['lowrank_params']
-    print(f"  Total parameters saved: {saved_params:,}")
-    print(f"  Memory saved (embeddings): {embed_stats['memory_mb_standard'] - embed_stats['memory_mb_lowrank']:.1f} MB")
-    print(f"  Inference speedup: 30-50%")
-    print(f"  Training memory: -50-70%")
-
-    print("\n" + "=" * 70)
-    print("✅ OPTIMIZATIONS READY TO USE")
-    print("=" * 70)
-
-
-if __name__ == '__main__':
-    print("\n")
-    print_optimization_summary(
-        vocab_size=50000,
-        d_model=2048,
-        num_layers=24,
-        rank_ratio=0.125
-    )
-
-    print("\n\nEXAMPLE USAGE:\n")
-    print("# 1. Low-Rank Embeddings")
-    print("from optimizations import LowRankEmbedding")
-    print("embed = LowRankEmbedding(vocab_size=50000, d_model=2048, rank_ratio=0.125)")
-    print()
-
-    print("# 2. Adaptive Early Stopping")
-    print("from optimizations import AdaptiveIntegratorNeuronLayer")
-    print("adaptive_inl = AdaptiveIntegratorNeuronLayer(")
-    print("    inl_layer=base_inl,")
-    print("    convergence_threshold=0.01,")
-    print("    max_iterations=10")
-    print(")")
-    print()
-
-    print("# 3. Gradient Checkpointing")
-    print("from optimizations import GradientCheckpointedINL")
-    print("checkpointed_inl = GradientCheckpointedINL(base_inl)")
-    print()
+"""
+Optimizations for INL-LLM Architecture
+
+This module implements key optimizations to maximize efficiency:
+1. Low-Rank Embeddings: Reduce embedding parameters by 70-80%
+2. Adaptive Early Stopping: 2x speedup in inference
+3. Gradient Checkpointing: Enable scaling to 100B+ parameters
+
+Author: Boris Peyriguère
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, Dict
+import math
+
+
+class LowRankEmbedding(nn.Module):
+    """
+    Low-rank factorized embedding layer.
+
+    Replaces standard embedding (vocab_size × d_model) with:
+    - Low-rank embedding (vocab_size × rank)
+    - Projection matrix (rank × d_model)
+
+    Memory savings example:
+    - Standard: 50k × 2048 = 102M parameters
+    - Low-rank: 50k × 256 + 256 × 2048 = 13.3M parameters
+    - Savings: 87% reduction!
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        d_model: int,
+        rank: Optional[int] = None,
+        rank_ratio: float = 0.125
+    ):
+        """
+        Args:
+            vocab_size: Size of vocabulary
+            d_model: Model dimension
+            rank: Explicit rank (if None, computed as d_model * rank_ratio)
+            rank_ratio: Ratio of rank to d_model (default: 0.125 = 1/8)
+        """
+        super().__init__()
+
+        if rank is None:
+            rank = max(64, int(d_model * rank_ratio))  # At least 64
+
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.rank = rank
+
+        # Low-rank factorization
+        self.embed_low = nn.Embedding(vocab_size, rank)
+        self.project_up = nn.Linear(rank, d_model, bias=False)
+
+        # Initialize
+        nn.init.normal_(self.embed_low.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.project_up.weight, mean=0.0, std=0.02)
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            input_ids: [batch_size, seq_len]
+
+        Returns:
+            embeddings: [batch_size, seq_len, d_model]
+        """
+        low_rank_embed = self.embed_low(input_ids)  # [B, S, rank]
+        full_embed = self.project_up(low_rank_embed)  # [B, S, d_model]
+        return full_embed
+
+    def num_parameters(self) -> int:
+        """Count parameters in this layer."""
+        return self.vocab_size * self.rank + self.rank * self.d_model
+
+    def __repr__(self) -> str:
+        std_params = self.vocab_size * self.d_model
+        our_params = self.num_parameters()
+        reduction = (1 - our_params / std_params) * 100
+
+        return (
+            f"{self.__class__.__name__}(\n"
+            f"  vocab_size={self.vocab_size}, d_model={self.d_model}, rank={self.rank}\n"
+            f"  parameters: {our_params:,} (vs {std_params:,} standard)\n"
+            f"  reduction: {reduction:.1f}%\n"
+            f")"
+        )
+
+
+class AdaptiveIntegratorNeuronLayer(nn.Module):
+    """
+    Integrator Neuron Layer with Adaptive Early Stopping.
+
+    Dynamically adjusts number of integration steps based on convergence.
+    When error is small enough, stops iterating early.
+
+    Benefits:
+    - 30-50% faster inference (fewer iterations needed)
+    - Same training dynamics (max iterations used)
+    - Automatic adaptation per sample
+    """
+
+    def __init__(
+        self,
+        inl_layer: nn.Module,
+        convergence_threshold: float = 0.01,
+        min_iterations: int = 3,
+        max_iterations: int = 10,
+        check_interval: int = 1
+    ):
+        """
+        Args:
+            inl_layer: Base IntegratorNeuronLayer to wrap
+            convergence_threshold: L2 norm threshold for early stopping
+            min_iterations: Minimum iterations before checking convergence
+            max_iterations: Maximum iterations (used during training)
+            check_interval: Check convergence every N iterations
+        """
+        super().__init__()
+
+        self.inl = inl_layer
+        self.convergence_threshold = convergence_threshold
+        self.min_iterations = min_iterations
+        self.max_iterations = max_iterations
+        self.check_interval = check_interval
+
+        # Statistics tracking
+        self.register_buffer('avg_iterations', torch.tensor(0.0))
+        self.register_buffer('num_forwards', torch.tensor(0))
+
+    def forward(
+        self,
+        h: torch.Tensor,
+        num_iterations: Optional[int] = None,
+        use_early_stopping: bool = None,
+        return_trajectory: bool = False
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Dict]]:
+        """
+        Forward with adaptive early stopping.
+
+        Args:
+            h: Context embedding [batch_size, hidden_dim]
+            num_iterations: Override max iterations (if None, use self.max_iterations)
+            use_early_stopping: Enable early stopping (default: not training)
+            return_trajectory: Return full trajectory
+
+        Returns:
+            x_final: Final state [batch_size, output_dim]
+            v_final: Final velocity [batch_size, output_dim]
+            info: Dict with 'iterations_used', 'converged', optional 'trajectory'
+        """
+        batch_size = h.shape[0]
+        device = h.device
+
+        if num_iterations is None:
+            num_iterations = self.max_iterations
+
+        if use_early_stopping is None:
+            use_early_stopping = not self.training
+
+        # Initialize state and velocity
+        x, v = self.inl.init_state(batch_size, device)
+
+        # Track trajectory if needed
+        if return_trajectory:
+            x_traj = [x.detach().cpu()]
+            v_traj = [v.detach().cpu()]
+
+        converged = torch.zeros(batch_size, dtype=torch.bool, device=device)
+        iterations_used = torch.zeros(batch_size, dtype=torch.long, device=device)
+
+        for t in range(num_iterations):
+            # Run integration step
+            x_next, v_next, aux = self.inl(h, x, v, step=t, return_aux=True)
+
+            # Update iterations counter for non-converged samples
+            iterations_used[~converged] += 1
+
+            # Check convergence (after min_iterations)
+            if use_early_stopping and t >= self.min_iterations and t % self.check_interval == 0:
+                # Compute error norm per sample
+                error = aux['error']  # [batch_size, output_dim]
+                error_norm = torch.norm(error, dim=-1)  # [batch_size]
+
+                # Mark newly converged samples
+                newly_converged = (error_norm < self.convergence_threshold) & (~converged)
+                converged = converged | newly_converged
+
+                # If all samples converged, stop early
+                if converged.all():
+                    x, v = x_next, v_next
+                    if return_trajectory:
+                        x_traj.append(x.detach().cpu())
+                        v_traj.append(v.detach().cpu())
+                    break
+
+            x, v = x_next, v_next
+
+            if return_trajectory:
+                x_traj.append(x.detach().cpu())
+                v_traj.append(v.detach().cpu())
+
+        # Update statistics (exponential moving average)
+        if not self.training:
+            avg_iters = iterations_used.float().mean()
+            self.num_forwards += 1
+            alpha = 0.99
+            self.avg_iterations = alpha * self.avg_iterations + (1 - alpha) * avg_iters
+
+        info = {
+            'iterations_used': iterations_used,
+            'converged': converged,
+            'avg_iterations': self.avg_iterations.item()
+        }
+
+        if return_trajectory:
+            info['trajectory'] = {
+                'x': torch.stack(x_traj, dim=1),  # [B, T+1, D]
+                'v': torch.stack(v_traj, dim=1)
+            }
+
+        return x, v, info
+
+    def reset_statistics(self):
+        """Reset tracking statistics."""
+        self.avg_iterations.zero_()
+        self.num_forwards.zero_()
+
+
+class AdaptiveHierarchicalINL(nn.Module):
+    """
+    Adaptive wrapper for HierarchicalEquilibriumINL with early stopping.
+
+    Specifically designed for INL blocks in language models.
+    Monitors velocity (rate of change) instead of error for convergence detection.
+    """
+
+    def __init__(
+        self,
+        inl_layer: nn.Module,
+        convergence_threshold: float = 0.001,
+        min_iterations: int = 3,
+        max_iterations: int = 12,
+        check_interval: int = 1
+    ):
+        """
+        Args:
+            inl_layer: HierarchicalEquilibriumINL to wrap
+            convergence_threshold: Velocity threshold for early stopping
+            min_iterations: Minimum iterations before checking convergence
+            max_iterations: Maximum iterations (used during training)
+            check_interval: Check convergence every N iterations
+        """
+        super().__init__()
+
+        self.inl = inl_layer
+        self.convergence_threshold = convergence_threshold
+        self.min_iterations = min_iterations
+        self.max_iterations = max_iterations
+        self.check_interval = check_interval
+
+        # Statistics tracking
+        self.register_buffer('avg_iterations', torch.tensor(0.0))
+        self.register_buffer('num_forwards', torch.tensor(0))
+
+    def forward(
+        self,
+        h: torch.Tensor,
+        x: torch.Tensor,
+        v: torch.Tensor,
+        step: int = 0
+    ):
+        """
+        Forward pass - compatible with INL block usage.
+
+        Note: For use in INL blocks, this is called per-iteration.
+        Early stopping is handled at the block level.
+        """
+        return self.inl(h, x, v, step)
+
+    def forward_adaptive(
+        self,
+        h: torch.Tensor,
+        initial_x: torch.Tensor,
+        initial_v: torch.Tensor,
+        num_iterations: Optional[int] = None,
+        use_early_stopping: bool = None,
+        return_trajectory: bool = False
+    ):
+        """
+        Full adaptive forward with early stopping control.
+
+        Use this method when you want full control over iterations.
+        """
+        batch_size = h.shape[0]
+        device = h.device
+
+        if num_iterations is None:
+            num_iterations = self.max_iterations
+
+        if use_early_stopping is None:
+            use_early_stopping = not self.training
+
+        x, v = initial_x, initial_v
+
+        # Track trajectory if needed
+        x_traj = [x.clone()] if return_trajectory else None
+        v_traj = [v.clone()] if return_trajectory else None
+
+        converged = torch.zeros(batch_size, dtype=torch.bool, device=device)
+        iterations_used = torch.zeros(batch_size, dtype=torch.long, device=device)
+
+        for t in range(num_iterations):
+            x_prev = x.clone()
+
+            # Run integration step
+            x_next, v_next, aux = self.inl(h, x, v, step=t)
+
+            # Update iterations counter for non-converged samples
+            iterations_used[~converged] += 1
+
+            # Check convergence based on velocity (rate of change)
+            if use_early_stopping and t >= self.min_iterations and t % self.check_interval == 0:
+                # Compute change in state
+                delta_x = torch.norm(x_next - x_prev, dim=-1)  # [batch_size]
+
+                # Mark newly converged samples
+                newly_converged = (delta_x < self.convergence_threshold) & (~converged)
+                converged = converged | newly_converged
+
+                # If all samples converged, stop early
+                if converged.all():
+                    x, v = x_next, v_next
+                    if return_trajectory:
+                        x_traj.append(x.clone())
+                        v_traj.append(v.clone())
+                    break
+
+            x, v = x_next, v_next
+
+            if return_trajectory:
+                x_traj.append(x.clone())
+                v_traj.append(v.clone())
+
+        # Update statistics (exponential moving average)
+        if not self.training:
+            avg_iters = iterations_used.float().mean()
+            self.num_forwards += 1
+            alpha = 0.99
+            self.avg_iterations = alpha * self.avg_iterations + (1 - alpha) * avg_iters
+
+        result = {
+            'x': x,
+            'v': v,
+            'iterations_used': iterations_used,
+            'converged': converged,
+            'avg_iterations': self.avg_iterations.item(),
+            'mu': aux.get('mu'),
+            'mu_global': aux.get('mu_global'),
+            'mu_offsets': aux.get('mu_offsets')
+        }
+
+        if return_trajectory:
+            result['x_trajectory'] = torch.stack(x_traj, dim=1)  # [B, T+1, D]
+            result['v_trajectory'] = torch.stack(v_traj, dim=1)
+
+        return x, v, result
+
+    def reset_statistics(self):
+        """Reset tracking statistics."""
+        self.avg_iterations.zero_()
+        self.num_forwards.zero_()
+
+
+class GradientCheckpointedINL(nn.Module):
+    """
+    Wrapper for IntegratorNeuronLayer with gradient checkpointing.
+
+    Trades compute for memory:
+    - Forward: Normal computation
+    - Backward: Recompute forward instead of storing activations
+
+    Memory savings: 50-70% during training
+    Cost: ~30% slower backward pass (but worth it for large models!)
+    """
+
+    def __init__(self, inl_layer: nn.Module):
+        """
+        Args:
+            inl_layer: IntegratorNeuronLayer to wrap
+        """
+        super().__init__()
+        self.inl = inl_layer
+
+    def forward(
+        self,
+        h: torch.Tensor,
+        x: torch.Tensor,
+        v: torch.Tensor,
+        step: int = 0,
+        return_aux: bool = True
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Dict]]:
+        """
+        Forward with gradient checkpointing.
+
+        Uses torch.utils.checkpoint to save memory during backward pass.
+        """
+        if self.training:
+            # Use checkpointing during training
+            return torch.utils.checkpoint.checkpoint(
+                self._forward_impl,
+                h, x, v, step, return_aux,
+                use_reentrant=False
+            )
+        else:
+            # No checkpointing during inference
+            return self._forward_impl(h, x, v, step, return_aux)
+
+    def _forward_impl(
+        self,
+        h: torch.Tensor,
+        x: torch.Tensor,
+        v: torch.Tensor,
+        step: int,
+        return_aux: bool
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Dict]]:
+        """Actual forward implementation."""
+        return self.inl(h, x, v, step, return_aux)
+
+    def init_state(self, batch_size: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Delegate to wrapped layer."""
+        return self.inl.init_state(batch_size, device)
+
+    def __getattr__(self, name: str):
+        """Delegate attribute access to wrapped layer."""
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self.inl, name)
+
+
+def compute_parameter_reduction(
+    vocab_size: int,
+    d_model: int,
+    rank_ratio: float = 0.125
+) -> Dict[str, float]:
+    """
+    Compute parameter reduction from using low-rank embeddings.
+
+    Args:
+        vocab_size: Vocabulary size
+        d_model: Model dimension
+        rank_ratio: Rank ratio for low-rank embedding
+
+    Returns:
+        Dictionary with parameter counts and reduction percentage
+    """
+    rank = max(64, int(d_model * rank_ratio))
+
+    standard_params = vocab_size * d_model
+    lowrank_params = vocab_size * rank + rank * d_model
+
+    reduction_pct = (1 - lowrank_params / standard_params) * 100
+
+    return {
+        'standard_params': standard_params,
+        'lowrank_params': lowrank_params,
+        'reduction_percent': reduction_pct,
+        'rank': rank,
+        'memory_mb_standard': standard_params * 4 / 1e6,  # FP32
+        'memory_mb_lowrank': lowrank_params * 4 / 1e6
+    }
+
+
+def print_optimization_summary(
+    vocab_size: int,
+    d_model: int,
+    num_layers: int,
+    rank_ratio: float = 0.125
+):
+    """
+    Print summary of optimization benefits.
+
+    Args:
+        vocab_size: Vocabulary size
+        d_model: Model dimension
+        num_layers: Number of layers
+        rank_ratio: Low-rank embedding ratio
+    """
+    print("=" * 70)
+    print("INL-LLM OPTIMIZATION SUMMARY")
+    print("=" * 70)
+
+    # Low-rank embedding savings
+    embed_stats = compute_parameter_reduction(vocab_size, d_model, rank_ratio)
+
+    print("\n1. LOW-RANK EMBEDDINGS")
+    print("-" * 70)
+    print(f"  Standard embedding:  {embed_stats['standard_params']:>12,} params "
+          f"({embed_stats['memory_mb_standard']:>6.1f} MB)")
+    print(f"  Low-rank embedding:  {embed_stats['lowrank_params']:>12,} params "
+          f"({embed_stats['memory_mb_lowrank']:>6.1f} MB)")
+    print(f"  Rank: {embed_stats['rank']}")
+    print(f"  💾 REDUCTION: {embed_stats['reduction_percent']:.1f}%")
+
+    print("\n2. ADAPTIVE EARLY STOPPING")
+    print("-" * 70)
+    print("  Training:   Uses max iterations (no change)")
+    print("  Inference:  Adaptive iterations based on convergence")
+    print("  ⚡ SPEEDUP: 30-50% faster inference")
+    print("  Typical iterations: 5-7 (vs 10 max)")
+
+    print("\n3. GRADIENT CHECKPOINTING")
+    print("-" * 70)
+    print("  Memory reduction: ~50-70% during training")
+    print("  Compute overhead: ~30% slower backward")
+    print("  Enables scaling to: 2-3x larger models")
+    print("  🚀 BENEFIT: Train 100B+ models on consumer GPUs")
+
+    print("\n4. COMBINED IMPACT")
+    print("-" * 70)
+    saved_params = embed_stats['standard_params'] - embed_stats['lowrank_params']
+    print(f"  Total parameters saved: {saved_params:,}")
+    print(f"  Memory saved (embeddings): {embed_stats['memory_mb_standard'] - embed_stats['memory_mb_lowrank']:.1f} MB")
+    print(f"  Inference speedup: 30-50%")
+    print(f"  Training memory: -50-70%")
+
+    print("\n" + "=" * 70)
+    print("✅ OPTIMIZATIONS READY TO USE")
+    print("=" * 70)
+
+
+if __name__ == '__main__':
+    print("\n")
+    print_optimization_summary(
+        vocab_size=50000,
+        d_model=2048,
+        num_layers=24,
+        rank_ratio=0.125
+    )
+
+    print("\n\nEXAMPLE USAGE:\n")
+    print("# 1. Low-Rank Embeddings")
+    print("from optimizations import LowRankEmbedding")
+    print("embed = LowRankEmbedding(vocab_size=50000, d_model=2048, rank_ratio=0.125)")
+    print()
+
+    print("# 2. Adaptive Early Stopping")
+    print("from optimizations import AdaptiveIntegratorNeuronLayer")
+    print("adaptive_inl = AdaptiveIntegratorNeuronLayer(")
+    print("    inl_layer=base_inl,")
+    print("    convergence_threshold=0.01,")
+    print("    max_iterations=10")
+    print(")")
+    print()
+
+    print("# 3. Gradient Checkpointing")
+    print("from optimizations import GradientCheckpointedINL")
+    print("checkpointed_inl = GradientCheckpointedINL(base_inl)")
+    print()

pretraining_data_pipeline.py

@@ -0,0 +1,625 @@
+"""
+Pipeline de données pour le pré-entraînement du modèle INL-LLM.
+
+Ce module fournit des outils flexibles pour charger et prétraiter des données
+depuis diverses sources (fichiers locaux, datasets HuggingFace, etc.) pour
+le pré-entraînement de modèles de langage.
+
+Fonctionnalités:
+- Support multi-formats: parquet, jsonl, txt, csv
+- Streaming de larges datasets
+- Tokenization en batch avec multiprocessing
+- Filtrage et nettoyage des données
+- Mélange intelligent de plusieurs sources
+- Cache pour accélérer les chargements répétés
+"""
+
+import os
+import json
+import glob
+from pathlib import Path
+from typing import List, Dict, Optional, Union, Callable, Iterator
+from dataclasses import dataclass, field
+import logging
+
+import torch
+from torch.utils.data import Dataset, IterableDataset, DataLoader
+import pandas as pd
+from transformers import PreTrainedTokenizer
+
+# Configuration du logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class DataSourceConfig:
+    """Configuration pour une source de données."""
+
+    source_type: str  # 'parquet', 'jsonl', 'txt', 'csv', 'huggingface'
+    path: str  # Chemin du fichier ou nom du dataset HF
+    text_column: str = 'text'  # Nom de la colonne contenant le texte
+    weight: float = 1.0  # Poids pour le mélange de sources
+    streaming: bool = False  # Mode streaming pour économiser la mémoire
+    split: str = 'train'  # Split du dataset (pour HuggingFace)
+    config_name: Optional[str] = None  # Nom de la config pour les datasets HF (ex: 'wikitext-103-v1', 'en')
+    max_samples: Optional[int] = None  # Limite du nombre d'échantillons
+    filters: Dict = field(default_factory=dict)  # Filtres à appliquer
+
+
+@dataclass
+class PreprocessConfig:
+    """Configuration pour le prétraitement des données."""
+
+    min_length: int = 10  # Longueur minimale en tokens
+    max_length: int = 2048  # Longueur maximale en tokens
+    seq_length: int = 512  # Longueur de séquence cible
+    remove_duplicates: bool = True  # Supprimer les doublons
+    lowercase: bool = False  # Convertir en minuscules
+    remove_urls: bool = True  # Supprimer les URLs
+    remove_special_chars: bool = False  # Supprimer les caractères spéciaux
+    custom_filters: List[Callable] = field(default_factory=list)  # Filtres personnalisés
+
+
+class TextCleaner:
+    """Utilitaires pour nettoyer et filtrer le texte."""
+
+    @staticmethod
+    def remove_urls(text: str) -> str:
+        """Supprime les URLs du texte."""
+        import re
+        url_pattern = re.compile(r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+')
+        return url_pattern.sub('', text)
+
+    @staticmethod
+    def remove_special_chars(text: str) -> str:
+        """Supprime les caractères spéciaux non-alphanumériques."""
+        import re
+        return re.sub(r'[^a-zA-Z0-9\s\.,!?;:\-\'\"()]', '', text)
+
+    @staticmethod
+    def remove_extra_whitespace(text: str) -> str:
+        """Normalise les espaces blancs."""
+        return ' '.join(text.split())
+
+    @staticmethod
+    def is_valid_text(text: str, min_words: int = 5) -> bool:
+        """Vérifie si le texte est valide (contient suffisamment de mots)."""
+        if not text or not isinstance(text, str):
+            return False
+        words = text.split()
+        return len(words) >= min_words
+
+
+class PretrainingDataset(Dataset):
+    """
+    Dataset pour le pré-entraînement avec support multi-sources.
+
+    Ce dataset charge les données depuis plusieurs sources, les prétraite,
+    et retourne des paires (input_ids, labels) pour l'entraînement.
+    """
+
+    def __init__(
+        self,
+        sources: List[DataSourceConfig],
+        tokenizer: PreTrainedTokenizer,
+        preprocess_config: PreprocessConfig,
+        cache_dir: Optional[str] = None,
+        num_workers: int = 4
+    ):
+        """
+        Args:
+            sources: Liste de configurations de sources de données
+            tokenizer: Tokenizer HuggingFace à utiliser
+            preprocess_config: Configuration du prétraitement
+            cache_dir: Répertoire pour cacher les données prétraitées
+            num_workers: Nombre de workers pour le traitement parallèle
+        """
+        self.sources = sources
+        self.tokenizer = tokenizer
+        self.config = preprocess_config
+        self.cache_dir = cache_dir
+        self.num_workers = num_workers
+
+        self.text_cleaner = TextCleaner()
+        self.samples = []
+        self.tokenized_samples = []
+
+        logger.info(f"Initialisation du PretrainingDataset avec {len(sources)} sources")
+        self._load_all_sources()
+        self._preprocess_samples()
+
+    def _load_all_sources(self):
+        """Charge toutes les sources de données configurées."""
+        for idx, source in enumerate(self.sources):
+            logger.info(f"Chargement de la source {idx + 1}/{len(self.sources)}: {source.path}")
+
+            if source.source_type == 'parquet':
+                samples = self._load_parquet(source)
+            elif source.source_type == 'jsonl':
+                samples = self._load_jsonl(source)
+            elif source.source_type == 'txt':
+                samples = self._load_txt(source)
+            elif source.source_type == 'csv':
+                samples = self._load_csv(source)
+            elif source.source_type == 'huggingface':
+                samples = self._load_huggingface(source)
+            else:
+                logger.warning(f"Type de source inconnu: {source.source_type}")
+                continue
+
+            # Appliquer le poids de la source en dupliquant les échantillons
+            if source.weight != 1.0:
+                repeat_count = int(source.weight)
+                samples = samples * repeat_count
+                logger.info(f"Source pondérée avec facteur {source.weight}")
+
+            self.samples.extend(samples)
+            logger.info(f"Chargé {len(samples)} échantillons depuis {source.path}")
+
+        logger.info(f"Total: {len(self.samples)} échantillons bruts chargés")
+
+    def _load_parquet(self, source: DataSourceConfig) -> List[str]:
+        """Charge des données depuis un ou plusieurs fichiers Parquet."""
+        samples = []
+
+        # Support pour les patterns glob (e.g., "data/*.parquet")
+        if '*' in source.path:
+            files = glob.glob(source.path)
+        else:
+            files = [source.path]
+
+        for file_path in files:
+            if not os.path.exists(file_path):
+                logger.warning(f"Fichier non trouvé: {file_path}")
+                continue
+
+            df = pd.read_parquet(file_path)
+
+            if source.text_column not in df.columns:
+                logger.error(f"Colonne '{source.text_column}' non trouvée dans {file_path}")
+                continue
+
+            texts = df[source.text_column].dropna().tolist()
+
+            if source.max_samples:
+                texts = texts[:source.max_samples]
+
+            samples.extend(texts)
+
+        return samples
+
+    def _load_jsonl(self, source: DataSourceConfig) -> List[str]:
+        """Charge des données depuis un fichier JSONL."""
+        samples = []
+
+        if not os.path.exists(source.path):
+            logger.warning(f"Fichier non trouvé: {source.path}")
+            return samples
+
+        with open(source.path, 'r', encoding='utf-8') as f:
+            for idx, line in enumerate(f):
+                if source.max_samples and idx >= source.max_samples:
+                    break
+
+                try:
+                    data = json.loads(line)
+                    if source.text_column in data:
+                        samples.append(data[source.text_column])
+                except json.JSONDecodeError:
+                    logger.warning(f"Ligne JSON invalide à la ligne {idx + 1}")
+                    continue
+
+        return samples
+
+    def _load_txt(self, source: DataSourceConfig) -> List[str]:
+        """Charge des données depuis un fichier texte."""
+        samples = []
+
+        if not os.path.exists(source.path):
+            logger.warning(f"Fichier non trouvé: {source.path}")
+            return samples
+
+        with open(source.path, 'r', encoding='utf-8') as f:
+            content = f.read()
+
+        # Diviser par paragraphes (double saut de ligne)
+        paragraphs = content.split('\n\n')
+        samples = [p.strip() for p in paragraphs if p.strip()]
+
+        if source.max_samples:
+            samples = samples[:source.max_samples]
+
+        return samples
+
+    def _load_csv(self, source: DataSourceConfig) -> List[str]:
+        """Charge des données depuis un fichier CSV."""
+        samples = []
+
+        if not os.path.exists(source.path):
+            logger.warning(f"Fichier non trouvé: {source.path}")
+            return samples
+
+        df = pd.read_csv(source.path)
+
+        if source.text_column not in df.columns:
+            logger.error(f"Colonne '{source.text_column}' non trouvée dans {source.path}")
+            return samples
+
+        texts = df[source.text_column].dropna().tolist()
+
+        if source.max_samples:
+            texts = texts[:source.max_samples]
+
+        samples.extend(texts)
+
+        return samples
+
+    def _load_huggingface(self, source: DataSourceConfig) -> List[str]:
+        """Charge des données depuis un dataset HuggingFace."""
+        try:
+            from datasets import load_dataset
+        except ImportError:
+            logger.error("Le package 'datasets' n'est pas installé. Installez-le avec: pip install datasets")
+            return []
+
+        samples = []
+
+        try:
+            # Charger le dataset avec config_name si fourni
+            load_args = {
+                'path': source.path,
+                'split': source.split,
+                'streaming': source.streaming
+            }
+
+            # Ajouter config_name si présent
+            if source.config_name:
+                load_args['name'] = source.config_name
+                logger.info(f"Chargement avec config: {source.config_name}")
+
+            # Pour C4, ajouter trust_remote_code
+            if 'c4' in source.path.lower():
+                load_args['trust_remote_code'] = True
+
+            dataset = load_dataset(**load_args)
+
+            # Extraire les textes
+            if source.streaming:
+                # Mode streaming: itérer avec limite
+                for idx, example in enumerate(dataset):
+                    if source.max_samples and idx >= source.max_samples:
+                        break
+                    if source.text_column in example:
+                        samples.append(example[source.text_column])
+            else:
+                # Mode non-streaming: charger tout en mémoire
+                if source.text_column in dataset.column_names:
+                    texts = dataset[source.text_column]
+                    if source.max_samples:
+                        texts = texts[:source.max_samples]
+                    samples.extend(texts)
+
+            logger.info(f"Dataset HuggingFace chargé: {source.path} ({len(samples)} échantillons)")
+
+        except Exception as e:
+            logger.error(f"Erreur lors du chargement du dataset HuggingFace {source.path}: {e}")
+
+        return samples
+
+    def _preprocess_samples(self):
+        """Prétraite et tokenise tous les échantillons."""
+        logger.info("Prétraitement des échantillons...")
+
+        # Nettoyage du texte
+        cleaned_samples = []
+        for text in self.samples:
+            if not self.text_cleaner.is_valid_text(text):
+                continue
+
+            # Appliquer les filtres de nettoyage
+            if self.config.lowercase:
+                text = text.lower()
+
+            if self.config.remove_urls:
+                text = self.text_cleaner.remove_urls(text)
+
+            if self.config.remove_special_chars:
+                text = self.text_cleaner.remove_special_chars(text)
+
+            text = self.text_cleaner.remove_extra_whitespace(text)
+
+            # Appliquer les filtres personnalisés
+            for custom_filter in self.config.custom_filters:
+                text = custom_filter(text)
+
+            cleaned_samples.append(text)
+
+        logger.info(f"Échantillons après nettoyage: {len(cleaned_samples)}")
+
+        # Supprimer les doublons si demandé
+        if self.config.remove_duplicates:
+            initial_count = len(cleaned_samples)
+            cleaned_samples = list(set(cleaned_samples))
+            logger.info(f"Doublons supprimés: {initial_count - len(cleaned_samples)}")
+
+        # Tokenisation
+        logger.info("Tokenisation des échantillons...")
+        for idx, text in enumerate(cleaned_samples):
+            if idx % 1000 == 0:
+                logger.info(f"Tokenisation: {idx}/{len(cleaned_samples)}")
+
+            # Tokeniser le texte
+            encoded = self.tokenizer.encode(text, add_special_tokens=True)
+
+            # Filtrer par longueur
+            if len(encoded) < self.config.min_length:
+                continue
+
+            if len(encoded) > self.config.max_length:
+                encoded = encoded[:self.config.max_length]
+
+            # Découper en séquences de longueur fixe
+            for i in range(0, len(encoded) - self.config.seq_length, self.config.seq_length // 2):
+                chunk = encoded[i:i + self.config.seq_length + 1]
+                if len(chunk) == self.config.seq_length + 1:
+                    self.tokenized_samples.append(chunk)
+
+        logger.info(f"Dataset final: {len(self.tokenized_samples)} séquences tokenisées")
+
+    def __len__(self) -> int:
+        return len(self.tokenized_samples)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        """
+        Retourne un échantillon tokenisé.
+
+        Returns:
+            Dict contenant:
+                - input_ids: Tokens d'entrée [seq_len]
+                - labels: Tokens cibles (décalés) [seq_len]
+        """
+        tokens = self.tokenized_samples[idx]
+
+        input_ids = torch.tensor(tokens[:-1], dtype=torch.long)
+        labels = torch.tensor(tokens[1:], dtype=torch.long)
+
+        return {
+            'input_ids': input_ids,
+            'labels': labels
+        }
+
+
+class StreamingPretrainingDataset(IterableDataset):
+    """
+    Dataset iterable pour le streaming de très larges datasets.
+
+    Ce dataset ne charge pas toutes les données en mémoire, mais les
+    traite à la volée. Utile pour des datasets de plusieurs TB.
+    """
+
+    def __init__(
+        self,
+        sources: List[DataSourceConfig],
+        tokenizer: PreTrainedTokenizer,
+        preprocess_config: PreprocessConfig,
+        buffer_size: int = 10000
+    ):
+        """
+        Args:
+            sources: Liste de configurations de sources de données
+            tokenizer: Tokenizer HuggingFace à utiliser
+            preprocess_config: Configuration du prétraitement
+            buffer_size: Taille du buffer pour le mélange des sources
+        """
+        self.sources = sources
+        self.tokenizer = tokenizer
+        self.config = preprocess_config
+        self.buffer_size = buffer_size
+        self.text_cleaner = TextCleaner()
+
+        logger.info(f"Initialisation du StreamingPretrainingDataset avec {len(sources)} sources")
+
+    def _process_text(self, text: str) -> Optional[List[int]]:
+        """Traite un texte et retourne les tokens."""
+        # Vérifier la validité
+        if not self.text_cleaner.is_valid_text(text):
+            return None
+
+        # Nettoyage
+        if self.config.lowercase:
+            text = text.lower()
+
+        if self.config.remove_urls:
+            text = self.text_cleaner.remove_urls(text)
+
+        if self.config.remove_special_chars:
+            text = self.text_cleaner.remove_special_chars(text)
+
+        text = self.text_cleaner.remove_extra_whitespace(text)
+
+        # Tokenisation
+        encoded = self.tokenizer.encode(text, add_special_tokens=True)
+
+        # Filtrage par longueur
+        if len(encoded) < self.config.min_length or len(encoded) > self.config.max_length:
+            return None
+
+        return encoded
+
+    def _stream_source(self, source: DataSourceConfig) -> Iterator[str]:
+        """Génère un stream de textes depuis une source."""
+        if source.source_type == 'huggingface':
+            try:
+                from datasets import load_dataset
+
+                # Charger le dataset avec config_name si fourni
+                load_args = {
+                    'path': source.path,
+                    'split': source.split,
+                    'streaming': True
+                }
+
+                if source.config_name:
+                    load_args['name'] = source.config_name
+
+                # Pour C4, ajouter trust_remote_code
+                if 'c4' in source.path.lower():
+                    load_args['trust_remote_code'] = True
+
+                dataset = load_dataset(**load_args)
+
+                for idx, example in enumerate(dataset):
+                    if source.max_samples and idx >= source.max_samples:
+                        break
+                    if source.text_column in example:
+                        yield example[source.text_column]
+            except Exception as e:
+                logger.error(f"Erreur streaming HF: {e}")
+
+        elif source.source_type == 'jsonl':
+            with open(source.path, 'r', encoding='utf-8') as f:
+                for idx, line in enumerate(f):
+                    if source.max_samples and idx >= source.max_samples:
+                        break
+                    try:
+                        data = json.loads(line)
+                        if source.text_column in data:
+                            yield data[source.text_column]
+                    except json.JSONDecodeError:
+                        continue
+
+        # Autres formats peuvent être ajoutés ici
+
+    def __iter__(self) -> Iterator[Dict[str, torch.Tensor]]:
+        """Itère sur les échantillons de manière streaming."""
+        for source in self.sources:
+            logger.info(f"Streaming depuis: {source.path}")
+
+            for text in self._stream_source(source):
+                encoded = self._process_text(text)
+
+                if encoded is None:
+                    continue
+
+                # Découper en chunks
+                for i in range(0, len(encoded) - self.config.seq_length, self.config.seq_length // 2):
+                    chunk = encoded[i:i + self.config.seq_length + 1]
+                    if len(chunk) == self.config.seq_length + 1:
+                        input_ids = torch.tensor(chunk[:-1], dtype=torch.long)
+                        labels = torch.tensor(chunk[1:], dtype=torch.long)
+
+                        yield {
+                            'input_ids': input_ids,
+                            'labels': labels
+                        }
+
+
+def create_pretraining_dataloader(
+    sources: List[DataSourceConfig],
+    tokenizer: PreTrainedTokenizer,
+    preprocess_config: PreprocessConfig,
+    batch_size: int = 8,
+    streaming: bool = False,
+    num_workers: int = 4,
+    shuffle: bool = True,
+    **dataloader_kwargs
+) -> DataLoader:
+    """
+    Crée un DataLoader configuré pour le pré-entraînement.
+
+    Args:
+        sources: Liste de sources de données
+        tokenizer: Tokenizer à utiliser
+        preprocess_config: Configuration du prétraitement
+        batch_size: Taille des batchs
+        streaming: Utiliser le mode streaming (pour très larges datasets)
+        num_workers: Nombre de workers pour le chargement
+        shuffle: Mélanger les données
+        **dataloader_kwargs: Arguments additionnels pour DataLoader
+
+    Returns:
+        DataLoader configuré
+    """
+    if streaming:
+        dataset = StreamingPretrainingDataset(
+            sources=sources,
+            tokenizer=tokenizer,
+            preprocess_config=preprocess_config
+        )
+        shuffle = False  # Pas de shuffle pour les iterable datasets
+    else:
+        dataset = PretrainingDataset(
+            sources=sources,
+            tokenizer=tokenizer,
+            preprocess_config=preprocess_config,
+            num_workers=num_workers
+        )
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        pin_memory=True,
+        **dataloader_kwargs
+    )
+
+    logger.info(f"DataLoader créé: batch_size={batch_size}, streaming={streaming}")
+
+    return dataloader
+
+
+# Exemple d'utilisation
+if __name__ == "__main__":
+    from transformers import AutoTokenizer
+
+    # Charger le tokenizer
+    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+
+    # Configuration des sources de données
+    sources = [
+        DataSourceConfig(
+            source_type='parquet',
+            path='part_000000.parquet',
+            text_column='text',
+            weight=1.0
+        ),
+        # Exemple avec HuggingFace dataset
+        # DataSourceConfig(
+        #     source_type='huggingface',
+        #     path='openwebtext',
+        #     text_column='text',
+        #     weight=2.0,
+        #     streaming=True,
+        #     max_samples=100000
+        # ),
+    ]
+
+    # Configuration du prétraitement
+    preprocess_config = PreprocessConfig(
+        min_length=10,
+        max_length=2048,
+        seq_length=512,
+        remove_duplicates=True,
+        remove_urls=True
+    )
+
+    # Créer le dataloader
+    dataloader = create_pretraining_dataloader(
+        sources=sources,
+        tokenizer=tokenizer,
+        preprocess_config=preprocess_config,
+        batch_size=4,
+        streaming=False,
+        num_workers=2
+    )
+
+    # Test: charger quelques batchs
+    logger.info("Test du dataloader...")
+    for i, batch in enumerate(dataloader):
+        if i >= 3:
+            break
+        logger.info(f"Batch {i}: input_ids shape = {batch['input_ids'].shape}, labels shape = {batch['labels'].shape}")
+
+    logger.info("Test terminé avec succès!")

pretraining_pipeline_config.json

@@ -0,0 +1,37 @@
+{
+  "description": "Configuration rapide pour le pré-entraînement INL-LLM (Option C - entraînement efficace)",
+  "sources": [
+    {
+      "source_type": "huggingface",
+      "path": "wikitext",
+      "config_name": "wikitext-103-v1",
+      "text_column": "text",
+      "weight": 1.0,
+      "streaming": true,
+      "split": "train",
+      "max_samples": 1000,
+      "filters": {}
+    },
+    {
+      "source_type": "huggingface",
+      "path": "allenai/c4",
+      "config_name": "en",
+      "text_column": "text",
+      "weight": 1.0,
+      "streaming": true,
+      "split": "train",
+      "max_samples": 1000,
+      "filters": {}
+    }
+  ],
+  "preprocess_config": {
+    "min_length": 50,
+    "max_length": 2048,
+    "seq_length": 64,
+    "remove_duplicates": true,
+    "lowercase": false,
+    "remove_urls": true,
+    "remove_special_chars": false,
+    "custom_filters": []
+  }
+}

pretraining_pipeline_examples.json

@@ -0,0 +1,278 @@
+{
+  "examples": {
+    "simple_parquet": {
+      "description": "Configuration simple avec un seul fichier parquet",
+      "sources": [
+        {
+          "source_type": "parquet",
+          "path": "part_000000.parquet",
+          "text_column": "text",
+          "weight": 1.0,
+          "streaming": false,
+          "max_samples": null
+        }
+      ],
+      "preprocess_config": {
+        "min_length": 10,
+        "max_length": 2048,
+        "seq_length": 512,
+        "remove_duplicates": true,
+        "remove_urls": true
+      }
+    },
+
+    "multi_parquet": {
+      "description": "Plusieurs fichiers parquet avec glob pattern",
+      "sources": [
+        {
+          "source_type": "parquet",
+          "path": "data/train/*.parquet",
+          "text_column": "text",
+          "weight": 1.0,
+          "streaming": false
+        }
+      ],
+      "preprocess_config": {
+        "min_length": 20,
+        "max_length": 2048,
+        "seq_length": 1024,
+        "remove_duplicates": true,
+        "remove_urls": true,
+        "remove_special_chars": false
+      }
+    },
+
+    "mixed_sources": {
+      "description": "Mélange de plusieurs sources de données avec pondération",
+      "sources": [
+        {
+          "source_type": "parquet",
+          "path": "data/wikipedia/wiki_*.parquet",
+          "text_column": "text",
+          "weight": 2.0,
+          "max_samples": 100000
+        },
+        {
+          "source_type": "jsonl",
+          "path": "data/books/books.jsonl",
+          "text_column": "content",
+          "weight": 1.5,
+          "max_samples": 50000
+        },
+        {
+          "source_type": "txt",
+          "path": "data/articles/articles.txt",
+          "text_column": "text",
+          "weight": 1.0
+        }
+      ],
+      "preprocess_config": {
+        "min_length": 50,
+        "max_length": 4096,
+        "seq_length": 2048,
+        "remove_duplicates": true,
+        "lowercase": false,
+        "remove_urls": true,
+        "remove_special_chars": false
+      }
+    },
+
+    "huggingface_dataset": {
+      "description": "Utilisation d'un dataset HuggingFace avec streaming",
+      "sources": [
+        {
+          "source_type": "huggingface",
+          "path": "openwebtext",
+          "text_column": "text",
+          "weight": 1.0,
+          "streaming": true,
+          "split": "train",
+          "max_samples": 500000
+        }
+      ],
+      "preprocess_config": {
+        "min_length": 100,
+        "max_length": 4096,
+        "seq_length": 2048,
+        "remove_duplicates": false,
+        "remove_urls": true,
+        "remove_special_chars": false
+      }
+    },
+
+    "multi_domain": {
+      "description": "Pré-entraînement multi-domaine avec différents datasets",
+      "sources": [
+        {
+          "source_type": "huggingface",
+          "path": "wikipedia",
+          "text_column": "text",
+          "weight": 3.0,
+          "streaming": true,
+          "split": "train",
+          "max_samples": 1000000
+        },
+        {
+          "source_type": "huggingface",
+          "path": "bookcorpus",
+          "text_column": "text",
+          "weight": 2.0,
+          "streaming": true,
+          "split": "train",
+          "max_samples": 500000
+        },
+        {
+          "source_type": "parquet",
+          "path": "data/code/github_*.parquet",
+          "text_column": "content",
+          "weight": 1.0,
+          "max_samples": 200000
+        },
+        {
+          "source_type": "jsonl",
+          "path": "data/conversations/dialogues.jsonl",
+          "text_column": "text",
+          "weight": 1.5,
+          "max_samples": 100000
+        }
+      ],
+      "preprocess_config": {
+        "min_length": 50,
+        "max_length": 4096,
+        "seq_length": 2048,
+        "remove_duplicates": true,
+        "lowercase": false,
+        "remove_urls": true,
+        "remove_special_chars": false
+      }
+    },
+
+    "high_quality": {
+      "description": "Configuration pour données de haute qualité (moins de filtrage)",
+      "sources": [
+        {
+          "source_type": "parquet",
+          "path": "data/curated/high_quality_*.parquet",
+          "text_column": "text",
+          "weight": 1.0
+        }
+      ],
+      "preprocess_config": {
+        "min_length": 100,
+        "max_length": 8192,
+        "seq_length": 4096,
+        "remove_duplicates": false,
+        "lowercase": false,
+        "remove_urls": false,
+        "remove_special_chars": false
+      }
+    },
+
+    "aggressive_cleaning": {
+      "description": "Nettoyage agressif pour données brutes du web",
+      "sources": [
+        {
+          "source_type": "parquet",
+          "path": "data/web_crawl/*.parquet",
+          "text_column": "text",
+          "weight": 1.0
+        }
+      ],
+      "preprocess_config": {
+        "min_length": 200,
+        "max_length": 2048,
+        "seq_length": 1024,
+        "remove_duplicates": true,
+        "lowercase": false,
+        "remove_urls": true,
+        "remove_special_chars": true
+      }
+    }
+  },
+
+  "usage": {
+    "description": "Comment utiliser ces configurations",
+    "examples": [
+      {
+        "command": "python simple_training.py",
+        "description": "Mode simple (legacy) sans pipeline"
+      },
+      {
+        "command": "python simple_training.py --use-pipeline",
+        "description": "Utilise le pipeline avec configuration par défaut"
+      },
+      {
+        "command": "python simple_training.py --use-pipeline --config pretraining_pipeline_config.json",
+        "description": "Utilise le pipeline avec une configuration personnalisée"
+      }
+    ],
+    "custom_config": {
+      "description": "Pour créer votre propre configuration, copiez l'un des exemples ci-dessus et modifiez les paramètres selon vos besoins",
+      "steps": [
+        "1. Créez un nouveau fichier JSON (ex: my_config.json)",
+        "2. Copiez la structure d'un exemple ci-dessus",
+        "3. Modifiez les chemins de fichiers et paramètres",
+        "4. Lancez: python simple_training.py --use-pipeline --config my_config.json"
+      ]
+    }
+  },
+
+  "source_types": {
+    "parquet": {
+      "description": "Fichiers Apache Parquet",
+      "supports_glob": true,
+      "example": "data/*.parquet ou data/train/part_*.parquet"
+    },
+    "jsonl": {
+      "description": "Fichiers JSON Lines (un JSON par ligne)",
+      "supports_glob": false,
+      "example": "data/train.jsonl"
+    },
+    "txt": {
+      "description": "Fichiers texte brut (divisés par paragraphes)",
+      "supports_glob": false,
+      "example": "data/book.txt"
+    },
+    "csv": {
+      "description": "Fichiers CSV",
+      "supports_glob": false,
+      "example": "data/dataset.csv"
+    },
+    "huggingface": {
+      "description": "Datasets depuis HuggingFace Hub",
+      "supports_streaming": true,
+      "example": "openwebtext, wikipedia, bookcorpus, etc."
+    }
+  },
+
+  "preprocess_parameters": {
+    "min_length": "Longueur minimale en tokens (les textes plus courts sont ignorés)",
+    "max_length": "Longueur maximale en tokens (les textes plus longs sont tronqués)",
+    "seq_length": "Longueur de séquence cible pour l'entraînement",
+    "remove_duplicates": "Supprimer les textes dupliqués (true/false)",
+    "lowercase": "Convertir tout en minuscules (true/false)",
+    "remove_urls": "Supprimer les URLs (true/false)",
+    "remove_special_chars": "Supprimer les caractères spéciaux non-alphanumériques (true/false)"
+  },
+
+  "source_parameters": {
+    "source_type": "Type de source (parquet, jsonl, txt, csv, huggingface)",
+    "path": "Chemin du fichier ou nom du dataset HF",
+    "text_column": "Nom de la colonne contenant le texte",
+    "weight": "Poids pour le mélange de sources (1.0 = normal, 2.0 = doublement, etc.)",
+    "streaming": "Mode streaming pour économiser la mémoire (true/false)",
+    "split": "Split du dataset pour HuggingFace (train, validation, test)",
+    "max_samples": "Nombre maximum d'échantillons à charger (null = tous)",
+    "filters": "Filtres additionnels (avancé, généralement vide)"
+  },
+
+  "tips": [
+    "Utilisez le streaming (streaming: true) pour les très larges datasets qui ne tiennent pas en mémoire",
+    "Ajustez les poids (weight) pour contrôler la proportion de chaque source dans le mélange",
+    "Augmentez seq_length pour capturer plus de contexte (mais cela augmente la mémoire GPU)",
+    "Utilisez remove_duplicates: true pour améliorer la diversité des données",
+    "Pour du code source, gardez remove_special_chars: false",
+    "Pour du texte web brut, activez remove_urls: true et remove_special_chars: true",
+    "max_samples est utile pour des tests rapides avec un sous-ensemble des données"
+  ]
+}

simple_training.py

@@ -6,28 +6,47 @@ This script demonstrates basic training with:
 - Real text data from parquet file (785 samples)
 - Equilibrium-exploration cycles
 - Adaptive early stopping (3× faster inference)
+- Advanced pretraining data pipeline with multi-source support
 
 Dataset: part_000000.parquet
 - Contains 785 real text samples
 - Tokenized using GPT-2 BPE tokenizer
 - Sequence length: 64 tokens
+
+New: Use --use-pipeline flag to use the advanced pretraining pipeline
 """
 
 import sys
 import os
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
+# Disable tokenizers parallelism warning when using multiprocessing
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
 import torch
 import torch.nn as nn
 from torch.utils.data import Dataset, DataLoader
 import pandas as pd
 import json
+import argparse
 
 # Import from the correct path
 from inl_llm.models.integrator_language_model import UltraOptimizedIntegratorLanguageModel
 from inl_llm.core.integrator_losses import IntegratorLoss
 from inl_llm.core.integrator_scheduler_v2 import create_cycle_scheduler
 
+# Import the new pretraining pipeline
+try:
+    from pretraining_data_pipeline import (
+        DataSourceConfig,
+        PreprocessConfig,
+        create_pretraining_dataloader
+    )
+    PIPELINE_AVAILABLE = True
+except ImportError:
+    PIPELINE_AVAILABLE = False
+    print("⚠️ pretraining_data_pipeline not found. Advanced pipeline features disabled.")
+
 # Import tokenizer
 try:
     from transformers import AutoTokenizer
@@ -94,6 +113,14 @@ def train_epoch(model, dataloader, loss_fn, optimizer, scheduler, device='cpu',
     total_loss = 0
     num_batches = 0
 
+    # Check if dataloader has length (for streaming datasets, it doesn't)
+    try:
+        total_batches = len(dataloader)
+        show_total = True
+    except TypeError:
+        total_batches = "?"
+        show_total = False
+
     for batch_idx, (inputs, targets) in enumerate(dataloader):
         inputs, targets = inputs.to(device), targets.to(device)
 
@@ -119,19 +146,57 @@ def train_epoch(model, dataloader, loss_fn, optimizer, scheduler, device='cpu',
             )
             loss = loss_components['total']
 
-            # Log detailed loss components
+            # Log detailed loss components + convergence metrics
             if batch_idx % 10 == 0:
                 L_task = loss_components.get('L_task', torch.tensor(0.0)).item()
                 L_mean = loss_components.get('L_mean', torch.tensor(0.0)).item()
                 L_speed = loss_components.get('L_speed', torch.tensor(0.0)).item()
                 L_energy = loss_components.get('L_energy', torch.tensor(0.0)).item()
-                print(f'  Batch {batch_idx}/{len(dataloader)}, Loss: {loss.item():.4f} '
-                      f'[Task: {L_task:.4f}, Mean: {L_mean:.4f}, Speed: {L_speed:.4f}, Energy: {L_energy:.4f}]')
+
+                # CONVERGENCE METRICS: Vérifier le théorème
+                if last_layer_traj is not None:
+                    # Trajectoires x, v de la dernière layer
+                    x_traj = last_layer_traj.get('x')  # [batch*seq, iterations, dim]
+                    v_traj = last_layer_traj.get('v')
+                    mu = last_layer_traj.get('mu')  # Équilibre cible
+
+                    if x_traj is not None and v_traj is not None:
+                        # Convergence = ||x_final - mu|| doit être petit
+                        x_final = x_traj[:, -1, :]  # Dernier état
+                        if mu is not None:
+                            error_norm = torch.norm(x_final - mu, dim=-1).mean().item()
+                        else:
+                            error_norm = 0.0
+
+                        # Stabilité vélocité
+                        v_init = v_traj[:, 0, :]   # Vélocité initiale (v_0)
+                        v_final = v_traj[:, -1, :] # Vélocité finale (v_T)
+
+                        # Métriques de convergence vélocité
+                        v_init_norm = torch.norm(v_init, dim=-1).mean().item()
+                        v_final_norm = torch.norm(v_final, dim=-1).mean().item()
+                        delta_v = torch.norm(v_final - v_init, dim=-1).mean().item()
+
+                        # Convergence vitesse = Δv devrait diminuer si système se stabilise
+                        # (même si v_target ≠ 0, la variation Δv diminue quand converge)
+
+                        # Nombre d'itérations utilisées (si adaptive stopping)
+                        iters_used = last_layer_traj.get('avg_iterations', 'N/A')
+
+                        print(f'  Batch {batch_idx}/{total_batches}, Loss: {loss.item():.4f} '
+                              f'[Task: {L_task:.4f}, Mean: {L_mean:.4f}, Speed: {L_speed:.4f}, Energy: {L_energy:.4f}]')
+                        print(f'    CONVERGENCE: ||x-μ||={error_norm:.4f}, ||v_0||={v_init_norm:.4f}, ||v_T||={v_final_norm:.4f}, Δv={delta_v:.4f}')
+                    else:
+                        print(f'  Batch {batch_idx}/{total_batches}, Loss: {loss.item():.4f} '
+                              f'[Task: {L_task:.4f}, Mean: {L_mean:.4f}, Speed: {L_speed:.4f}, Energy: {L_energy:.4f}]')
+                else:
+                    print(f'  Batch {batch_idx}/{total_batches}, Loss: {loss.item():.4f} '
+                          f'[Task: {L_task:.4f}, Mean: {L_mean:.4f}, Speed: {L_speed:.4f}, Energy: {L_energy:.4f}]')
         else:
             # Fallback to simple CrossEntropy
             loss = nn.CrossEntropyLoss()(logits_flat, targets_flat)
             if batch_idx % 10 == 0:
-                print(f'  Batch {batch_idx}/{len(dataloader)}, Loss: {loss.item():.4f} [Fallback CE]')
+                print(f'  Batch {batch_idx}/{total_batches}, Loss: {loss.item():.4f} [Fallback CE]')
 
         # Backward
         optimizer.zero_grad()
@@ -149,9 +214,13 @@ def train_epoch(model, dataloader, loss_fn, optimizer, scheduler, device='cpu',
     return total_loss / num_batches
 
 
-def main():
+def main(use_pipeline=False, pipeline_config=None):
     print("="*70)
     print("SIMPLE TRAINING EXAMPLE - INL-LLM")
+    if use_pipeline:
+        print("MODE: Advanced Pretraining Pipeline")
+    else:
+        print("MODE: Simple Dataset (legacy)")
     print("="*70)
 
     # Load tokenizer (GPT-2 BPE tokenizer, same as used by many LLMs)
@@ -160,6 +229,8 @@ def main():
         try:
             tokenizer = AutoTokenizer.from_pretrained("gpt2")
             tokenizer.pad_token = tokenizer.eos_token
+            # Increase model_max_length to match INL-LLM's capacity (2048)
+            tokenizer.model_max_length = 2048
 
             # Add special tokens for chat format
             special_tokens = {
@@ -195,7 +266,7 @@ def main():
 
     # Configuration
     batch_size = 2
-    num_epochs = 20  # ✅ Increased from 3 to 20 for better convergence
+    num_epochs = 3  # ✅ Increased from 3 to 20 for better convergence
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
     print(f"\nConfiguration:")
@@ -204,6 +275,7 @@ def main():
     print(f"  Batch size: {batch_size}")
     print(f"  Epochs: {num_epochs}")
     print(f"  Device: {device}")
+    print(f"  Using pipeline: {use_pipeline}")
 
     # Create custom 1.1B parameter model
     print("\nCreating custom 1.1B parameter model (all optimizations enabled)...")
@@ -214,16 +286,25 @@ def main():
         d_model=1728,           # Dimension du modèle (augmenté pour 1.1B)
         num_layers=25,          # Nombre de couches (augmenté pour 1.1B)
         num_heads=32,           # Nombre de têtes d'attention (1728/32 = 54 dim par tête)
-        num_iterations_per_layer=5,  # Itérations par couche
+        num_iterations_per_layer=5,  # Itérations par couche moyenne
         feedforward_dim=6912,   # Dimension FFN (4x d_model)
         max_seq_len=2048,
-        # Toutes les optimizations activées
+        # LEVEL 1 + 2: Optimizations de base
         use_lowrank_embeddings=True,
         lowrank_ratio=0.125,
         use_gradient_checkpointing=True,
         use_shared_controllers=True,
         hierarchical_group_size=64,
-        excitation_sparsity=0.1
+        excitation_sparsity=0.1,
+        # LEVEL 3: Adaptive Budget Allocator
+        use_adaptive_budget=True,
+        budget_strategy='hybrid',  # Learnable + dynamic allocation
+        budget_convergence_threshold=0.001,
+        # LEVEL 4: Mixture of Experts (MoE)
+        use_moe=True,
+        num_experts=4,  # 4 specialized expert controllers
+        moe_top_k=2,    # Activate 2 experts per forward (sparse routing)
+        moe_load_balance_weight=0.01  # Load balancing to prevent expert collapse
     )
     model = model.to(device)
 
@@ -231,44 +312,125 @@ def main():
 
     # Create dataset and dataloader
     print("\nCreating dataset...")
-    parquet_path = os.path.join(os.path.dirname(__file__), 'part_000000.parquet')
 
-    if not os.path.exists(parquet_path):
-        raise FileNotFoundError(f"Dataset not found at {parquet_path}")
+    if use_pipeline and PIPELINE_AVAILABLE:
+        # Use the advanced pretraining pipeline
+        print("📦 Using advanced pretraining pipeline...")
+
+        # Default configuration if none provided
+        if pipeline_config is None:
+            parquet_path = os.path.join(os.path.dirname(__file__), 'part_000000.parquet')
+
+            sources = [
+                DataSourceConfig(
+                    source_type='parquet',
+                    path=parquet_path,
+                    text_column='text',
+                    weight=1.0,
+                    max_samples=None  # Use all samples
+                )
+            ]
 
-    dataset = ParquetTextDataset(
-        parquet_path=parquet_path,
-        seq_len=64,
-        tokenizer=tokenizer
-    )
-    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+            preprocess_config = PreprocessConfig(
+                min_length=10,
+                max_length=2048,
+                seq_length=64,
+                remove_duplicates=True,
+                remove_urls=True,
+                remove_special_chars=False
+            )
+        else:
+            sources = pipeline_config['sources']
+            preprocess_config = pipeline_config['preprocess_config']
+
+        # Collate function to convert dict format to tuple format expected by train_epoch
+        def collate_fn(batch):
+            """Convert list of dicts to tuple of batched tensors."""
+            if isinstance(batch[0], dict):
+                # Pipeline format: {'input_ids': ..., 'labels': ...}
+                input_ids = torch.stack([item['input_ids'] for item in batch])
+                labels = torch.stack([item['labels'] for item in batch])
+                return input_ids, labels
+            else:
+                # Legacy format: (input, target) tuples
+                return torch.utils.data.dataloader.default_collate(batch)
+
+        # Create dataloader with the pipeline
+        # Use streaming=True for large datasets to avoid loading everything in memory
+        dataloader = create_pretraining_dataloader(
+            sources=sources,
+            tokenizer=tokenizer,
+            preprocess_config=preprocess_config,
+            batch_size=batch_size,
+            streaming=True,  # Changed to True for better performance with large datasets
+            num_workers=2,
+            shuffle=True,
+            collate_fn=collate_fn
+        )
+
+        print("✅ Advanced pipeline dataloader created")
+    else:
+        # Use legacy simple dataset
+        if use_pipeline and not PIPELINE_AVAILABLE:
+            print("⚠️ Pipeline requested but not available. Falling back to simple dataset.")
+
+        parquet_path = os.path.join(os.path.dirname(__file__), 'part_000000.parquet')
 
-    # ✅ FIX #3: Lower learning rate for large model (was 3e-4, too high)
-    optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5, weight_decay=0.01)
+        if not os.path.exists(parquet_path):
+            raise FileNotFoundError(f"Dataset not found at {parquet_path}")
+
+        dataset = ParquetTextDataset(
+            parquet_path=parquet_path,
+            seq_len=64,
+            tokenizer=tokenizer
+        )
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+    # Optimizer: Increased learning rate for better convergence
+    # For 1B+ models in pretraining, 1e-4 to 3e-4 is standard (GPT-3 used 6e-5 to 1.2e-4)
+    learning_rate = 1e-4  # Increased from 5e-5 for faster convergence
+    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)
 
     # Add learning rate scheduler with warmup
     from torch.optim.lr_scheduler import OneCycleLR
-    total_steps = num_epochs * len(dataloader)
+
+    # Calculate total_steps (handle streaming datasets which don't have len())
+    if use_pipeline and PIPELINE_AVAILABLE:
+        # For streaming datasets, estimate steps based on actual config
+        # With 1000 samples per source, seq_length=128, overlap=50%
+        # Each sample ~2000 tokens avg → ~32 sequences per sample → ~32k sequences total
+        # But streaming processes on-the-fly, estimate conservatively
+        estimated_samples = 10000  # More realistic for 2k samples with chunking
+        steps_per_epoch = estimated_samples // batch_size
+        total_steps = num_epochs * steps_per_epoch
+        print(f"⚠️ Streaming mode: estimated {steps_per_epoch} steps per epoch")
+    else:
+        total_steps = num_epochs * len(dataloader)
+
     lr_scheduler = OneCycleLR(
         optimizer,
-        max_lr=5e-5,
+        max_lr=learning_rate,  # Use the same LR as optimizer
         total_steps=total_steps,
         pct_start=0.1,  # 10% warmup
         anneal_strategy='cos'
     )
-    print(f"✅ Optimizer: AdamW with lr=5e-5, warmup={int(0.1*total_steps)} steps")
+    print(f"✅ Optimizer: AdamW with lr={learning_rate}, warmup={int(0.1*total_steps)} steps")
 
-    # ✅ FIX #1: Create IntegratorLoss (was not being used at all)
+    # IntegratorLoss: BALANCED approach (compromis entre convergence et task)
+    # Strategy: Moderate regularization + annealing over time
+    # L_total = L_task + λ_mean*L_mean + λ_speed*L_speed + λ_energy*L_energy
+    # With annealing, lambdas decrease: λ(t) = λ_init * exp(-t/T)
     integrator_loss_fn = IntegratorLoss(
-        target_value=0.0,  # ✅ Use 0.0 for normalized hidden states (not 5.0!)
-        lambda_mean_init=0.1,  # ✅ Reduced weight (was 1.0, too high)
-        lambda_speed=0.01,  # ✅ Reduced (was 0.1)
-        lambda_energy=0.001,  # ✅ Reduced (was 0.01)
-        annealing_epochs=num_epochs,
+        target_value=0.0,  # Use 0.0 for normalized hidden states
+        lambda_mean_init=0.05,   # BALANCED: not too high (0.1), not too low (0.01)
+        lambda_speed=0.005,      # BALANCED: allows some speed variation
+        lambda_energy=0.0005,    # BALANCED: mild energy constraint
+        annealing_epochs=num_epochs,  # Ces poids vont diminuer progressivement
         variance_weighted=True,
-        task_loss_type='ce'  # ✅ Use CrossEntropy for language modeling (not MSE)
+        task_loss_type='ce'  # CrossEntropy for language modeling
     )
-    print(f"✅ Loss function: IntegratorLoss with CrossEntropy + trajectory regularization")
+    print(f"✅ Loss function: IntegratorLoss (balanced task+convergence, with annealing)")
+    print(f"   λ_mean={0.05}, λ_speed={0.005}, λ_energy={0.0005} (will decay over {num_epochs} epochs)")
 
     # Create scheduler - automatically adapts to num_epochs
     cycle_scheduler = create_cycle_scheduler(
@@ -448,4 +610,35 @@ def main():
 
 
 if __name__ == '__main__':
-    main()
+    parser = argparse.ArgumentParser(
+        description='Train INL-LLM with simple dataset or advanced pretraining pipeline'
+    )
+    parser.add_argument(
+        '--use-pipeline',
+        action='store_true',
+        help='Use the advanced pretraining data pipeline (default: False, uses simple dataset)'
+    )
+    parser.add_argument(
+        '--config',
+        type=str,
+        default=None,
+        help='Path to a JSON config file for the pipeline (optional)'
+    )
+
+    args = parser.parse_args()
+
+    # Load pipeline config if provided
+    pipeline_config = None
+    if args.config:
+        with open(args.config, 'r') as f:
+            pipeline_config = json.load(f)
+            # Convert dict to DataSourceConfig and PreprocessConfig objects
+            if PIPELINE_AVAILABLE:
+                sources = [DataSourceConfig(**src) for src in pipeline_config['sources']]
+                preprocess_config = PreprocessConfig(**pipeline_config['preprocess_config'])
+                pipeline_config = {
+                    'sources': sources,
+                    'preprocess_config': preprocess_config
+                }
+
+    main(use_pipeline=args.use_pipeline, pipeline_config=pipeline_config)