models/hypernetwork.py

# models/hypernetwork.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple


class HypernetworkCore(nn.Module):
    """
    Core Hypernetwork that generates Delta LoRA weights with Hebbian modulation.

    Architecture:
    1. Input Fusion: Concatenate context, parametric, and usage vectors
    2. Shared Backbone: Extract high-level features
    3. Delta Generation: Generate free-form Delta_W
    4. Hebbian Alignment: Modulate Delta_W based on usage frequency (per-rank)

    Hebbian Mechanism:
    - High usage (hot ranks) → Delta_W aligns with W_old → accumulation
    - Low usage (cold ranks) → Delta_W is free → selective overwriting

    This enables "use it or lose it" parametric memory evolution.
    """

    def __init__(
            self,
            context_dim: int,
            param_dim: int,
            usage_dim: int,
            rank: int,
            hidden_dim: int,
            num_layers: int,
            hidden_state_dim: int = 256,
            alignment_scale: float = 0.1
    ):
        super().__init__()

        self.context_dim = context_dim
        self.param_dim = param_dim
        self.usage_dim = usage_dim
        self.rank = rank
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.alignment_scale = alignment_scale

        # Input dimension after concatenation
        input_dim = context_dim + param_dim + usage_dim  # 384 + 256 + 16 = 656

        print(f"✅ Hypernetwork Core initialized:")
        print(f"   - Input dims: context={context_dim}, param={param_dim}, usage={usage_dim}")
        print(f"   - Output: ({rank}, {hidden_dim}) per layer")
        print(f"   - Hidden state: {hidden_state_dim}")

        # Shared backbone MLP (processes all layers)
        self.backbone = nn.Sequential(
            nn.Linear(input_dim, hidden_state_dim * 2),
            nn.LayerNorm(hidden_state_dim * 2),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_state_dim * 2, hidden_state_dim),
            nn.LayerNorm(hidden_state_dim),
            nn.GELU(),
            nn.Dropout(0.1)
        )

        # Delta generation head (generates free-form Delta_W)
        self.delta_head = nn.Sequential(
            nn.Linear(hidden_state_dim, hidden_state_dim),
            nn.GELU(),
            nn.Linear(hidden_state_dim, rank * hidden_dim)
        )

        # Each rank gets its own alignment weight based on its usage
        # Input: usage scalar per rank (1) -> Output: alignment weight (1)
        self.alignment_head = nn.Sequential(
            nn.Linear(1, 32),  # Each rank has 1 usage value
            nn.LayerNorm(32),
            nn.GELU(),
            nn.Linear(32, 1),
            nn.Sigmoid()  # Output in [0, 1]
        )

        # Initialize weights
        self._init_weights()

        # Print parameter count
        total_params = sum(p.numel() for p in self.parameters())
        print(f"   - Total params: {total_params / 1e6:.2f}M")

    def _init_weights(self):
        """
        Initialize weights with Zero-Bias for delta_head.
        This ensures Delta_W = 0 at initialization, so Active LoRA = Old LoRA.
        """
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

        # CRITICAL: Zero-bias for delta_head to ensure Delta_W = 0 initially
        for module in reversed(list(self.delta_head.modules())):
            if isinstance(module, nn.Linear):
                nn.init.zeros_(module.bias)
                print(f"   ✅ Zero-bias initialized for delta_head output layer")
                break

    def forward(
            self,
            v_ctx: torch.Tensor,
            v_old: torch.Tensor,
            u: torch.Tensor,
            w_old: torch.Tensor
    ) -> torch.Tensor:
        """
        Generate Delta LoRA with Hebbian modulation.

        Args:
            v_ctx: Context embedding (B, context_dim)
            v_old: Parametric embedding (B, N, param_dim)
            u: Usage vector (B, N, rank)
            w_old: Old LoRA weights (B, N, rank, hidden_dim)

        Returns:
            delta_W: Delta LoRA weights (B, N, rank, hidden_dim)
        """
        B, N, _ = v_old.shape

        # Validate inputs
        assert v_ctx.shape == (B, self.context_dim)
        assert v_old.shape == (B, N, self.param_dim)
        assert u.shape == (B, N, self.rank)
        assert w_old.shape == (B, N, self.rank, self.hidden_dim)

        # Step 1: Expand context vector to match layer dimension
        v_ctx_expanded = v_ctx.unsqueeze(1).expand(-1, N, -1)  # (B, N, context_dim)

        # Step 2: Concatenate all inputs
        combined = torch.cat([v_ctx_expanded, v_old, u], dim=-1)  # (B, N, 656)

        # Step 3: Pass through shared backbone
        features = self.backbone(combined)  # (B, N, hidden_state_dim)

        # Step 4: Generate free-form Delta_W
        delta_W_free = self.delta_head(features)  # (B, N, rank * hidden_dim)
        delta_W_free = delta_W_free.view(B, N, self.rank, self.hidden_dim)  # (B, N, r, d)

        # Step 5: Generate Hebbian alignment weight PER RANK
        u_expanded = u.unsqueeze(-1)  # (B, N, rank, 1)
        alignment = self.alignment_head(u_expanded)  # (B, N, rank, 1)

        # Step 6: Normalize W_old to get direction (unit vectors)
        w_old_direction = F.normalize(w_old, dim=-1)  # (B, N, r, d), magnitude = 1

        # Step 7: Scale W_old direction to match Delta_W magnitude
        with torch.no_grad():
            delta_magnitude = delta_W_free.norm(dim=-1, keepdim=True).mean(dim=-2, keepdim=True)

        w_old_scaled = w_old_direction * delta_magnitude * self.alignment_scale

        # Step 8: Hebbian combination
        # Now shapes match: (B, N, rank, 1) * (B, N, rank, hidden_dim)
        delta_W_final = alignment * w_old_scaled + (1 - alignment) * delta_W_free

        return delta_W_final

    def get_alignment_weights(self, u: torch.Tensor) -> torch.Tensor:
        """
        Get alignment weights for diagnostic purposes.

        Args:
            u: Usage vector (B, N, rank)

        Returns:
            alignment: Alignment weights (B, N, rank, 1)
        """
        with torch.no_grad():
            u_expanded = u.unsqueeze(-1)  # (B, N, rank, 1)
            return self.alignment_head(u_expanded)

    def extra_repr(self) -> str:
        """String representation for debugging"""
        return (
            f"context_dim={self.context_dim}, "
            f"param_dim={self.param_dim}, "
            f"rank={self.rank}, "
            f"hidden_dim={self.hidden_dim}"
        )