lnp_ml/lnp_ml/modeling/layers/fusion.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, List, Literal, Optional, Tuple, Union


PoolingStrategy = Literal["concat", "avg", "max", "attention"]


class FusionLayer(nn.Module):
    """
    将多个 token 融合成单个向量。

    输入: Dict[str, Tensor] 或 [B, n_tokens, d_model]
    输出: [B, fusion_dim]

    策略:
        - concat: [B, n_tokens, d_model] -> [B, n_tokens * d_model]
        - avg: [B, n_tokens, d_model] -> [B, d_model]
        - max: [B, n_tokens, d_model] -> [B, d_model]
        - attention: [B, n_tokens, d_model] -> [B, d_model] (learnable attention pooling)
    """

    def __init__(
        self,
        d_model: int,
        n_tokens: int,
        strategy: PoolingStrategy = "attention",
    ) -> None:
        """
        Args:
            d_model: 每个 token 的维度
            n_tokens: token 数量（如 8）
            strategy: 融合策略
        """
        super().__init__()
        self.d_model = d_model
        self.n_tokens = n_tokens
        self.strategy = strategy

        if strategy == "concat":
            self.fusion_dim = n_tokens * d_model
        else:
            self.fusion_dim = d_model

        # Attention pooling: learnable query
        if strategy == "attention":
            self.attn_query = nn.Parameter(torch.randn(1, 1, d_model))
            self.attn_proj = nn.Linear(d_model, d_model)

    def forward(
        self,
        x: Union[Dict[str, torch.Tensor], torch.Tensor],
        return_attn_weights: bool = False,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        """
        Args:
            x: Dict[str, Tensor] 每个 [B, d_model]，或已 stack 的 [B, n_tokens, d_model]
            return_attn_weights: 若为 True 且策略为 attention，额外返回 attn_weights [B, n_tokens]

        Returns:
            return_attn_weights=False: [B, fusion_dim]
            return_attn_weights=True:  ([B, fusion_dim], [B, n_tokens])
        """
        if isinstance(x, dict):
            x = torch.stack(list(x.values()), dim=1)

        if self.strategy == "concat":
            out = x.flatten(start_dim=1)
            return (out, None) if return_attn_weights else out

        elif self.strategy == "avg":
            out = x.mean(dim=1)
            return (out, None) if return_attn_weights else out

        elif self.strategy == "max":
            out = x.max(dim=1).values
            return (out, None) if return_attn_weights else out

        elif self.strategy == "attention":
            return self._attention_pooling(x, return_attn_weights)

        else:
            raise ValueError(f"Unknown strategy: {self.strategy}")

    def _attention_pooling(
        self, x: torch.Tensor, return_attn_weights: bool = False,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        """
        Attention pooling: 用可学习 query 对 tokens 做加权求和

        Args:
            x: [B, n_tokens, d_model]
            return_attn_weights: 是否返回权重

        Returns:
            return_attn_weights=False: [B, d_model]
            return_attn_weights=True:  ([B, d_model], [B, n_tokens])
        """
        B = x.size(0)
        query = self.attn_query.expand(B, -1, -1)

        keys = self.attn_proj(x)
        scores = torch.bmm(query, keys.transpose(1, 2)) / (self.d_model ** 0.5)
        attn_weights = F.softmax(scores, dim=-1)  # [B, 1, n_tokens]

        out = torch.bmm(attn_weights, x).squeeze(1)  # [B, d_model]

        if return_attn_weights:
            return out, attn_weights.squeeze(1)  # [B, n_tokens]
        return out


class ResidualConcatFusion(nn.Module):
    """对真实 token 做 attention pooling，再用零初始化门把 MoE/LLM 旁路以残差方式加入。

    g_moe / g_llm 初始为 0 → +moe/+llm 起点严格等于 baseline；
    旁路只有确实有用时才会被训练打开，从机制上保证“加了不会更差”。
    """

    def __init__(self, d_model: int, strategy: PoolingStrategy = "attention") -> None:
        super().__init__()
        if strategy == "concat":
            raise ValueError("ResidualConcatFusion 不支持 concat（token 数随开关变化）")
        self.d_model = d_model
        self.pool = FusionLayer(d_model=d_model, n_tokens=1, strategy=strategy)
        self.fusion_dim = self.pool.fusion_dim
        # 零初始化门控（可学习标量），旁路初始不参与
        self.g_moe = nn.Parameter(torch.zeros(()))
        self.g_llm = nn.Parameter(torch.zeros(()))
        self.g_retr = nn.Parameter(torch.zeros(()))  # 检索旁路零初始化门控

    def forward(
        self,
        chem: torch.Tensor,
        tab: torch.Tensor,
        f_moe: Optional[torch.Tensor] = None,
        f_llm: Optional[torch.Tensor] = None,
        f_retr: Optional[torch.Tensor] = None,
        return_attn_weights: bool = False,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        # 只对真实 token（chem + tab）做注意力池化，旁路不参与 softmax 竞争
        seq = torch.cat([chem, tab], dim=1)  # [B, n_chem + n_cond, d_model]
        pooled = self.pool(seq, return_attn_weights=return_attn_weights)
        if return_attn_weights:
            pooled, attn = pooled

        out = pooled
        if f_moe is not None:
            out = out + self.g_moe * f_moe  # 残差 + 零初始化门
        if f_llm is not None:
            out = out + self.g_llm * f_llm
        if f_retr is not None:
            out = out + self.g_retr * f_retr  # 检索旁路，零初始化门保证起点=不开

        # 回归旁路：额外返回 pooled（纯 chem+tab，不含 f_llm/f_moe/f_retr）
        if return_attn_weights:
            return out, pooled, attn
        return out, pooled