feat: 回归旁路（regression bypass）+ --reg-bypass 开关

- fusion.py: ResidualConcatFusion 额外返回 pooled（纯 chem+tab 注意力池化）
- models.py: _backbone_from_stacked 存 _last_pooled；forward 按 reg_bypass 决定是否传入
- heads.py: MultiTaskHead 接收 x_numeric，回归头(size/delivery)走纯数值，分类/biodist 走 fused
- nested_cv_optuna.py: 新增 --reg-bypass 开关（typer），经 llm_kwargs 传入

fix: LNPModelWithoutMPNN 未将 reg_bypass 转发给 super().__init__()，
导致开关静默失效、on/off 产生完全相同结果（已验证修复：off→False, on→True）

lnp_ml/modeling/heads.py

@@ -1,125 +1,126 @@
-import torch
+import torch
-import torch.nn as nn
+import torch.nn as nn
-import torch.nn.functional as F
+import torch.nn.functional as F
-from typing import Dict
+from typing import Dict
-
+
-
+
-class RegressionHead(nn.Module):
+class RegressionHead(nn.Module):
-    """回归任务头：输出单个 float 值"""
+    """回归任务头：输出单个 float 值"""
-
+
-    def __init__(self, in_dim: int, hidden_dim: int = 128, dropout: float = 0.1) -> None:
+    def __init__(self, in_dim: int, hidden_dim: int = 128, dropout: float = 0.1) -> None:
-        super().__init__()
+        super().__init__()
-        self.net = nn.Sequential(
+        self.net = nn.Sequential(
-            nn.Linear(in_dim, hidden_dim),
+            nn.Linear(in_dim, hidden_dim),
-            nn.ReLU(),
+            nn.ReLU(),
-            nn.Dropout(dropout),
+            nn.Dropout(dropout),
-            nn.Linear(hidden_dim, 1),
+            nn.Linear(hidden_dim, 1),
-        )
+        )
-
+
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """[B, in_dim] -> [B, 1]"""
+        """[B, in_dim] -> [B, 1]"""
-        return self.net(x)
+        return self.net(x)
-
+
-
+
-class ClassificationHead(nn.Module):
+class ClassificationHead(nn.Module):
-    """分类任务头：输出 logits"""
+    """分类任务头：输出 logits"""
-
+
-    def __init__(
+    def __init__(
-        self, in_dim: int, num_classes: int, hidden_dim: int = 128, dropout: float = 0.1
+        self, in_dim: int, num_classes: int, hidden_dim: int = 128, dropout: float = 0.1
-    ) -> None:
+    ) -> None:
-        super().__init__()
+        super().__init__()
-        self.net = nn.Sequential(
+        self.net = nn.Sequential(
-            nn.Linear(in_dim, hidden_dim),
+            nn.Linear(in_dim, hidden_dim),
-            nn.ReLU(),
+            nn.ReLU(),
-            nn.Dropout(dropout),
+            nn.Dropout(dropout),
-            nn.Linear(hidden_dim, num_classes),
+            nn.Linear(hidden_dim, num_classes),
-        )
+        )
-
+
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """[B, in_dim] -> [B, num_classes] (logits)"""
+        """[B, in_dim] -> [B, num_classes] (logits)"""
-        return self.net(x)
+        return self.net(x)
-
+
-
+
-class DistributionHead(nn.Module):
+class DistributionHead(nn.Module):
-    """分布任务头：输出和为 1 的概率分布（用于 Biodistribution）"""
+    """分布任务头：输出和为 1 的概率分布（用于 Biodistribution）"""
-
+
-    def __init__(
+    def __init__(
-        self, in_dim: int, num_outputs: int, hidden_dim: int = 128, dropout: float = 0.1
+        self, in_dim: int, num_outputs: int, hidden_dim: int = 128, dropout: float = 0.1
-    ) -> None:
+    ) -> None:
-        super().__init__()
+        super().__init__()
-        self.net = nn.Sequential(
+        self.net = nn.Sequential(
-            nn.Linear(in_dim, hidden_dim),
+            nn.Linear(in_dim, hidden_dim),
-            nn.ReLU(),
+            nn.ReLU(),
-            nn.Dropout(dropout),
+            nn.Dropout(dropout),
-            nn.Linear(hidden_dim, num_outputs),
+            nn.Linear(hidden_dim, num_outputs),
-        )
+        )
-
+
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """[B, in_dim] -> [B, num_outputs] (softmax, sum=1)"""
+        """[B, in_dim] -> [B, num_outputs] (softmax, sum=1)"""
-        logits = self.net(x)
+        logits = self.net(x)
-        return F.softmax(logits, dim=-1)
+        return F.softmax(logits, dim=-1)
-
+
-
+
-class MultiTaskHead(nn.Module):
+class MultiTaskHead(nn.Module):
-    """
+    """
-    多任务预测头，根据任务配置自动创建对应的子头。
+    多任务预测头，根据任务配置自动创建对应的子头。
-
+
-    输出:
+    输出:
-        - size: [B, 1] 回归
+        - size: [B, 1] 回归
-        - pdi: [B, 2] 分类 logits
+        - pdi: [B, 2] 分类 logits
-        - ee: [B, 3] 分类 logits
+        - ee: [B, 3] 分类 logits
-        - delivery: [B, 1] 回归
+        - delivery: [B, 1] 回归
-        - biodist: [B, 7] softmax 分布
+        - biodist: [B, 7] softmax 分布
-        - toxic: [B, 2] 二分类 logits
+        - toxic: [B, 2] 二分类 logits
-    """
+    """
-
+
-    def __init__(self, in_dim: int, hidden_dim: int = 128, dropout: float = 0.1) -> None:
+    def __init__(self, in_dim: int, hidden_dim: int = 128, dropout: float = 0.1) -> None:
-        super().__init__()
+        super().__init__()
-
+
-        # size: 回归 (log-transformed)
+        # size: 回归 (log-transformed)
-        size_dropout = min(0.5, dropout + 0.2)
+        size_dropout = min(0.5, dropout + 0.2)
-        self.size_head = RegressionHead(in_dim, hidden_dim, size_dropout)
+        self.size_head = RegressionHead(in_dim, hidden_dim, size_dropout)
-
+
-        # PDI: 2 分类
+        # PDI: 2 分类
-        self.pdi_head = ClassificationHead(in_dim, num_classes=2, hidden_dim=hidden_dim, dropout=dropout)
+        self.pdi_head = ClassificationHead(in_dim, num_classes=2, hidden_dim=hidden_dim, dropout=dropout)
-
+
-        # Encapsulation Efficiency: 3 分类
+        # Encapsulation Efficiency: 3 分类
-        self.ee_head = ClassificationHead(in_dim, num_classes=3, hidden_dim=hidden_dim, dropout=dropout)
+        self.ee_head = ClassificationHead(in_dim, num_classes=3, hidden_dim=hidden_dim, dropout=dropout)
-
+
-        # quantified_delivery: 回归 (z-scored)
+        # quantified_delivery: 回归 (z-scored)
-        self.delivery_head = RegressionHead(in_dim, hidden_dim, dropout)
+        self.delivery_head = RegressionHead(in_dim, hidden_dim, dropout)
-
+
-        # Biodistribution: 7 输出，softmax (sum=1)
+        # Biodistribution: 7 输出，softmax (sum=1)
-        self.biodist_head = DistributionHead(in_dim, num_outputs=7, hidden_dim=hidden_dim, dropout=dropout)
+        self.biodist_head = DistributionHead(in_dim, num_outputs=7, hidden_dim=hidden_dim, dropout=dropout)
-
+
-        # toxic: 二分类
+        # toxic: 二分类
-        self.toxic_head = ClassificationHead(in_dim, num_classes=2, hidden_dim=hidden_dim, dropout=dropout)
+        self.toxic_head = ClassificationHead(in_dim, num_classes=2, hidden_dim=hidden_dim, dropout=dropout)
-
+
-        # 不确定性加权（Kendall 2018）：每个任务一个可学习 log σ²，初始 0
+        # 不确定性加权（Kendall 2018）：每个任务一个可学习 log σ²，初始 0
-        self.log_vars = nn.ParameterDict({
+        self.log_vars = nn.ParameterDict({
-            t: nn.Parameter(torch.zeros(()))
+            t: nn.Parameter(torch.zeros(()))
-            for t in ["size", "delivery", "pdi", "ee", "toxic", "biodist"]
+            for t in ["size", "delivery", "pdi", "ee", "toxic", "biodist"]
-        })
+        })
-
+
-    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
+    def forward(self, x: torch.Tensor, x_numeric: torch.Tensor = None) -> Dict[str, torch.Tensor]:
-        """
+        """
-        Args:
+        Args:
-            x: [B, in_dim] fusion 层输出
+            x: [B, in_dim] fusion 层输出
-
+
-        Returns:
+        Returns:
-            Dict with keys:
+            Dict with keys:
-                - "size": [B, 1]
+                - "size": [B, 1]
-                - "pdi": [B, 2] logits
+                - "pdi": [B, 2] logits
-                - "ee": [B, 3] logits
+                - "ee": [B, 3] logits
-                - "delivery": [B, 1]
+                - "delivery": [B, 1]
-                - "biodist": [B, 7] probabilities (sum=1)
+                - "biodist": [B, 7] probabilities (sum=1)
-                - "toxic": [B, 2] logits
+                - "toxic": [B, 2] logits
-        """
+        """
-        return {
+        xn = x if x_numeric is None else x_numeric
-            "size": self.size_head(x),
+        return {
-            "pdi": self.pdi_head(x),
+            "size": self.size_head(xn),
-            "ee": self.ee_head(x),
+            "pdi": self.pdi_head(x),
-            "delivery": self.delivery_head(x),
+            "ee": self.ee_head(x),
-            "biodist": self.biodist_head(x),
+            "delivery": self.delivery_head(xn),
-            "toxic": self.toxic_head(x),
+            "biodist": self.biodist_head(x),
-        }
+            "toxic": self.toxic_head(x),
+        }

lnp_ml/modeling/layers/fusion.py

@@ -153,4 +153,7 @@ class ResidualConcatFusion(nn.Module):
         if f_retr is not None:
             out = out + self.g_retr * f_retr  # 检索旁路，零初始化门保证起点=不开
 
-        return (out, attn) if return_attn_weights else out
+        # 回归旁路：额外返回 pooled（纯 chem+tab，不含 f_llm/f_moe/f_retr）
+        if return_attn_weights:
+            return out, pooled, attn
+        return out, pooled

lnp_ml/modeling/models.py

@@ -99,6 +99,8 @@ class LNPModel(nn.Module):
         moe_top_k: int = 2,
         moe_expert_hidden_mult: int = 2,
         moe_jitter_noise: float = 0.0,
+        # ============ 回归旁路开关 ============
+        reg_bypass: str = "on",
         # ============ LLM 相关 ============
         use_llm: bool = False,
         llm_model_path: str = DEFAULT_MOLT5_PATH,
@@ -174,6 +176,7 @@ class LNPModel(nn.Module):
 
         # ============ LLM Prompt (可选) ============
         self.use_llm = use_llm
+        self.reg_bypass = (str(reg_bypass).lower() == "on")
         if use_llm:
             self.llm_prompt = LLMPromptEncoder(
                 d_model=d_model,
@@ -270,7 +273,9 @@ class LNPModel(nn.Module):
             _feats_t = torch.as_tensor(_feats, dtype=chem.dtype, device=chem.device)
             f_retr = self.retr_proj(_feats_t)
 
-        return self.fusion(chem, tab, f_moe=f_moe, f_llm=f_llm, f_retr=f_retr)
+        fused, pooled = self.fusion(chem, tab, f_moe=f_moe, f_llm=f_llm, f_retr=f_retr)
+        self._last_pooled = pooled  # 纯数值向量，供回归 head 使用
+        return fused
 
     def forward_from_projected(
         self,
@@ -357,7 +362,7 @@ class LNPModel(nn.Module):
             delivery [B,1], biodist [B,7], toxic [B,2]
         """
         fused = self.forward_backbone(smiles, tabular)
-        return self.head(fused)
+        return self.head(fused, getattr(self, "_last_pooled", None) if getattr(self, "reg_bypass", True) else None)
 
     def clear_cache(self) -> None:
         """清空所有 encoder 的缓存"""
@@ -443,6 +448,8 @@ class LNPModelWithoutMPNN(LNPModel):
         moe_top_k: int = 2,
         moe_expert_hidden_mult: int = 2,
         moe_jitter_noise: float = 0.0,
+        # ============ 回归旁路开关 ============
+        reg_bypass: str = "on",
         # ============ LLM 相关 ============
         use_llm: bool = False,
         llm_model_path: str = DEFAULT_MOLT5_PATH,
@@ -472,6 +479,7 @@ class LNPModelWithoutMPNN(LNPModel):
             mpnn_checkpoint=None,
             mpnn_ensemble_paths=None,
             input_dims=dims,
+            reg_bypass=reg_bypass,
             use_moe=use_moe,
             moe_n_experts=moe_n_experts,
             moe_top_k=moe_top_k,

lnp_ml/modeling/nested_cv_optuna.py

@@ -930,6 +930,8 @@ def main(
     moe_jitter_noise: float = 0.0,
     # Set Transformer
     set_transformer_block: str = "sab",
+    # 回归旁路（消融开关）
+    reg_bypass: str = "on",
     # LLM（消融开关）
     use_llm: bool = False,
     use_rag: bool = False,
@@ -967,6 +969,7 @@ def main(
     set_global_seed(seed)
 
     llm_kwargs = dict(
+        reg_bypass=reg_bypass,
         use_rag=use_rag,
         rag_top_k=rag_top_k,
         use_llm=use_llm,