models/parametric_encoder.py

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


class ParametricEncoder(nn.Module):
    """
    Compresses LoRA weight matrices into compact feature vectors using 1D-CNN.
    Captures local spatial patterns in weight matrices for efficient processing.

    Architecture:
        Input: (B, N, r, d) -> Reshape to (B*N, r, d)
        -> Conv1D + BatchNorm1D layers -> AdaptiveAvgPool1D -> Linear projection
        Output: (B, N, s_param)
    """

    def __init__(
            self,
            rank: int,
            hidden_dim: int,
            state_dim: int = 256,
            num_layers: int = 8
    ):
        super().__init__()

        self.rank = rank
        self.hidden_dim = hidden_dim
        self.state_dim = state_dim
        self.num_layers = num_layers

        # 1D-CNN architecture with BatchNorm1d (standard for Conv1d)
        self.cnn = nn.Sequential(
            # First conv layer: extract local patterns
            nn.Conv1d(
                in_channels=rank,
                out_channels=64,
                kernel_size=3,
                padding=1,
                bias=False  # Bias not needed before BatchNorm
            ),
            nn.BatchNorm1d(64),
            nn.GELU(),

            # Second conv layer: higher-level patterns
            nn.Conv1d(
                in_channels=64,
                out_channels=128,
                kernel_size=3,
                padding=1,
                bias=False
            ),
            nn.BatchNorm1d(128),
            nn.GELU(),

            # Third conv layer: abstract features
            nn.Conv1d(
                in_channels=128,
                out_channels=256,
                kernel_size=3,
                padding=1,
                bias=False
            ),
            nn.BatchNorm1d(256),
            nn.GELU(),

            # Global average pooling: compress to single vector
            nn.AdaptiveAvgPool1d(1)
        )

        # Final projection to desired state dimension
        self.projection = nn.Sequential(
            nn.Linear(256, state_dim),
            nn.LayerNorm(state_dim),  # ✅ LayerNorm OK here (after Linear)
            nn.GELU()
        )

        self._init_weights()

        print(f"✅ Parametric Encoder initialized:")
        print(f"   - Input: ({rank}, {hidden_dim}) per layer")
        print(f"   - Output: {state_dim} features per layer")
        print(f"   - Total params: {sum(p.numel() for p in self.parameters()) / 1e6:.2f}M")

    def _init_weights(self):
        """Initialize weights using appropriate initialization strategies"""
        for m in self.modules():
            if isinstance(m, nn.Conv1d):
                # Xavier works well with GELU
                nn.init.xavier_normal_(m.weight)
                if hasattr(m, 'bias') and m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, (nn.BatchNorm1d, nn.LayerNorm)):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)

    def forward(self, lora_weights: torch.Tensor) -> torch.Tensor:
        """
        Compress LoRA weight matrices into feature vectors.

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

        Returns:
            features: Tensor of shape (B, N, s_param) or (N, s_param)
        """
        # Handle both 3D and 4D inputs
        if lora_weights.dim() == 3:
            batch_size = 1
            N, r, d = lora_weights.shape
            x = lora_weights.unsqueeze(0)
        else:
            batch_size, N, r, d = lora_weights.shape
            x = lora_weights

        # Validate dimensions
        assert r == self.rank, f"Expected rank {self.rank}, got {r}"
        assert d == self.hidden_dim, f"Expected hidden_dim {self.hidden_dim}, got {d}"
        assert N == self.num_layers, f"Expected num_layers {self.num_layers}, got {N}"

        # Reshape: (B, N, r, d) -> (B*N, r, d) for Conv1d
        x = x.reshape(batch_size * N, r, d)

        # Pass through CNN
        x = self.cnn(x)  # (B*N, 256, 1)

        # Squeeze last dimension
        x = x.squeeze(-1)  # (B*N, 256)

        # Project to state dimension
        x = self.projection(x)  # (B*N, s_param)

        # Reshape back: (B*N, s_param) -> (B, N, s_param)
        x = x.reshape(batch_size, N, self.state_dim)

        if lora_weights.dim() == 3:
            x = x.squeeze(0)

        return x

    def get_output_dim(self) -> int:
        """Get output feature dimension"""
        return self.state_dim