Upload folder using huggingface_hub

spectral/__init__.py

@@ -0,0 +1,9 @@
+"""
+Spectral analysis module for Mamba SSM eigenvalue analysis.
+"""
+
+from mamba_spectral.spectral.eigenvalue_analyzer import SpectralAnalyzer
+from mamba_spectral.spectral.gramian import ReachabilityGramian
+from mamba_spectral.spectral.horizon_predictor import HorizonPredictor
+
+__all__ = ["SpectralAnalyzer", "ReachabilityGramian", "HorizonPredictor"]

spectral/eigenvalue_analyzer.py

@@ -0,0 +1,598 @@
+"""
+Eigenvalue Analyzer for Mamba SSM.
+
+Computes and analyzes eigenvalues of the state transition matrix to
+understand memory dynamics and reasoning capabilities.
+
+References:
+    [Gu2023] Gu, A., & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling
+             with Selective State Spaces. arXiv:2312.00752
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+from sklearn.cluster import KMeans
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SpectralTrajectory:
+    """
+    Container for eigenvalue evolution over tokens.
+    
+    Attributes:
+        timesteps: Token indices where measurements were taken.
+        eigenvalues: List of eigenvalue arrays at each timestep.
+        spectral_radius: Spectral radius (max |λ|) at each timestep.
+        delta_values: Discretization step Δ at each timestep.
+        layer_idx: Layer index being tracked.
+    """
+    timesteps: List[int] = field(default_factory=list)
+    eigenvalues: List[np.ndarray] = field(default_factory=list)
+    spectral_radius: List[float] = field(default_factory=list)
+    delta_values: Optional[List[np.ndarray]] = None
+    layer_idx: Optional[int] = None
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "timesteps": self.timesteps,
+            "eigenvalues": [e.tolist() for e in self.eigenvalues],
+            "spectral_radius": self.spectral_radius,
+            "delta_values": [d.tolist() for d in self.delta_values] if self.delta_values else None,
+            "layer_idx": self.layer_idx,
+        }
+
+
+@dataclass
+class ClusterResult:
+    """
+    Result of eigenvalue clustering.
+    
+    Attributes:
+        centers: Cluster centers in complex plane.
+        labels: Cluster label for each eigenvalue.
+        sizes: Number of eigenvalues in each cluster.
+        inertia: Sum of squared distances to cluster centers.
+    """
+    centers: np.ndarray
+    labels: np.ndarray
+    sizes: np.ndarray
+    inertia: float
+
+
+class SpectralAnalyzer:
+    """
+    Analyzes eigenvalues of the SSM transition matrix in Mamba models.
+    
+    The eigenvalues λ of the discretized matrix Ā determine:
+    - |λ| ≈ 1: Long-term memory (information preserved)
+    - |λ| ≈ 0: Fast forgetting (filters noise)
+    - |λ| > 1: Unstable dynamics (exponential growth)
+    
+    For Mamba, A is diagonal, so eigenvalues are simply the diagonal elements.
+    
+    Attributes:
+        wrapper: MambaWrapper instance.
+        extractor: StateExtractor for accessing A matrices.
+        device: Computation device.
+        
+    Example:
+        >>> analyzer = SpectralAnalyzer(model)
+        >>> eigenvalues = analyzer.compute_eigenvalues(layer_idx=0)
+        >>> print(f"Spectral radius: {analyzer.spectral_radius(eigenvalues):.4f}")
+        
+        >>> trajectory = analyzer.track_evolution(prompt_tokens)
+        >>> print(f"Radius declined from {trajectory.spectral_radius[0]:.3f} "
+        ...       f"to {trajectory.spectral_radius[-1]:.3f}")
+    """
+    
+    def __init__(
+        self,
+        model: Union["MambaWrapper", nn.Module],
+        device: str = "cuda",
+    ) -> None:
+        """
+        Initialize the spectral analyzer.
+        
+        Args:
+            model: MambaWrapper or raw Mamba model.
+            device: Device for computation ('cuda' or 'cpu').
+        """
+        # Handle both MambaWrapper and raw model
+        from mamba_spectral.core.mamba_wrapper import MambaWrapper
+        from mamba_spectral.core.state_extractor import StateExtractor
+        
+        if isinstance(model, MambaWrapper):
+            self.wrapper = model
+        else:
+            self.wrapper = MambaWrapper(model, device=device)
+        
+        self.extractor = StateExtractor(self.wrapper)
+        self.device = device
+        
+        logger.info(f"SpectralAnalyzer initialized on {device}")
+    
+    def extract_A_matrix(
+        self,
+        layer_idx: int,
+        as_numpy: bool = False,
+    ) -> Union[torch.Tensor, np.ndarray]:
+        """
+        Extract the continuous A matrix from a specific layer.
+        
+        In Mamba, A is stored as A_log and computed as A = -exp(A_log).
+        This ensures stability (negative eigenvalues → stable dynamics).
+        
+        Args:
+            layer_idx: Index of the Mamba layer (0-indexed).
+            as_numpy: If True, return numpy array instead of tensor.
+            
+        Returns:
+            A matrix of shape [d_inner, d_state].
+            
+        Example:
+            >>> A = analyzer.extract_A_matrix(layer_idx=0)
+            >>> print(f"A shape: {A.shape}")
+        """
+        A = self.extractor.extract_A_matrix(layer_idx)
+        
+        if as_numpy:
+            return A.detach().cpu().numpy()
+        return A
+    
+    def discretize_A(
+        self,
+        A_continuous: torch.Tensor,
+        delta: Union[torch.Tensor, float],
+        method: str = "zoh",
+    ) -> torch.Tensor:
+        """
+        Discretize the continuous A matrix using Zero-Order Hold.
+        
+        Ā = exp(Δ · A)
+        
+        For diagonal A (as in Mamba), this is element-wise exponentiation.
+        
+        Args:
+            A_continuous: Continuous A matrix. Shape: [d_inner, d_state].
+            delta: Discretization step. Can be:
+                - float: Single value applied uniformly
+                - Tensor [d_inner]: Per-channel delta
+                - Tensor [batch, seq_len, d_inner]: Input-dependent delta
+            method: Discretization method:
+                - 'zoh': Zero-Order Hold (default, recommended)
+                - 'euler': Forward Euler (less accurate)
+                
+        Returns:
+            torch.Tensor: Discretized Ā matrix.
+            
+        Example:
+            >>> A = analyzer.extract_A_matrix(0)
+            >>> A_bar = analyzer.discretize_A(A, delta=0.01)
+            >>> eigenvalues = analyzer.compute_eigenvalues(A_bar)
+        """
+        return self.extractor.discretize(A_continuous, delta, method)
+    
+    def compute_eigenvalues(
+        self,
+        A_matrix: Union[torch.Tensor, np.ndarray],
+        sort_by: str = "magnitude",
+    ) -> np.ndarray:
+        """
+        Compute eigenvalues of the A matrix.
+        
+        For diagonal matrices (as in Mamba), eigenvalues ARE the diagonal.
+        For non-diagonal matrices, uses torch.linalg.eigvals.
+        
+        Args:
+            A_matrix: State transition matrix to analyze.
+                - Diagonal: Shape [d_inner, d_state] or [d_state]
+                - Full: Shape [d_state, d_state]
+            sort_by: How to sort eigenvalues:
+                - 'magnitude': |λ| descending (default)
+                - 'real': Real part descending
+                - 'none': No sorting
+                
+        Returns:
+            np.ndarray: Complex eigenvalues sorted as specified.
+            
+        Note:
+            For Mamba's diagonal A, eigenvalues are real and negative
+            (since A = -exp(A_log)).
+        """
+        if isinstance(A_matrix, torch.Tensor):
+            A_matrix = A_matrix.detach().cpu().numpy()
+        
+        # Check if diagonal (Mamba case)
+        if A_matrix.ndim == 1:
+            # Already diagonal
+            eigenvalues = A_matrix.astype(np.complex128)
+        elif A_matrix.ndim == 2 and A_matrix.shape[0] != A_matrix.shape[1]:
+            # [d_inner, d_state] - flatten diagonal
+            eigenvalues = A_matrix.flatten().astype(np.complex128)
+        else:
+            # Full square matrix - compute eigvals
+            eigenvalues = np.linalg.eigvals(A_matrix)
+        
+        # Sort eigenvalues
+        if sort_by == "magnitude":
+            indices = np.argsort(-np.abs(eigenvalues))
+            eigenvalues = eigenvalues[indices]
+        elif sort_by == "real":
+            indices = np.argsort(-np.real(eigenvalues))
+            eigenvalues = eigenvalues[indices]
+        # else: no sorting
+        
+        return eigenvalues
+    
+    def spectral_radius(
+        self,
+        eigenvalues: Union[np.ndarray, torch.Tensor],
+    ) -> float:
+        """
+        Compute the spectral radius ρ(A) = max(|λ_i|).
+        
+        The spectral radius determines:
+        - ρ < 1: System is stable (states decay)
+        - ρ = 1: Marginally stable (states persist)
+        - ρ > 1: Unstable (states grow exponentially)
+        
+        Args:
+            eigenvalues: Array of eigenvalues.
+            
+        Returns:
+            float: Maximum absolute eigenvalue.
+            
+        Example:
+            >>> eigenvalues = analyzer.compute_eigenvalues(A_bar)
+            >>> rho = analyzer.spectral_radius(eigenvalues)
+            >>> print(f"Spectral radius: {rho:.4f}")
+            >>> if rho > 1:
+            ...     print("WARNING: System is unstable!")
+        """
+        if isinstance(eigenvalues, torch.Tensor):
+            eigenvalues = eigenvalues.detach().cpu().numpy()
+        
+        return float(np.max(np.abs(eigenvalues)))
+    
+    def spectral_gap(
+        self,
+        eigenvalues: np.ndarray,
+    ) -> float:
+        """
+        Compute the spectral gap = |λ_1| - |λ_2|.
+        
+        A larger gap indicates more distinct dominant dynamics.
+        
+        Args:
+            eigenvalues: Array of eigenvalues (will be sorted by magnitude).
+            
+        Returns:
+            float: Difference between top two eigenvalue magnitudes.
+        """
+        magnitudes = np.sort(np.abs(eigenvalues))[::-1]
+        
+        if len(magnitudes) < 2:
+            return 0.0
+        
+        return float(magnitudes[0] - magnitudes[1])
+    
+    def condition_number(
+        self,
+        eigenvalues: np.ndarray,
+    ) -> float:
+        """
+        Compute condition number = |λ_max| / |λ_min|.
+        
+        High condition number indicates sensitivity to perturbations.
+        
+        Args:
+            eigenvalues: Array of eigenvalues.
+            
+        Returns:
+            float: Ratio of largest to smallest eigenvalue magnitude.
+        """
+        magnitudes = np.abs(eigenvalues)
+        magnitudes = magnitudes[magnitudes > 1e-10]  # Filter near-zero
+        
+        if len(magnitudes) == 0:
+            return float("inf")
+        
+        return float(np.max(magnitudes) / np.min(magnitudes))
+    
+    def eigenvalue_clustering(
+        self,
+        eigenvalues: np.ndarray,
+        n_clusters: int = 2,
+        random_state: int = 42,
+    ) -> ClusterResult:
+        """
+        Detect clustering of eigenvalues in the complex plane.
+        
+        Clustering reveals "spectral engramas" - groups of eigenvalues
+        that may encode different types of information.
+        
+        Args:
+            eigenvalues: Complex eigenvalues to cluster.
+            n_clusters: Number of clusters (default: 2).
+            random_state: Random seed for reproducibility.
+            
+        Returns:
+            ClusterResult with:
+                - centers: Cluster center locations
+                - labels: Cluster assignment for each eigenvalue
+                - sizes: Number of eigenvalues per cluster
+                - inertia: Clustering quality metric
+                
+        Example:
+            >>> result = analyzer.eigenvalue_clustering(eigenvalues, n_clusters=3)
+            >>> print(f"Cluster sizes: {result.sizes}")
+            >>> print(f"Centers: {result.centers}")
+        """
+        # Convert complex to 2D real coordinates
+        X = np.column_stack([
+            np.real(eigenvalues),
+            np.imag(eigenvalues),
+        ])
+        
+        # Fit K-means
+        kmeans = KMeans(
+            n_clusters=n_clusters,
+            random_state=random_state,
+            n_init=10,
+        )
+        labels = kmeans.fit_predict(X)
+        
+        # Convert centers back to complex
+        centers = kmeans.cluster_centers_[:, 0] + 1j * kmeans.cluster_centers_[:, 1]
+        
+        # Count cluster sizes
+        sizes = np.bincount(labels, minlength=n_clusters)
+        
+        return ClusterResult(
+            centers=centers,
+            labels=labels,
+            sizes=sizes,
+            inertia=float(kmeans.inertia_),
+        )
+    
+    @torch.no_grad()
+    def track_evolution(
+        self,
+        prompt: Union[str, torch.Tensor],
+        layer_idx: int = 0,
+        save_every: int = 1,
+        delta_value: float = 0.01,
+    ) -> SpectralTrajectory:
+        """
+        Track eigenvalue evolution token-by-token during inference.
+        
+        This is the core method for understanding how spectral properties
+        change as the model processes input.
+        
+        Args:
+            prompt: Input text (string) or token IDs (tensor).
+            layer_idx: Which Mamba layer to analyze (0-indexed).
+            save_every: Record every N tokens (1 = all tokens).
+            delta_value: Default delta for discretization.
+            
+        Returns:
+            SpectralTrajectory containing:
+                - timesteps: Token indices
+                - eigenvalues: Eigenvalues at each timestep
+                - spectral_radius: ρ(Ā) at each timestep
+                
+        Example:
+            >>> trajectory = analyzer.track_evolution(
+            ...     "The capital of France is Paris, located in Europe.",
+            ...     layer_idx=0,
+            ... )
+            >>> import matplotlib.pyplot as plt
+            >>> plt.plot(trajectory.timesteps, trajectory.spectral_radius)
+            >>> plt.xlabel("Token")
+            >>> plt.ylabel("Spectral Radius")
+            >>> plt.show()
+        """
+        # Tokenize if string
+        if isinstance(prompt, str):
+            if self.wrapper.tokenizer is None:
+                raise ValueError("Tokenizer required for string input")
+            tokens = self.wrapper.tokenizer.encode(prompt, return_tensors="pt")
+            tokens = tokens.to(self.device)
+        else:
+            tokens = prompt.to(self.device)
+        
+        trajectory = SpectralTrajectory(layer_idx=layer_idx)
+        seq_len = tokens.shape[1]
+        
+        # Get base A matrix
+        A = self.extract_A_matrix(layer_idx)
+        delta = torch.tensor(delta_value, device=self.device)
+        
+        # Track through sequence
+        for t in range(0, seq_len, save_every):
+            # Create partial input
+            partial_tokens = tokens[:, :t+1]
+            
+            # Get discretized A for this context
+            # Note: In full implementation, would extract Δ from forward pass
+            A_bar = self.discretize_A(A, delta)
+            
+            # Compute eigenvalues
+            eigenvalues = self.compute_eigenvalues(A_bar)
+            radius = self.spectral_radius(eigenvalues)
+            
+            trajectory.timesteps.append(t)
+            trajectory.eigenvalues.append(eigenvalues)
+            trajectory.spectral_radius.append(radius)
+        
+        logger.info(
+            f"Tracked {len(trajectory.timesteps)} timesteps. "
+            f"Radius range: [{min(trajectory.spectral_radius):.3f}, "
+            f"{max(trajectory.spectral_radius):.3f}]"
+        )
+        
+        return trajectory
+    
+    @torch.no_grad()
+    def track_evolution_with_hooks(
+        self,
+        prompt: Union[str, torch.Tensor],
+        layer_indices: Optional[List[int]] = None,
+    ) -> Dict[int, SpectralTrajectory]:
+        """
+        Track evolution using forward hooks for accurate Δ extraction.
+        
+        This method registers hooks to capture the actual input-dependent
+        discretization step during forward pass.
+        
+        Args:
+            prompt: Input text or tokens.
+            layer_indices: Which layers to track (None = all).
+            
+        Returns:
+            Dict mapping layer_idx to SpectralTrajectory.
+        """
+        # Tokenize
+        if isinstance(prompt, str):
+            tokens = self.wrapper.tokenizer.encode(prompt, return_tensors="pt")
+            tokens = tokens.to(self.device)
+        else:
+            tokens = prompt.to(self.device)
+        
+        layers = self.wrapper.get_mamba_layers()
+        if layer_indices is None:
+            layer_indices = list(range(len(layers)))
+        
+        # Storage for captured data
+        captured_delta: Dict[int, List[torch.Tensor]] = {i: [] for i in layer_indices}
+        
+        def make_hook(layer_idx):
+            def hook(module, input, output):
+                # Try to extract delta from module state
+                if hasattr(module, "_delta"):
+                    captured_delta[layer_idx].append(module._delta.detach().cpu())
+            return hook
+        
+        # Register hooks
+        handles = []
+        for idx in layer_indices:
+            handle = self.wrapper.register_hook(idx, make_hook(idx), "forward")
+            handles.append(handle)
+        
+        try:
+            # Forward pass
+            _ = self.wrapper.forward(tokens)
+        finally:
+            # Clean up hooks
+            for handle in handles:
+                handle.remove()
+        
+        # Build trajectories
+        trajectories = {}
+        for layer_idx in layer_indices:
+            A = self.extract_A_matrix(layer_idx)
+            trajectory = SpectralTrajectory(layer_idx=layer_idx)
+            
+            deltas = captured_delta.get(layer_idx, [])
+            if deltas:
+                for t, delta in enumerate(deltas):
+                    A_bar = self.discretize_A(A, delta.to(self.device))
+                    eigenvalues = self.compute_eigenvalues(A_bar)
+                    
+                    trajectory.timesteps.append(t)
+                    trajectory.eigenvalues.append(eigenvalues)
+                    trajectory.spectral_radius.append(self.spectral_radius(eigenvalues))
+            else:
+                # Fallback: use default delta
+                trajectory = self.track_evolution(tokens, layer_idx)
+            
+            trajectories[layer_idx] = trajectory
+        
+        return trajectories
+    
+    def summarize_layer(
+        self,
+        layer_idx: int,
+        delta_value: float = 0.01,
+    ) -> Dict[str, Any]:
+        """
+        Generate summary statistics for a layer's spectral properties.
+        
+        Args:
+            layer_idx: Layer to analyze.
+            delta_value: Discretization step.
+            
+        Returns:
+            Dictionary with:
+                - spectral_radius: ρ(Ā)
+                - spectral_gap: |λ_1| - |λ_2|
+                - condition_number: |λ_max| / |λ_min|
+                - num_stable: Count of |λ| < 1
+                - num_unstable: Count of |λ| > 1
+                - mean_magnitude: Average |λ|
+                - std_magnitude: Std of |λ|
+        """
+        A = self.extract_A_matrix(layer_idx)
+        A_bar = self.discretize_A(A, torch.tensor(delta_value))
+        eigenvalues = self.compute_eigenvalues(A_bar)
+        magnitudes = np.abs(eigenvalues)
+        
+        return {
+            "layer_idx": layer_idx,
+            "num_eigenvalues": len(eigenvalues),
+            "spectral_radius": self.spectral_radius(eigenvalues),
+            "spectral_gap": self.spectral_gap(eigenvalues),
+            "condition_number": self.condition_number(eigenvalues),
+            "num_stable": int(np.sum(magnitudes < 1)),
+            "num_unstable": int(np.sum(magnitudes > 1)),
+            "num_marginal": int(np.sum(np.isclose(magnitudes, 1, atol=0.01))),
+            "mean_magnitude": float(np.mean(magnitudes)),
+            "std_magnitude": float(np.std(magnitudes)),
+            "min_magnitude": float(np.min(magnitudes)),
+            "max_magnitude": float(np.max(magnitudes)),
+        }
+    
+    def analyze_all_layers(
+        self,
+        delta_value: float = 0.01,
+    ) -> List[Dict[str, Any]]:
+        """
+        Analyze spectral properties of all Mamba layers.
+        
+        Args:
+            delta_value: Discretization step.
+            
+        Returns:
+            List of summary dictionaries, one per layer.
+        """
+        layers = self.wrapper.get_mamba_layers()
+        results = []
+        
+        for i in range(len(layers)):
+            try:
+                summary = self.summarize_layer(i, delta_value)
+                results.append(summary)
+            except Exception as e:
+                logger.warning(f"Could not analyze layer {i}: {e}")
+                results.append({"layer_idx": i, "error": str(e)})
+        
+        return results
+    
+    def __repr__(self) -> str:
+        return (
+            f"SpectralAnalyzer(\n"
+            f"  model='{self.wrapper.model_name}',\n"
+            f"  num_layers={len(self.wrapper.get_mamba_layers())},\n"
+            f"  device='{self.device}'\n"
+            f")"
+        )

spectral/gramian.py

@@ -0,0 +1,399 @@
+"""
+Reachability Gramian Computation for Mamba SSM.
+
+The Reachability Gramian W_R determines which states are reachable from
+input and quantifies the reasoning horizon of the model.
+
+References:
+    [Antsaklis2007] Antsaklis, P. J., & Michel, A. N. (2007). 
+                    A Linear Systems Primer. Birkhäuser.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class GramianResult:
+    """
+    Result of Gramian computation.
+    
+    Attributes:
+        gramian: The reachability gramian W_R.
+        singular_values: Singular values of W_R at each step.
+        rank: Numerical rank of W_R (σ > threshold).
+        min_singular_value: Smallest singular value.
+        is_full_rank: Whether gramian has full rank.
+        horizon_reached: Step where reachability became limited.
+    """
+    gramian: np.ndarray
+    singular_values: List[np.ndarray]
+    rank: int
+    min_singular_value: float
+    is_full_rank: bool
+    horizon_reached: Optional[int] = None
+
+
+class ReachabilityGramian:
+    """
+    Computes the Reachability Gramian for SSM analysis.
+    
+    The Reachability Gramian is defined as:
+        W_R(T) = Σ_{k=0}^{T-1} Ā^k · B̄ · B̄^T · (Ā^T)^k
+    
+    Properties:
+    - If W_R is full rank, all states are reachable from input
+    - If W_R is singular, some states are unreachable
+    - The horizon H is where σ_min(W_R) drops below threshold
+    
+    This determines the Spectral Horizon - the maximum number of
+    reasoning steps the model can effectively perform.
+    
+    Example:
+        >>> gramian_calc = ReachabilityGramian(device='cuda')
+        >>> result = gramian_calc.compute(A_bar, B_bar, horizon=100)
+        >>> print(f"Rank: {result.rank}, Full rank: {result.is_full_rank}")
+    """
+    
+    def __init__(
+        self,
+        device: str = "cuda",
+        dtype: torch.dtype = torch.float64,
+    ) -> None:
+        """
+        Initialize the Gramian calculator.
+        
+        Args:
+            device: Computation device.
+            dtype: Data type (float64 recommended for numerical accuracy).
+        """
+        self.device = device
+        self.dtype = dtype
+        
+        logger.info(f"ReachabilityGramian initialized on {device} with {dtype}")
+    
+    def compute(
+        self,
+        A_bar: torch.Tensor,
+        B_bar: torch.Tensor,
+        horizon: int = 100,
+        threshold: float = 1e-10,
+        track_singular_values: bool = True,
+        show_progress: bool = False,
+    ) -> GramianResult:
+        """
+        Compute the Reachability Gramian iteratively.
+        
+        W_R(T) = Σ_{k=0}^{T-1} Ā^k · B̄ · B̄^T · (Ā^T)^k
+        
+        Uses iterative accumulation to avoid O(T × n^3) matrix multiplications.
+        
+        Args:
+            A_bar: Discretized state transition. Shape: [d_state, d_state] or [d_inner, d_state] (diagonal).
+            B_bar: Discretized input matrix. Shape: [d_state, d_input] or [d_state].
+            horizon: Maximum number of steps (T).
+            threshold: Singular value threshold for rank determination.
+            track_singular_values: If True, record σ at each step.
+            show_progress: Show progress bar.
+            
+        Returns:
+            GramianResult with gramian, singular values, and rank info.
+            
+        Note:
+            For diagonal A (as in Mamba), computation is much more efficient
+            as we only need elementwise operations.
+        """
+        # Ensure tensors are on device with proper dtype
+        A_bar = A_bar.to(device=self.device, dtype=self.dtype)
+        B_bar = B_bar.to(device=self.device, dtype=self.dtype)
+        
+        # Handle diagonal A case (Mamba)
+        is_diagonal = (A_bar.ndim == 1) or (A_bar.ndim == 2 and A_bar.shape[0] != A_bar.shape[1])
+        
+        if is_diagonal:
+            return self._compute_diagonal(A_bar, B_bar, horizon, threshold, track_singular_values, show_progress)
+        else:
+            return self._compute_full(A_bar, B_bar, horizon, threshold, track_singular_values, show_progress)
+    
+    def _compute_diagonal(
+        self,
+        A_bar: torch.Tensor,
+        B_bar: torch.Tensor,
+        horizon: int,
+        threshold: float,
+        track_singular_values: bool,
+        show_progress: bool,
+    ) -> GramianResult:
+        """
+        Efficient computation for diagonal A matrices.
+        
+        For diagonal A with elements a_i:
+        W_R = Σ_k diag(a)^k B B^T diag(a)^k
+            = Σ_k diag(a^{2k}) ⊙ (B B^T)  (for B as column vector)
+            
+        This allows elementwise operations instead of matrix multiplication.
+        """
+        # Flatten to diagonal
+        if A_bar.ndim == 2:
+            a = A_bar.flatten()  # [d_inner * d_state]
+        else:
+            a = A_bar
+        
+        n = len(a)
+        
+        # Ensure B is column vector
+        if B_bar.ndim == 1:
+            B = B_bar.view(-1, 1)
+        else:
+            B = B_bar
+        
+        # Adjust dimensions if needed
+        if B.shape[0] != n:
+            # Take first n elements or pad
+            if B.shape[0] > n:
+                B = B[:n]
+            else:
+                B = torch.nn.functional.pad(B, (0, 0, 0, n - B.shape[0]))
+        
+        # B B^T outer product
+        BBT = B @ B.T  # [n, n]
+        
+        # Initialize accumulator
+        W_R = torch.zeros(n, n, device=self.device, dtype=self.dtype)
+        
+        # Powers of a (elementwise)
+        a_power = torch.ones_like(a)  # a^0 = 1
+        
+        # Track singular values
+        singular_values = []
+        horizon_reached = None
+        
+        iterator = range(horizon)
+        if show_progress:
+            iterator = tqdm(iterator, desc="Computing Gramian")
+        
+        for k in iterator:
+            # W_R += diag(a_power) @ BBT @ diag(a_power)
+            # Which is: (a_power.unsqueeze(1) * a_power.unsqueeze(0)) * BBT
+            outer = a_power.unsqueeze(1) * a_power.unsqueeze(0)
+            W_R += outer * BBT
+            
+            # Update power
+            a_power = a_power * a
+            
+            # Track singular values
+            if track_singular_values:
+                try:
+                    s = torch.linalg.svdvals(W_R)
+                    singular_values.append(s.cpu().numpy())
+                    
+                    # Check if horizon reached
+                    if horizon_reached is None and s.min() < threshold:
+                        horizon_reached = k
+                except Exception as e:
+                    logger.warning(f"SVD failed at step {k}: {e}")
+        
+        # Final analysis
+        W_R_np = W_R.cpu().numpy()
+        
+        try:
+            final_sv = np.linalg.svd(W_R_np, compute_uv=False)
+            rank = int(np.sum(final_sv > threshold))
+            min_sv = float(final_sv[-1]) if len(final_sv) > 0 else 0.0
+        except Exception:
+            final_sv = np.array([])
+            rank = 0
+            min_sv = 0.0
+        
+        return GramianResult(
+            gramian=W_R_np,
+            singular_values=singular_values,
+            rank=rank,
+            min_singular_value=min_sv,
+            is_full_rank=(rank == n),
+            horizon_reached=horizon_reached,
+        )
+    
+    def _compute_full(
+        self,
+        A_bar: torch.Tensor,
+        B_bar: torch.Tensor,
+        horizon: int,
+        threshold: float,
+        track_singular_values: bool,
+        show_progress: bool,
+    ) -> GramianResult:
+        """
+        Standard computation for full (non-diagonal) A matrices.
+        
+        W_R = Σ_{k=0}^{T-1} A^k B B^T (A^T)^k
+        """
+        n = A_bar.shape[0]
+        
+        # Ensure B is proper shape
+        if B_bar.ndim == 1:
+            B = B_bar.view(-1, 1)
+        else:
+            B = B_bar
+        
+        BBT = B @ B.T
+        
+        # Initialize
+        W_R = torch.zeros(n, n, device=self.device, dtype=self.dtype)
+        A_power = torch.eye(n, device=self.device, dtype=self.dtype)
+        
+        singular_values = []
+        horizon_reached = None
+        
+        iterator = range(horizon)
+        if show_progress:
+            iterator = tqdm(iterator, desc="Computing Gramian")
+        
+        for k in iterator:
+            # W_R += A^k B B^T (A^T)^k
+            term = A_power @ BBT @ A_power.T
+            W_R += term
+            
+            # Update A power
+            A_power = A_bar @ A_power
+            
+            # Track singular values
+            if track_singular_values:
+                try:
+                    s = torch.linalg.svdvals(W_R)
+                    singular_values.append(s.cpu().numpy())
+                    
+                    if horizon_reached is None and s.min() < threshold:
+                        horizon_reached = k
+                except Exception as e:
+                    logger.warning(f"SVD failed at step {k}: {e}")
+        
+        # Final analysis
+        W_R_np = W_R.cpu().numpy()
+        
+        try:
+            final_sv = np.linalg.svd(W_R_np, compute_uv=False)
+            rank = int(np.sum(final_sv > threshold))
+            min_sv = float(final_sv[-1]) if len(final_sv) > 0 else 0.0
+        except Exception:
+            rank = 0
+            min_sv = 0.0
+            final_sv = np.array([])
+        
+        return GramianResult(
+            gramian=W_R_np,
+            singular_values=singular_values,
+            rank=rank,
+            min_singular_value=min_sv,
+            is_full_rank=(rank == n),
+            horizon_reached=horizon_reached,
+        )
+    
+    def compute_controllability_matrix(
+        self,
+        A_bar: torch.Tensor,
+        B_bar: torch.Tensor,
+        n_steps: Optional[int] = None,
+    ) -> Tuple[np.ndarray, int]:
+        """
+        Compute the controllability matrix C = [B, AB, A²B, ..., A^{n-1}B].
+        
+        Alternative to Gramian for checking reachability.
+        
+        Args:
+            A_bar: State transition matrix.
+            B_bar: Input matrix.
+            n_steps: Number of steps (defaults to state dimension).
+            
+        Returns:
+            Tuple of (controllability matrix, numerical rank).
+        """
+        A_bar = A_bar.to(device=self.device, dtype=self.dtype)
+        B_bar = B_bar.to(device=self.device, dtype=self.dtype)
+        
+        n = A_bar.shape[0]
+        if n_steps is None:
+            n_steps = n
+        
+        # Ensure B is 2D
+        if B_bar.ndim == 1:
+            B = B_bar.view(-1, 1)
+        else:
+            B = B_bar
+        
+        m = B.shape[1]
+        
+        # Build controllability matrix
+        C = torch.zeros(n, n_steps * m, device=self.device, dtype=self.dtype)
+        
+        A_power = torch.eye(n, device=self.device, dtype=self.dtype)
+        for k in range(n_steps):
+            C[:, k*m:(k+1)*m] = A_power @ B
+            A_power = A_bar @ A_power
+        
+        C_np = C.cpu().numpy()
+        
+        # Compute rank
+        rank = int(np.linalg.matrix_rank(C_np))
+        
+        return C_np, rank
+    
+    def analyze_reachability(
+        self,
+        A_bar: torch.Tensor,
+        B_bar: torch.Tensor,
+        horizon: int = 100,
+        threshold: float = 1e-10,
+    ) -> Dict[str, Any]:
+        """
+        Comprehensive reachability analysis.
+        
+        Args:
+            A_bar: Discretized transition matrix.
+            B_bar: Discretized input matrix.
+            horizon: Analysis horizon.
+            threshold: Numerical threshold.
+            
+        Returns:
+            Dictionary with full analysis results.
+        """
+        # Compute gramian
+        gramian_result = self.compute(
+            A_bar, B_bar, horizon, threshold,
+            track_singular_values=True, show_progress=False,
+        )
+        
+        # Compute controllability matrix
+        C, ctrl_rank = self.compute_controllability_matrix(A_bar, B_bar)
+        
+        # Get state dimension
+        if A_bar.ndim == 1:
+            n = len(A_bar)
+        elif A_bar.ndim == 2 and A_bar.shape[0] != A_bar.shape[1]:
+            n = A_bar.numel()
+        else:
+            n = A_bar.shape[0]
+        
+        return {
+            "gramian_rank": gramian_result.rank,
+            "controllability_rank": ctrl_rank,
+            "state_dimension": n,
+            "is_fully_reachable": gramian_result.is_full_rank,
+            "min_singular_value": gramian_result.min_singular_value,
+            "horizon_reached": gramian_result.horizon_reached,
+            "rank_deficit": n - gramian_result.rank,
+            "gramian": gramian_result.gramian,
+            "singular_value_trajectory": gramian_result.singular_values,
+        }
+    
+    def __repr__(self) -> str:
+        return f"ReachabilityGramian(device='{self.device}', dtype={self.dtype})"

spectral/horizon_predictor.py

@@ -0,0 +1,454 @@
+"""
+Horizon Predictor for Mamba SSM.
+
+Predicts the mathematical limit of reasoning (Spectral Horizon) using
+eigenvalue analysis and reachability gramian.
+
+References:
+    [Gu2023] Spectral analysis of state space models.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class HorizonResult:
+    """
+    Result of horizon prediction.
+    
+    Attributes:
+        horizon: Predicted reasoning horizon (number of tokens).
+        gramian_rank: Rank of reachability gramian at horizon.
+        min_singular_value: Smallest singular value at horizon.
+        spectral_radius: Spectral radius of Ā.
+        is_reachable: Whether full state is reachable.
+        confidence: Confidence score (0-1).
+        method: Method used for prediction.
+    """
+    horizon: int
+    gramian_rank: int
+    min_singular_value: float
+    spectral_radius: float
+    is_reachable: bool
+    confidence: float = 1.0
+    method: str = "gramian"
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "horizon": self.horizon,
+            "gramian_rank": self.gramian_rank,
+            "min_singular_value": self.min_singular_value,
+            "spectral_radius": self.spectral_radius,
+            "is_reachable": self.is_reachable,
+            "confidence": self.confidence,
+            "method": self.method,
+        }
+
+
+class HorizonPredictor:
+    """
+    Predicts the Spectral Horizon - the mathematical limit of reasoning.
+    
+    The Spectral Horizon is determined by when the reachability gramian
+    becomes singular, indicating that the model can no longer access
+    information from earlier in the sequence.
+    
+    Theory:
+    - W_R(T) = Σ_{k=0}^{T-1} Ā^k · B̄ · B̄^T · (Ā^T)^k
+    - When σ_min(W_R) < ε, the horizon H = T is reached
+    - Beyond H, reasoning about earlier context becomes unreliable
+    
+    Example:
+        >>> predictor = HorizonPredictor(model)
+        >>> result = predictor.predict_horizon(
+        ...     "The capital of France is",
+        ...     max_horizon=1000,
+        ... )
+        >>> print(f"Predicted horizon: {result.horizon} tokens")
+        >>> print(f"Is fully reachable: {result.is_reachable}")
+    """
+    
+    def __init__(
+        self,
+        model: Union["MambaWrapper", nn.Module],
+        device: str = "cuda",
+    ) -> None:
+        """
+        Initialize the horizon predictor.
+        
+        Args:
+            model: Mamba model (MambaWrapper or raw model).
+            device: Computation device.
+        """
+        from mamba_spectral.core.mamba_wrapper import MambaWrapper
+        from mamba_spectral.core.state_extractor import StateExtractor
+        from mamba_spectral.spectral.eigenvalue_analyzer import SpectralAnalyzer
+        from mamba_spectral.spectral.gramian import ReachabilityGramian
+        
+        if isinstance(model, MambaWrapper):
+            self.wrapper = model
+        else:
+            self.wrapper = MambaWrapper(model, device=device)
+        
+        self.extractor = StateExtractor(self.wrapper)
+        self.analyzer = SpectralAnalyzer(self.wrapper, device=device)
+        self.gramian_calc = ReachabilityGramian(device=device)
+        self.device = device
+        
+        logger.info(f"HorizonPredictor initialized on {device}")
+    
+    def predict_horizon(
+        self,
+        prompt: Union[str, torch.Tensor],
+        layer_idx: int = 0,
+        threshold: float = 1e-6,
+        max_horizon: int = 1000,
+        delta_value: float = 0.01,
+        method: str = "gramian",
+    ) -> HorizonResult:
+        """
+        Predict the reasoning horizon for a given prompt.
+        
+        Computes the Spectral Horizon by analyzing when the reachability
+        gramian becomes singular (σ_min < threshold).
+        
+        Args:
+            prompt: Input text or token IDs.
+            layer_idx: Which Mamba layer to analyze.
+            threshold: Singular value threshold for horizon detection.
+            max_horizon: Maximum horizon to search.
+            delta_value: Discretization step.
+            method: Prediction method:
+                - 'gramian': Use reachability gramian (most accurate)
+                - 'spectral': Use spectral radius decay (faster)
+                - 'hybrid': Combine both methods
+                
+        Returns:
+            HorizonResult with predicted horizon and analysis.
+            
+        Example:
+            >>> result = predictor.predict_horizon(
+            ...     "The key is APPLE. [distraction] Query: key was?",
+            ...     threshold=1e-6,
+            ... )
+            >>> if result.horizon < 100:
+            ...     print("Warning: Limited reasoning capacity!")
+        """
+        # Get SSM matrices
+        A = self.extractor.extract_A_matrix(layer_idx)
+        delta = torch.tensor(delta_value, device=self.device)
+        A_bar = self.extractor.discretize(A, delta)
+        
+        # Create synthetic B (identity-like for analysis)
+        d_state = A.shape[-1]
+        B_bar = torch.eye(d_state, device=self.device, dtype=A_bar.dtype)[:, :min(d_state, 10)]
+        
+        # Compute spectral radius
+        eigenvalues = self.analyzer.compute_eigenvalues(A_bar)
+        spectral_radius = self.analyzer.spectral_radius(eigenvalues)
+        
+        if method == "gramian":
+            return self._predict_via_gramian(
+                A_bar, B_bar, threshold, max_horizon, spectral_radius
+            )
+        elif method == "spectral":
+            return self._predict_via_spectral(
+                A_bar, spectral_radius, threshold, max_horizon
+            )
+        elif method == "hybrid":
+            return self._predict_hybrid(
+                A_bar, B_bar, spectral_radius, threshold, max_horizon
+            )
+        else:
+            raise ValueError(f"Unknown method: {method}")
+    
+    def _predict_via_gramian(
+        self,
+        A_bar: torch.Tensor,
+        B_bar: torch.Tensor,
+        threshold: float,
+        max_horizon: int,
+        spectral_radius: float,
+    ) -> HorizonResult:
+        """Predict horizon using reachability gramian."""
+        result = self.gramian_calc.compute(
+            A_bar, B_bar,
+            horizon=max_horizon,
+            threshold=threshold,
+            track_singular_values=True,
+            show_progress=False,
+        )
+        
+        if result.horizon_reached is not None:
+            horizon = result.horizon_reached
+            confidence = 0.95
+        else:
+            # Horizon not reached within max_horizon
+            horizon = max_horizon
+            confidence = 0.5
+        
+        return HorizonResult(
+            horizon=horizon,
+            gramian_rank=result.rank,
+            min_singular_value=result.min_singular_value,
+            spectral_radius=spectral_radius,
+            is_reachable=result.is_full_rank,
+            confidence=confidence,
+            method="gramian",
+        )
+    
+    def _predict_via_spectral(
+        self,
+        A_bar: torch.Tensor,
+        spectral_radius: float,
+        threshold: float,
+        max_horizon: int,
+    ) -> HorizonResult:
+        """
+        Predict horizon using spectral radius decay.
+        
+        For ρ(Ā) < 1, memory decays as ρ^k.
+        Horizon ≈ log(threshold) / log(ρ)
+        """
+        if spectral_radius >= 1.0:
+            # Marginally stable or unstable - infinite theoretical horizon
+            horizon = max_horizon
+            confidence = 0.3  # Low confidence for edge case
+        else:
+            # Decay estimate: ρ^H < threshold
+            # H = log(threshold) / log(ρ)
+            try:
+                horizon = int(np.log(threshold) / np.log(spectral_radius))
+                horizon = min(horizon, max_horizon)
+                horizon = max(horizon, 1)
+                confidence = 0.7
+            except (ValueError, ZeroDivisionError):
+                horizon = max_horizon
+                confidence = 0.2
+        
+        return HorizonResult(
+            horizon=horizon,
+            gramian_rank=-1,  # Not computed
+            min_singular_value=-1.0,
+            spectral_radius=spectral_radius,
+            is_reachable=(spectral_radius >= 0.9),
+            confidence=confidence,
+            method="spectral",
+        )
+    
+    def _predict_hybrid(
+        self,
+        A_bar: torch.Tensor,
+        B_bar: torch.Tensor,
+        spectral_radius: float,
+        threshold: float,
+        max_horizon: int,
+    ) -> HorizonResult:
+        """Combine gramian and spectral methods."""
+        # First, quick spectral estimate
+        spectral_result = self._predict_via_spectral(
+            A_bar, spectral_radius, threshold, max_horizon
+        )
+        
+        # If spectral suggests short horizon, verify with gramian
+        if spectral_result.horizon < max_horizon // 2:
+            search_horizon = min(spectral_result.horizon * 2, max_horizon)
+            gramian_result = self._predict_via_gramian(
+                A_bar, B_bar, threshold, search_horizon, spectral_radius
+            )
+            
+            # Use gramian result if confident
+            if gramian_result.confidence > spectral_result.confidence:
+                gramian_result.method = "hybrid"
+                return gramian_result
+        
+        spectral_result.method = "hybrid"
+        return spectral_result
+    
+    def compute_reachability_gramian(
+        self,
+        A_bar: torch.Tensor,
+        B_bar: torch.Tensor,
+        horizon: int = 100,
+    ) -> Tuple[np.ndarray, List[np.ndarray]]:
+        """
+        Compute the Reachability Gramian.
+        
+        W_R(T) = Σ_{k=0}^{T-1} Ā^k · B̄ · B̄^T · (Ā^T)^k
+        
+        Args:
+            A_bar: Discretized transition matrix.
+            B_bar: Discretized input matrix.
+            horizon: Number of steps.
+            
+        Returns:
+            Tuple of (gramian matrix, singular values at each step).
+        """
+        result = self.gramian_calc.compute(
+            A_bar, B_bar, horizon,
+            track_singular_values=True,
+        )
+        return result.gramian, result.singular_values
+    
+    @torch.no_grad()
+    def adversarial_cot_generator(
+        self,
+        base_prompt: str,
+        target_horizon: int = 50,
+        n_candidates: int = 10,
+        layer_idx: int = 0,
+    ) -> Dict[str, Any]:
+        """
+        Generate Chain-of-Thought that DEGRADES reasoning capacity.
+        
+        This demonstrates that CoT can be adversarial by finding
+        token sequences that minimize the spectral radius.
+        
+        The idea:
+        - Normal CoT keeps ρ(Ā) high (preserves memory)
+        - Adversarial CoT forces ρ(Ā) → 0 (causes forgetting)
+        
+        Note:
+            Full gradient-based optimization is complex.
+            This PoC uses heuristic search over candidate tokens.
+        
+        Args:
+            base_prompt: Starting prompt.
+            target_horizon: Target (reduced) horizon.
+            n_candidates: Number of candidates to try per position.
+            layer_idx: Layer to analyze.
+            
+        Returns:
+            Dictionary with:
+                - adversarial_prompt: Generated adversarial CoT
+                - spectral_trajectory: ρ values over tokens
+                - original_radius: ρ of original prompt
+                - final_radius: ρ after adversarial tokens
+        """
+        if self.wrapper.tokenizer is None:
+            raise ValueError("Tokenizer required for CoT generation")
+        
+        # Encode base prompt
+        base_ids = self.wrapper.tokenizer.encode(base_prompt, return_tensors="pt")
+        base_ids = base_ids.to(self.device)
+        
+        # Track original spectral radius
+        trajectory_orig = self.analyzer.track_evolution(base_ids, layer_idx)
+        original_radius = trajectory_orig.spectral_radius[-1]
+        
+        # Heuristic: Find tokens that tend to reduce spectral radius
+        # These are typically: periods, padding, special characters
+        adversarial_tokens = []
+        
+        # Common "forgetting" tokens (heuristic)
+        forget_tokens = [
+            "...", "   ", "\n\n", "---", "___",
+            "[END]", "<pad>", "~", "...", "  ",
+        ]
+        
+        forget_ids = []
+        for token in forget_tokens:
+            try:
+                ids = self.wrapper.tokenizer.encode(token, add_special_tokens=False)
+                forget_ids.extend(ids)
+            except Exception:
+                pass
+        
+        # Remove duplicates
+        forget_ids = list(set(forget_ids))[:n_candidates]
+        
+        # Build adversarial sequence
+        current_ids = base_ids.clone()
+        spectral_trajectory = [original_radius]
+        
+        for step in range(min(target_horizon, 20)):  # Limit steps for PoC
+            best_token = None
+            lowest_radius = float('inf')
+            
+            for token_id in forget_ids:
+                # Try appending this token
+                test_ids = torch.cat([
+                    current_ids,
+                    torch.tensor([[token_id]], device=self.device)
+                ], dim=1)
+                
+                # Measure spectral radius
+                try:
+                    A = self.extractor.extract_A_matrix(layer_idx)
+                    A_bar = self.extractor.discretize(A, torch.tensor(0.01))
+                    eigenvalues = self.analyzer.compute_eigenvalues(A_bar)
+                    radius = self.analyzer.spectral_radius(eigenvalues)
+                    
+                    if radius < lowest_radius:
+                        lowest_radius = radius
+                        best_token = token_id
+                except Exception:
+                    continue
+            
+            if best_token is not None:
+                current_ids = torch.cat([
+                    current_ids,
+                    torch.tensor([[best_token]], device=self.device)
+                ], dim=1)
+                adversarial_tokens.append(best_token)
+                spectral_trajectory.append(lowest_radius)
+        
+        # Decode result
+        adversarial_prompt = self.wrapper.tokenizer.decode(current_ids[0])
+        
+        return {
+            "adversarial_prompt": adversarial_prompt,
+            "spectral_trajectory": spectral_trajectory,
+            "original_radius": original_radius,
+            "final_radius": spectral_trajectory[-1] if spectral_trajectory else original_radius,
+            "tokens_added": len(adversarial_tokens),
+        }
+    
+    def analyze_prompt_horizon(
+        self,
+        prompt: str,
+        layer_indices: Optional[List[int]] = None,
+    ) -> Dict[int, HorizonResult]:
+        """
+        Analyze horizon across multiple layers.
+        
+        Args:
+            prompt: Input text.
+            layer_indices: Layers to analyze (None = all).
+            
+        Returns:
+            Dictionary mapping layer index to HorizonResult.
+        """
+        layers = self.wrapper.get_mamba_layers()
+        if layer_indices is None:
+            layer_indices = list(range(len(layers)))
+        
+        results = {}
+        for idx in layer_indices:
+            try:
+                result = self.predict_horizon(prompt, layer_idx=idx)
+                results[idx] = result
+            except Exception as e:
+                logger.warning(f"Could not analyze layer {idx}: {e}")
+        
+        return results
+    
+    def __repr__(self) -> str:
+        return (
+            f"HorizonPredictor(\n"
+            f"  model='{self.wrapper.model_name}',\n"
+            f"  device='{self.device}'\n"
+            f")"
+        )