models/usage_tracker.py

# models/usage_tracker.py
import torch
import torch.nn as nn
from typing import Optional


class UsageTracker(nn.Module):
    """
    Hebbian-style usage tracker implementing "soft LFU" (Least Frequently Used).

    Tracks the "hotness" of each LoRA rank over time:
    - Ranks with high usage (frequently activated) are preserved and accumulated
    - Ranks with low usage (rarely activated) decay over time and can be overwritten

    This mechanism enables:
    1. Selective forgetting (Adaptive Parametric Forgetting)
    2. Knowledge accumulation for frequently used concepts
    3. Catastrophic forgetting prevention for stable memories

    Unlike Redis LFU (hard eviction), this uses soft decay to preserve
    distributed representations in continuous parametric space.
    """

    def __init__(
            self,
            num_layers: int,
            rank: int,
            decay_gamma: float = 0.95,
            epsilon: float = 1e-8,
            moving_average_alpha: float = 0.3
    ):
        """
        Args:
            num_layers: Number of target LoRA layers (N)
            rank: LoRA rank (r)
            decay_gamma: Temporal decay factor (0 < γ < 1).
                        Lower = faster forgetting.
            epsilon: Small constant for numerical stability
            moving_average_alpha: Smoothing factor for activation updates.
                                 Higher = more responsive to recent changes.
        """
        super().__init__()

        assert 0 < decay_gamma < 1, "decay_gamma must be in (0, 1)"
        assert 0 < moving_average_alpha <= 1, "alpha must be in (0, 1]"

        self.num_layers = num_layers
        self.rank = rank
        self.decay_gamma = decay_gamma
        self.epsilon = epsilon
        self.moving_average_alpha = moving_average_alpha

        # Usage statistics buffer (not a learnable parameter)
        # Will be dynamically resized based on batch size
        self.register_buffer(
            'usage_stats',
            torch.zeros(1, num_layers, rank),
            persistent=False  # Not saved in state_dict
        )

        # Track update count for monitoring
        self.register_buffer('update_count', torch.tensor(0), persistent=False)

        print(f"✅ Usage Tracker initialized:")
        print(f"   - Layers: {num_layers}, Rank: {rank}")
        print(f"   - Decay γ: {decay_gamma}")
        print(f"   - Moving average α: {moving_average_alpha}")

    @torch.no_grad()
    def compute_activation(self, lora_weights: torch.Tensor) -> torch.Tensor:
        """
        Compute activation magnitude for each rank using L2 norm.

        Args:
            lora_weights: Tensor of shape (B, N, r, d)

        Returns:
            activation: Tensor of shape (B, N, r), L2 norm of each rank vector
        """
        # L2 norm along hidden dimension (d)
        # Result: magnitude of each rank vector
        activation = torch.norm(lora_weights, p=2, dim=-1)  # (B, N, r)

        return activation

    @torch.no_grad()
    def apply_decay(self, usage: torch.Tensor) -> torch.Tensor:
        """
        Apply temporal decay to usage vector (Ebbinghaus forgetting curve).

        Args:
            usage: Current usage vector (B, N, r)

        Returns:
            decayed_usage: Usage after temporal decay
        """
        return usage * self.decay_gamma

    @torch.no_grad()
    def normalize(self, usage: torch.Tensor) -> torch.Tensor:
        """
        Normalize usage vector to [0, 1] range per sample.

        Uses min-max normalization with epsilon for stability.

        Args:
            usage: Usage vector (B, N, r)

        Returns:
            normalized_usage: Normalized usage in [0, 1]
        """
        # Per-sample normalization
        batch_size = usage.shape[0]
        usage_flat = usage.reshape(batch_size, -1)  # (B, N*r)

        min_val = usage_flat.min(dim=1, keepdim=True).values
        max_val = usage_flat.max(dim=1, keepdim=True).values

        # Avoid division by zero
        range_val = (max_val - min_val).clamp(min=self.epsilon)

        normalized = (usage_flat - min_val) / range_val
        normalized = normalized.reshape_as(usage)

        return normalized

    @torch.no_grad()
    def update(
            self,
            old_usage: Optional[torch.Tensor],
            lora_weights: torch.Tensor,
            normalize: bool = True
    ) -> torch.Tensor:
        """
        Update usage vector based on current LoRA weights.

        Implements Hebbian consolidation:
        1. Apply temporal decay to old usage (forgetting)
        2. Compute activation magnitude from current weights
        3. Combine via exponential moving average
        4. Optionally normalize to [0, 1]

        Args:
            old_usage: Previous usage vector (B, N, r), or None for first update
            lora_weights: Current LoRA weights (B, N, r, d)
            normalize: Whether to normalize output to [0, 1]

        Returns:
            new_usage: Updated usage vector (B, N, r)
        """
        batch_size = lora_weights.shape[0]

        # Step 1: Initialize or decay old usage
        if old_usage is None:
            # First update: start from zero
            decayed_usage = torch.zeros_like(
                lora_weights[..., 0],  # (B, N, r)
                device=lora_weights.device,
                dtype=lora_weights.dtype
            )
        else:
            # Validate shape
            assert old_usage.shape == (batch_size, self.num_layers, self.rank), \
                f"Expected usage shape {(batch_size, self.num_layers, self.rank)}, got {old_usage.shape}"

            # Apply temporal decay (Ebbinghaus forgetting)
            decayed_usage = self.apply_decay(old_usage)

        # Step 2: Compute current activation magnitude
        activation = self.compute_activation(lora_weights)  # (B, N, r)

        # Step 3: Exponential moving average (Hebbian consolidation)
        # new_usage = α * activation + (1 - α) * decayed_usage
        new_usage = (
                self.moving_average_alpha * activation +
                (1 - self.moving_average_alpha) * decayed_usage
        )

        # Step 4: Normalize if requested
        if normalize:
            new_usage = self.normalize(new_usage)

        # Update internal statistics
        self.update_count += 1

        return new_usage

    @torch.no_grad()
    def reset(self, batch_size: int, device: torch.device, dtype: torch.dtype):
        """
        Reset usage tracker to initial state.

        Args:
            batch_size: Batch size for new usage tensor
            device: Target device
            dtype: Target dtype
        """
        self.usage_stats = torch.zeros(
            batch_size, self.num_layers, self.rank,
            device=device, dtype=dtype
        )
        self.update_count = torch.tensor(0, device=device)

    @torch.no_grad()
    def get_statistics(self, usage: torch.Tensor) -> dict:
        """
        Compute diagnostic statistics for monitoring.

        Args:
            usage: Usage vector (B, N, r)

        Returns:
            Dictionary with statistics
        """
        with torch.no_grad():
            return {
                'mean': usage.mean().item(),
                'std': usage.std().item(),
                'min': usage.min().item(),
                'max': usage.max().item(),
                'sparsity': (usage < 0.1).float().mean().item(),  # % of "cold" ranks
                'hot_ratio': (usage > 0.7).float().mean().item(),  # % of "hot" ranks
            }

    def extra_repr(self) -> str:
        """String representation for debugging"""
        return (
            f"num_layers={self.num_layers}, "
            f"rank={self.rank}, "
            f"γ={self.decay_gamma}, "
            f"α={self.moving_average_alpha}"
        )