models/mocae.py

"""Mixture of Calibrated Experts (MoCaE) — Stage 2 of AlignX.

Router  : α = softmax(W_r · h_q + b_r)
Experts : 3 dedicated FFN heads (helpful, harmless, honest)
Fractal : FD_a = log(N) / log(1/ε) on sparse/dense activations from T_a
Natural : k-means on token-level activations from E_{a*}(h_q)
Score   : s_a = λ1·FD_a + λ2·score(N_a),  ŝ_a = softmax(s)
Output  : h_final = Σ ŝ_a · E_a(h_q)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Optional

from .fractal_calibrator import FractalCalibrator
from .natural_calibrator import NaturalCalibrator


def _dequantize_ffn_weights(ffn_module) -> dict:
    """Return a float16 state dict from an FFN that may be 4-bit quantized."""
    sd = {}
    for name, param in ffn_module.named_parameters():
        if hasattr(param, 'quant_state'):
            # bitsandbytes Params4bit — dequantize to float16
            import bitsandbytes as bnb
            dq = bnb.functional.dequantize_4bit(param.data, param.quant_state).to(torch.float16)
            sd[name] = dq
        else:
            sd[name] = param.data.to(torch.float16)
    return sd


class ExpertFFN(nn.Module):
    """Single alignment-axis expert: a gated FFN (SiLU gate, identical to LLaMA MLP)."""

    def __init__(self, hidden_dim: int, intermediate_dim: int):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_dim, intermediate_dim, bias=False)
        self.up_proj   = nn.Linear(hidden_dim, intermediate_dim, bias=False)
        self.down_proj = nn.Linear(intermediate_dim, hidden_dim, bias=False)
        self.act = nn.SiLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x))


class MoCaE(nn.Module):
    """Mixture of Calibrated Experts layer.

    Replaces the final FFN block in the transformer decoder.
    Expects the 3 task-feature matrices T_a (precomputed in Stage 1) to be
    registered as non-trainable buffers before training Stage 2.
    """

    AXES = ["helpful", "harmless", "honest"]

    def __init__(
        self,
        hidden_dim: int,
        intermediate_dim: int,
        k: int = 256,
        num_experts: int = 3,
        top_k: int = 1,
        lambda1: float = 0.6,
        lambda2: float = 0.4,
        temperature: float = 1.0,
        epsilon: float = 0.05,
        n_clusters: int = 8,
    ):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.intermediate_dim = intermediate_dim
        self.k = k
        self.num_experts = num_experts
        self.top_k = top_k
        self.lambda1 = lambda1
        self.lambda2 = lambda2
        self.temperature = temperature

        # Router: W_r ∈ R^{3×d}, b_r ∈ R^3
        self.router = nn.Linear(hidden_dim, num_experts, bias=True)

        # Three expert FFN heads (one per alignment axis)
        self.experts = nn.ModuleList(
            [ExpertFFN(hidden_dim, intermediate_dim) for _ in range(num_experts)]
        )

        # Calibrators
        self.fractal_cal = FractalCalibrator(epsilon=epsilon, latent_dim=k)
        self.natural_cal = NaturalCalibrator(K=n_clusters, hidden_dim=hidden_dim)

        # Task-feature matrix buffers (set via register_task_matrices)
        # T_a: (k,) for each axis a
        for i, axis in enumerate(self.AXES):
            self.register_buffer(f"T_{axis}", torch.zeros(k))

        self._task_matrices_set = False

    # ------------------------------------------------------------------
    # Public helpers
    # ------------------------------------------------------------------

    def register_task_matrices(self, T_helpful, T_harmless, T_honest):
        """Register precomputed task-feature matrices as non-trainable buffers."""
        self.T_helpful.copy_(T_helpful)
        self.T_harmless.copy_(T_harmless)
        self.T_honest.copy_(T_honest)
        self._task_matrices_set = True

    def get_task_matrices(self) -> List[torch.Tensor]:
        return [self.T_helpful, self.T_harmless, self.T_honest]

    def init_experts_from_ffn(self, ffn_module):
        """Initialise all three experts from an existing FFN module's weights.

        Called at build time so experts start from the fine-tuned model's
        final FFN layer rather than from scratch.
        """
        src = _dequantize_ffn_weights(ffn_module)
        for expert in self.experts:
            expert.load_state_dict(src, strict=False)

    def init_expert_from_ffn(self, expert_idx: int, ffn_module):
        """Initialise a single expert from a specific fine-tuned model's FFN."""
        src = _dequantize_ffn_weights(ffn_module)
        self.experts[expert_idx].load_state_dict(src, strict=False)

    # ------------------------------------------------------------------
    # Forward pass
    # ------------------------------------------------------------------

    def forward(
        self,
        h_q: torch.Tensor,
        use_calibration: bool = True,
        use_fractal: bool = True,
        use_natural: bool = True,
    ) -> torch.Tensor:
        """
        h_q: (batch, seq_len, hidden_dim)
        Returns h_final: (batch, seq_len, hidden_dim)
        """
        batch, seq_len, d = h_q.shape

        # --- Router ---
        # Use mean-pooled query representation for routing
        h_mean = h_q.mean(dim=1)  # (batch, hidden_dim)
        alpha = F.softmax(self.router(h_mean) / self.temperature, dim=-1)  # (batch, 3)

        # --- Expert outputs ---
        expert_outputs = []
        for expert in self.experts:
            z_a = expert(h_q)  # (batch, seq_len, hidden_dim)
            expert_outputs.append(z_a)

        if not use_calibration:
            # Plain routing: weighted sum by router
            expert_stack = torch.stack(expert_outputs, dim=-1)  # (B, S, D, 3)
            h_final = (expert_stack * alpha[:, None, None, :]).sum(-1)
            return h_final

        # --- Calibration scores ---
        task_matrices = self.get_task_matrices()  # [T_helpful, T_harmless, T_honest]
        cal_scores = []
        for i in range(self.num_experts):
            s_a = torch.zeros(1, device=h_q.device, dtype=h_q.dtype)
            if use_fractal and self._task_matrices_set:
                fd = self.fractal_cal(task_matrices[i].to(h_q.device))
                s_a = s_a + self.lambda1 * fd
            if use_natural:
                nat = self.natural_cal(expert_outputs[i])
                s_a = s_a + self.lambda2 * nat
            if not (use_fractal or use_natural):
                s_a = alpha[:, i].mean().unsqueeze(0)
            cal_scores.append(s_a)

        # s_hat: normalize via softmax  (3,)
        s_vec = torch.stack(cal_scores).squeeze()  # (3,)
        s_hat = F.softmax(s_vec, dim=0)           # (3,)

        # --- Weighted sum of expert outputs ---
        expert_stack = torch.stack(expert_outputs, dim=-1)  # (B, S, D, 3)
        h_final = (expert_stack * s_hat[None, None, None, :]).sum(-1)

        return h_final


class AlignXLayer(nn.Module):
    """Drop-in replacement for a transformer FFN layer.

    Wraps MoCaE and adds a residual connection + optional layer norm.
    """

    def __init__(self, mocae: MoCaE, hidden_dim: int, eps: float = 1e-5):
        super().__init__()
        self.mocae = mocae
        self.layer_norm = nn.LayerNorm(hidden_dim, eps=eps)

    def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
        residual = hidden_states
        out = self.mocae(hidden_states)
        out = self.layer_norm(out + residual)
        return out