feat(training): size 头正则化(dropout+weight_decay)与 size z-score，更新消融汇总

.gitignore

@@ -188,4 +188,11 @@ cython_debug/
 .ruff_cache/
 
 # PyPI configuration file
-.pypirc
+.pypirc
+
+logs/
+*.out
+models/**/*.pt
+models/**/*.png
+models/molt5-base/
+models/abl/

lnp_ml/dataset.py

@@ -118,10 +118,14 @@ def process_dataframe(df: pd.DataFrame) -> pd.DataFrame:
             df[col] = df[col].fillna(0.0).astype(float)
     
     # 5. 处理 target 列
-    # size: 已经 log 过，填充缺失值
+    # size: 已经 log 过；再做全局 z-score，使其 loss 量纲与 delivery 一致
+    #（R² 对仿射变换不变，这里只为多任务不确定性加权更稳，不是 R² 解药）
     if "size" in df.columns:
         df["size"] = pd.to_numeric(df["size"], errors="coerce")
-    
+        _mu, _sd = df["size"].mean(), df["size"].std()
+        if _sd and _sd > 0:
+            df["size"] = (df["size"] - _mu) / _sd
+
     # quantified_delivery: 已经 z-score 过
     if "quantified_delivery" in df.columns:
         df["quantified_delivery"] = pd.to_numeric(df["quantified_delivery"], errors="coerce")
@@ -216,7 +220,8 @@ class LNPDataset(Dataset):
         # PDI: one-hot -> class index
         if all(col in self.df.columns for col in TARGET_CLASSIFICATION_PDI):
             pdi_onehot = self.df[TARGET_CLASSIFICATION_PDI].values
-            self.pdi = np.argmax(pdi_onehot, axis=1).astype(np.int64)
+            pdi_4 = np.argmax(pdi_onehot, axis=1)
+            self.pdi = (pdi_4 >= 1).astype(np.int64)
             self.pdi_valid = pdi_onehot.sum(axis=1) > 0
         else:
             self.pdi = None

lnp_ml/modeling/heads.py

@@ -66,7 +66,7 @@ class MultiTaskHead(nn.Module):
 
     输出:
         - size: [B, 1] 回归
-        - pdi: [B, 4] 分类 logits
+        - pdi: [B, 2] 分类 logits
         - ee: [B, 3] 分类 logits
         - delivery: [B, 1] 回归
         - biodist: [B, 7] softmax 分布
@@ -77,10 +77,11 @@ class MultiTaskHead(nn.Module):
         super().__init__()
 
         # size: 回归 (log-transformed)
-        self.size_head = RegressionHead(in_dim, hidden_dim, dropout)
+        size_dropout = min(0.5, dropout + 0.2)
+        self.size_head = RegressionHead(in_dim, hidden_dim, size_dropout)
 
-        # PDI: 4 分类
+        # PDI: 2 分类
-        self.pdi_head = ClassificationHead(in_dim, num_classes=4, hidden_dim=hidden_dim, dropout=dropout)
+        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)
@@ -108,7 +109,7 @@ class MultiTaskHead(nn.Module):
         Returns:
             Dict with keys:
                 - "size": [B, 1]
-                - "pdi": [B, 4] logits
+                - "pdi": [B, 2] logits
                 - "ee": [B, 3] logits
                 - "delivery": [B, 1]
                 - "biodist": [B, 7] probabilities (sum=1)

lnp_ml/modeling/layers/__init__.py

@@ -1,6 +1,16 @@
 from lnp_ml.modeling.layers.token_projector import TokenProjector
 from lnp_ml.modeling.layers.bidirectional_cross_attention import CrossModalAttention
-from lnp_ml.modeling.layers.fusion import FusionLayer
+from lnp_ml.modeling.layers.set_transformer import SetTransformer
+from lnp_ml.modeling.layers.fusion import FusionLayer, ResidualConcatFusion
 from lnp_ml.modeling.layers.moe import MoEBlock
+from lnp_ml.modeling.layers.llm_prompt import LLMPromptEncoder
 
-__all__ = ["TokenProjector", "CrossModalAttention", "FusionLayer", "MoEBlock"]
+__all__ = [
+    "TokenProjector",
+    "CrossModalAttention",
+    "SetTransformer",
+    "FusionLayer",
+    "ResidualConcatFusion",
+    "MoEBlock",
+    "LLMPromptEncoder",
+]

lnp_ml/modeling/trainer_balanced.py

@@ -29,12 +29,11 @@ class LossWeightsBalanced:
     biodist: float = 1.0
     toxic: float = 1.0
     moe_lb: float = 0.01  # MoE load-balancing 系数（仅在 use_moe=True 时生效）
-    moe_lb: float = 0.01  # MoE load-balancing 系数（仅在 use_moe=True 时生效）
 
 
 def compute_class_weights_from_loader(
     loader: DataLoader,
-    n_pdi_classes: int = 4,
+    n_pdi_classes: int = 2,
     n_ee_classes: int = 3,
     n_toxic_classes: int = 2,
     smoothing: float = 0.1,
@@ -113,7 +112,6 @@ def compute_multitask_loss_balanced(
     task_weights: Optional[LossWeightsBalanced] = None,
     class_weights: Optional[ClassWeights] = None,
     model: Optional[nn.Module] = None,
-    model: Optional[nn.Module] = None,
 ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
     """
     计算带类权重的多任务损失。
@@ -205,7 +203,6 @@ def train_epoch_balanced(
     model.train()
     total_loss = 0.0
     task_losses = {k: 0.0 for k in ["size", "pdi", "ee", "delivery", "biodist", "toxic", "moe_lb"]}
-    task_losses = {k: 0.0 for k in ["size", "pdi", "ee", "delivery", "biodist", "toxic", "moe_lb"]}
     n_batches = 0
     
     for batch in tqdm(loader, desc="Training", leave=False):
@@ -218,7 +215,6 @@ def train_epoch_balanced(
         outputs = model(smiles, tabular)
         loss, losses = compute_multitask_loss_balanced(
             outputs, targets, mask, task_weights, class_weights, model=model,
-            outputs, targets, mask, task_weights, class_weights, model=model,
         )
         
         loss.backward()
@@ -248,7 +244,6 @@ def validate_balanced(
     model.eval()
     total_loss = 0.0
     task_losses = {k: 0.0 for k in ["size", "pdi", "ee", "delivery", "biodist", "toxic", "moe_lb"]}
-    task_losses = {k: 0.0 for k in ["size", "pdi", "ee", "delivery", "biodist", "toxic", "moe_lb"]}
     n_batches = 0
     
     # 用于计算准确率
@@ -264,7 +259,6 @@ def validate_balanced(
         outputs = model(smiles, tabular)
         loss, losses = compute_multitask_loss_balanced(
             outputs, targets, mask, task_weights, class_weights, model=model,
-            outputs, targets, mask, task_weights, class_weights, model=model,
         )
         
         total_loss += loss.item()
@@ -315,6 +309,7 @@ def build_optimizer(
     lr: float,
     weight_decay: float,
     backbone_lr_ratio: float = 1.0,
+    size_wd_mult: float = 5.0,
 ) -> torch.optim.AdamW:
     """
     构建 AdamW 优化器，支持分层学习率。
@@ -322,17 +317,30 @@ def build_optimizer(
     仅收集 requires_grad=True 的参数（冻结的 MolT5 encoder 被排除）。
     当 backbone_lr_ratio < 1.0 时，backbone 参数使用 lr * backbone_lr_ratio，
     其余参数（task heads 等）使用 lr。backbone_lr_ratio = 1.0 等价于统一学习率。
+    size_head 信号弱、易过拟合 -> 单独施加 size_wd_mult 倍 weight_decay。
     """
     trainable = [(n, p) for n, p in model.named_parameters() if p.requires_grad]
+    size_wd = weight_decay * size_wd_mult
+
+    def is_size(name: str) -> bool:
+        return "size_head" in name
+
+    size_params = [p for n, p in trainable if is_size(n)]
 
     if backbone_lr_ratio >= 1.0:
+        other = [p for n, p in trainable if not is_size(n)]
         return torch.optim.AdamW(
-            [p for _, p in trainable], lr=lr, weight_decay=weight_decay
+            [
+                {"params": other, "weight_decay": weight_decay},
+                {"params": size_params, "weight_decay": size_wd},
+            ],
+            lr=lr,
         )
 
-    backbone_params = []
+    backbone_params, head_params = [], []
-    head_params = []
     for name, param in trainable:
+        if is_size(name):
+            continue
         if name.startswith(BACKBONE_PREFIXES) and not name.startswith(FROM_SCRATCH_PREFIXES):
             backbone_params.append(param)
         else:
@@ -340,10 +348,10 @@ def build_optimizer(
 
     return torch.optim.AdamW(
         [
-            {"params": backbone_params, "lr": lr * backbone_lr_ratio},
+            {"params": backbone_params, "lr": lr * backbone_lr_ratio, "weight_decay": weight_decay},
-            {"params": head_params, "lr": lr},
+            {"params": head_params, "lr": lr, "weight_decay": weight_decay},
+            {"params": size_params, "lr": lr, "weight_decay": size_wd},
         ],
-        weight_decay=weight_decay,
     )
 
 

reports/ablation_summary.csv

@@ -1,5 +1,5 @@
 variant,run_dir,size.mse_mean,size.mse_std,size.rmse_mean,size.rmse_std,size.mae_mean,size.mae_std,size.r2_mean,size.r2_std,delivery.mse_mean,delivery.mse_std,delivery.rmse_mean,delivery.rmse_std,delivery.mae_mean,delivery.mae_std,delivery.r2_mean,delivery.r2_std,pdi.accuracy_mean,pdi.accuracy_std,pdi.precision_mean,pdi.precision_std,pdi.recall_mean,pdi.recall_std,pdi.f1_mean,pdi.f1_std,ee.accuracy_mean,ee.accuracy_std,ee.precision_mean,ee.precision_std,ee.recall_mean,ee.recall_std,ee.f1_mean,ee.f1_std,toxic.accuracy_mean,toxic.accuracy_std,toxic.precision_mean,toxic.precision_std,toxic.recall_mean,toxic.recall_std,toxic.f1_mean,toxic.f1_std,biodist.kl_divergence_mean,biodist.kl_divergence_std,biodist.js_divergence_mean,biodist.js_divergence_std
-baseline,models/abl/baseline/20260617_180556,0.2768235036159284,0.1346583260723207,0.508798491275714,0.13396864891266763,0.3427187344475763,0.06271329639776796,-0.9110994850991861,1.3609378589681,0.834075853950265,0.13581578656540713,0.9102708371721865,0.0740463162089242,0.634870878942305,0.04810398332825185,0.14309063048291837,0.06678544282427759,0.573809523809524,0.07315376902732006,0.40644301832628565,0.06081224559746048,0.591894166836878,0.15212645465643837,0.40503518503422653,0.07270554354315185,0.65,0.05460640448180816,0.5988843624524457,0.055997857124168736,0.6355952320322067,0.07902933872282923,0.5997699900861396,0.06505269720858664,0.9500396342534485,0.014290090778642328,0.7457142857142857,0.03149343955006944,0.9737706334802525,0.007575117962473199,0.81483431023254,0.0312513648447909,0.42024315728866957,0.04170766696370389,0.10843433388487847,0.010645869413863087
+baseline,models/abl/baseline/20260619_135407,0.9082383737055625,0.5422182144883982,0.9104755717833242,0.28155391471508867,0.5205343839682207,0.03982129636498266,0.05610322557822627,0.19284604391819896,0.7824895463125138,0.14115183472004894,0.8812359118158511,0.07689482452414646,0.6179461847044858,0.04656336248770346,0.196211615771416,0.07653951730854655,0.7063492063492063,0.06354164730554951,0.6606473176789238,0.05894477341890824,0.6895783697291987,0.06686890252898299,0.6613711473169848,0.06575616485061392,0.6444444444444444,0.05513445851732825,0.5885476935334627,0.06114570693094277,0.6115032213351541,0.07118927254312313,0.5872237144865128,0.06396547883545951,0.9534482971536368,0.023380286362927025,0.770595238095238,0.09360247273707213,0.9755667905395673,0.012330101400262774,0.8305779855175,0.07184597861130607,0.38188324071409857,0.04334164661110034,0.09809163862360361,0.012995176709790921
-+moe,models/abl/moe/20260617_191948,0.23034984029037256,0.09049026510650632,0.46822986556281965,0.105406988788203,0.32189987007849846,0.08014105320611228,-0.4728185006647179,0.5427422370821634,0.8054494222584121,0.16018870426337528,0.8929394867547544,0.09004718347937392,0.6145264492845487,0.045144810651025415,0.17842704120419808,0.06422601059093863,0.5896825396825397,0.05682229307831347,0.35469393906374463,0.055592769200553505,0.522735863838455,0.15050332076418682,0.3604048512722805,0.06427235489398077,0.6547619047619047,0.05481364015819684,0.6019542483742392,0.06379027980318479,0.6352859185632295,0.08160957226119422,0.6021181253673275,0.06960400235208303,0.9500779484296937,0.017755203225104002,0.7600104427736006,0.07046117058026663,0.9577096476824245,0.05611117888622105,0.813010272994035,0.03449626596837632,0.3520444477507348,0.08432681015154675,0.08975023050906776,0.021944673828634763
++moe,models/abl/moe/20260619_143616,0.9361227190522554,0.5659367834445426,0.9216387923569987,0.2944565425915824,0.5212914871450111,0.05180236173161496,0.050383827826705994,0.12351338205322832,0.7247201037293581,0.190365922638638,0.8429619825495387,0.11889154597997831,0.5550159654608235,0.06560539845461018,0.2612041487030041,0.15699462642686918,0.7484126984126986,0.07725773665801877,0.6887813376873987,0.08302283869178602,0.6958623275386893,0.0818264031276976,0.6894523197380621,0.08461637088074195,0.653968253968254,0.0651277404410433,0.5934211709278002,0.0740252840701101,0.6238961617280946,0.08844799059005547,0.5959105380152682,0.07995388211928552,0.9501539593267181,0.021557649069279507,0.7709628237259817,0.09333048568985285,0.9577301731339845,0.05648181648201431,0.8191105536863431,0.057895632139112044,0.3162587404501918,0.07095308085906475,0.07871453331202961,0.02165881457825889
-+llm,models/abl/llm/20260617_203747,0.22339749912455512,0.13276860667147747,0.45253463978969516,0.13641810332564366,0.3048570797351743,0.07637999745755018,-0.2662118521181242,0.20796761595726418,0.810584411488471,0.17918852622038825,0.8948853549265088,0.09881605652183918,0.6158850481512493,0.04133793584599054,0.1761904116252725,0.0670662068919268,0.6,0.07271331378442132,0.382198316306454,0.059482007465697324,0.552886246650569,0.15151126540989962,0.3854764962481809,0.06295617241236845,0.6492063492063492,0.05989829290194668,0.5992000446776198,0.0738385570229797,0.6206719076130841,0.08494291371338275,0.5966089895417119,0.07618779152357238,0.9511890595408049,0.016071213375949578,0.7617961570593149,0.0689235340767993,0.958294443004062,0.05613748087593979,0.8151453855878632,0.03136647604932653,0.4240504874533223,0.049093469886439346,0.1077911642566912,0.015306583731384354
++llm,models/abl/llm/20260619_152140,0.9446909140698433,0.5246450869505823,0.9323618485203619,0.2745765785596742,0.534852125810107,0.06923968128710634,0.010675298699051424,0.17503260221449574,0.7944675704032856,0.18114607138886382,0.8853841868434346,0.10277359627295161,0.6019238276123685,0.035144473220882017,0.19083465917332137,0.10406116788019794,0.7166666666666667,0.06793020294136717,0.6712542207128174,0.06849859623450226,0.6927146000295703,0.06697261945160246,0.6686471365355745,0.06875001186427347,0.6531746031746032,0.04991804823338536,0.6004632345280062,0.05337084241208218,0.6391669855199268,0.06873793721898303,0.6051066395139207,0.05943186282345569,0.950116262605939,0.02064666686696538,0.7743065998329156,0.09235234794557083,0.9416486618845965,0.07565471465983832,0.8111862357555301,0.03737221833347304,0.387547460852652,0.0864564610528379,0.09686251581762508,0.024012654397470485
-+both,models/abl/both/20260617_232841,0.21249719869892153,0.083819203886272,0.4509291694638765,0.09570832161071456,0.29299021591121593,0.04711634555603193,-0.33611237791230353,0.3452040990881361,0.7425693610280117,0.15313439940788104,0.8571056675991812,0.08910238828097625,0.5779881956006352,0.040317114677088234,0.24162283519579483,0.08163767158765754,0.634126984126984,0.08009190991455786,0.36320050940978454,0.06018646749200912,0.5239036063335376,0.14740021344728976,0.37610884530892047,0.0760087699533074,0.6563492063492063,0.06431020501550125,0.6032694055441358,0.07290814030935075,0.6464310864983134,0.09580096959112427,0.6100615120789444,0.08078656411296503,0.9511890595408049,0.016071213375949578,0.7617961570593149,0.0689235340767993,0.958294443004062,0.05613748087593979,0.8151453855878632,0.03136647604932653,0.3412071963471256,0.07500564394419486,0.08684355805526371,0.020779098028044035
++both,models/abl/both/20260619_174020,0.9595393632675637,0.5340398065327858,0.9401911015958848,0.2749182710324654,0.5328599284215613,0.051988251756972186,-0.000542385625156198,0.11211912465762124,0.7993840715989491,0.1569950944154102,0.889711124724226,0.08830733911120735,0.5756929427656439,0.06030148450805026,0.17102632095715226,0.16310841506794083,0.7031746031746032,0.060150143877766404,0.6493500422415014,0.06021511831854895,0.672008899565058,0.06669948879965092,0.65127019538571,0.06496385143120126,0.6412698412698412,0.06474945414086074,0.5870717598152753,0.07093354974743615,0.6252169231328896,0.09493156871874915,0.591592552147522,0.07893156617942566,0.9522806888673845,0.015363196314866132,0.753095238095238,0.04231535854936122,0.974950666897128,0.008126538565775098,0.8212449416431714,0.038549704208112205,0.2516389180942744,0.06568800916844175,0.06192911639345754,0.01766285899355436