diff --git a/lnp_ml/modeling/heads.py b/lnp_ml/modeling/heads.py index 39e655d..2c4c89a 100644 --- a/lnp_ml/modeling/heads.py +++ b/lnp_ml/modeling/heads.py @@ -1,125 +1,126 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -from typing import Dict - - -class RegressionHead(nn.Module): - """回归任务头：输出单个 float 值""" - - def __init__(self, in_dim: int, hidden_dim: int = 128, dropout: float = 0.1) -> None: - super().__init__() - self.net = nn.Sequential( - nn.Linear(in_dim, hidden_dim), - nn.ReLU(), - nn.Dropout(dropout), - nn.Linear(hidden_dim, 1), - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """[B, in_dim] -> [B, 1]""" - return self.net(x) - - -class ClassificationHead(nn.Module): - """分类任务头：输出 logits""" - - def __init__( - self, in_dim: int, num_classes: int, hidden_dim: int = 128, dropout: float = 0.1 - ) -> None: - super().__init__() - self.net = nn.Sequential( - nn.Linear(in_dim, hidden_dim), - nn.ReLU(), - nn.Dropout(dropout), - nn.Linear(hidden_dim, num_classes), - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """[B, in_dim] -> [B, num_classes] (logits)""" - return self.net(x) - - -class DistributionHead(nn.Module): - """分布任务头：输出和为 1 的概率分布（用于 Biodistribution）""" - - def __init__( - self, in_dim: int, num_outputs: int, hidden_dim: int = 128, dropout: float = 0.1 - ) -> None: - super().__init__() - self.net = nn.Sequential( - nn.Linear(in_dim, hidden_dim), - nn.ReLU(), - nn.Dropout(dropout), - nn.Linear(hidden_dim, num_outputs), - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """[B, in_dim] -> [B, num_outputs] (softmax, sum=1)""" - logits = self.net(x) - return F.softmax(logits, dim=-1) - - -class MultiTaskHead(nn.Module): - """ - 多任务预测头，根据任务配置自动创建对应的子头。 - - 输出: - - size: [B, 1] 回归 - - pdi: [B, 2] 分类 logits - - ee: [B, 3] 分类 logits - - delivery: [B, 1] 回归 - - biodist: [B, 7] softmax 分布 - - toxic: [B, 2] 二分类 logits - """ - - def __init__(self, in_dim: int, hidden_dim: int = 128, dropout: float = 0.1) -> None: - super().__init__() - - # size: 回归 (log-transformed) - size_dropout = min(0.5, dropout + 0.2) - self.size_head = RegressionHead(in_dim, hidden_dim, size_dropout) - - # PDI: 2 分类 - self.pdi_head = ClassificationHead(in_dim, num_classes=2, hidden_dim=hidden_dim, dropout=dropout) - - # Encapsulation Efficiency: 3 分类 - self.ee_head = ClassificationHead(in_dim, num_classes=3, hidden_dim=hidden_dim, dropout=dropout) - - # quantified_delivery: 回归 (z-scored) - self.delivery_head = RegressionHead(in_dim, hidden_dim, dropout) - - # Biodistribution: 7 输出，softmax (sum=1) - self.biodist_head = DistributionHead(in_dim, num_outputs=7, hidden_dim=hidden_dim, dropout=dropout) - - # toxic: 二分类 - self.toxic_head = ClassificationHead(in_dim, num_classes=2, hidden_dim=hidden_dim, dropout=dropout) - - # 不确定性加权（Kendall 2018）：每个任务一个可学习 log σ²，初始 0 - self.log_vars = nn.ParameterDict({ - t: nn.Parameter(torch.zeros(())) - for t in ["size", "delivery", "pdi", "ee", "toxic", "biodist"] - }) - - def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]: - """ - Args: - x: [B, in_dim] fusion 层输出 - - Returns: - Dict with keys: - - "size": [B, 1] - - "pdi": [B, 2] logits - - "ee": [B, 3] logits - - "delivery": [B, 1] - - "biodist": [B, 7] probabilities (sum=1) - - "toxic": [B, 2] logits - """ - return { - "size": self.size_head(x), - "pdi": self.pdi_head(x), - "ee": self.ee_head(x), - "delivery": self.delivery_head(x), - "biodist": self.biodist_head(x), - "toxic": self.toxic_head(x), - } +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import Dict + + +class RegressionHead(nn.Module): + """回归任务头：输出单个 float 值""" + + def __init__(self, in_dim: int, hidden_dim: int = 128, dropout: float = 0.1) -> None: + super().__init__() + self.net = nn.Sequential( + nn.Linear(in_dim, hidden_dim), + nn.ReLU(), + nn.Dropout(dropout), + nn.Linear(hidden_dim, 1), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """[B, in_dim] -> [B, 1]""" + return self.net(x) + + +class ClassificationHead(nn.Module): + """分类任务头：输出 logits""" + + def __init__( + self, in_dim: int, num_classes: int, hidden_dim: int = 128, dropout: float = 0.1 + ) -> None: + super().__init__() + self.net = nn.Sequential( + nn.Linear(in_dim, hidden_dim), + nn.ReLU(), + nn.Dropout(dropout), + nn.Linear(hidden_dim, num_classes), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """[B, in_dim] -> [B, num_classes] (logits)""" + return self.net(x) + + +class DistributionHead(nn.Module): + """分布任务头：输出和为 1 的概率分布（用于 Biodistribution）""" + + def __init__( + self, in_dim: int, num_outputs: int, hidden_dim: int = 128, dropout: float = 0.1 + ) -> None: + super().__init__() + self.net = nn.Sequential( + nn.Linear(in_dim, hidden_dim), + nn.ReLU(), + nn.Dropout(dropout), + nn.Linear(hidden_dim, num_outputs), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """[B, in_dim] -> [B, num_outputs] (softmax, sum=1)""" + logits = self.net(x) + return F.softmax(logits, dim=-1) + + +class MultiTaskHead(nn.Module): + """ + 多任务预测头，根据任务配置自动创建对应的子头。 + + 输出: + - size: [B, 1] 回归 + - pdi: [B, 2] 分类 logits + - ee: [B, 3] 分类 logits + - delivery: [B, 1] 回归 + - biodist: [B, 7] softmax 分布 + - toxic: [B, 2] 二分类 logits + """ + + def __init__(self, in_dim: int, hidden_dim: int = 128, dropout: float = 0.1) -> None: + super().__init__() + + # size: 回归 (log-transformed) + size_dropout = min(0.5, dropout + 0.2) + self.size_head = RegressionHead(in_dim, hidden_dim, size_dropout) + + # PDI: 2 分类 + self.pdi_head = ClassificationHead(in_dim, num_classes=2, hidden_dim=hidden_dim, dropout=dropout) + + # Encapsulation Efficiency: 3 分类 + self.ee_head = ClassificationHead(in_dim, num_classes=3, hidden_dim=hidden_dim, dropout=dropout) + + # quantified_delivery: 回归 (z-scored) + self.delivery_head = RegressionHead(in_dim, hidden_dim, dropout) + + # Biodistribution: 7 输出，softmax (sum=1) + self.biodist_head = DistributionHead(in_dim, num_outputs=7, hidden_dim=hidden_dim, dropout=dropout) + + # toxic: 二分类 + self.toxic_head = ClassificationHead(in_dim, num_classes=2, hidden_dim=hidden_dim, dropout=dropout) + + # 不确定性加权（Kendall 2018）：每个任务一个可学习 log σ²，初始 0 + self.log_vars = nn.ParameterDict({ + t: nn.Parameter(torch.zeros(())) + for t in ["size", "delivery", "pdi", "ee", "toxic", "biodist"] + }) + + def forward(self, x: torch.Tensor, x_numeric: torch.Tensor = None) -> Dict[str, torch.Tensor]: + """ + Args: + x: [B, in_dim] fusion 层输出 + + Returns: + Dict with keys: + - "size": [B, 1] + - "pdi": [B, 2] logits + - "ee": [B, 3] logits + - "delivery": [B, 1] + - "biodist": [B, 7] probabilities (sum=1) + - "toxic": [B, 2] logits + """ + xn = x if x_numeric is None else x_numeric + return { + "size": self.size_head(xn), + "pdi": self.pdi_head(x), + "ee": self.ee_head(x), + "delivery": self.delivery_head(xn), + "biodist": self.biodist_head(x), + "toxic": self.toxic_head(x), + } diff --git a/lnp_ml/modeling/layers/fusion.py b/lnp_ml/modeling/layers/fusion.py index 3824eb8..7edeed5 100644 --- a/lnp_ml/modeling/layers/fusion.py +++ b/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 \ No newline at end of file + # 回归旁路：额外返回 pooled（纯 chem+tab，不含 f_llm/f_moe/f_retr） + if return_attn_weights: + return out, pooled, attn + return out, pooled \ No newline at end of file diff --git a/lnp_ml/modeling/models.py b/lnp_ml/modeling/models.py index edd9510..1428075 100644 --- a/lnp_ml/modeling/models.py +++ b/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, diff --git a/lnp_ml/modeling/nested_cv_optuna.py b/lnp_ml/modeling/nested_cv_optuna.py index 362dffe..7c34e63 100644 --- a/lnp_ml/modeling/nested_cv_optuna.py +++ b/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,