Fix BRCA training collapse with lab-grid defaults

Training on BRCA with the previous defaults (Adam, lr 1e-3, weight decay
0.01) collapsed: every sample got the same prediction (macro F1 0.133).

- Default to the configuration from the Bozdag Lab grid that was stable
across 8 seeds on BRCA: Adam, lr 0.01, L2 0.01, batch 32, projection
64, tanh/relu/tanh activations, 200 epochs (macro F1 0.749 +/- 0.018).
- Make the three activation functions and SGD momentum (default 0.03)
configurable in Hyperparameters. model.py and the artifact format are
unchanged; artifacts already store their activations.
- Make the automatic GenePT choice the lab-recommended single-cell file.
- Default to a 75/25 stratified train/validation split, with helper text
and a clear error when a class is too small to stratify.
- Fix the Results page training-history plot, which was never populated.
- Add stable element ids for e2e tests and a layout tweak for short
desktop screens.
- Tests: split, leakage, class weights, history plot, activations,
momentum and defaults (22 passing).

app.py

@@ -34,6 +34,7 @@ from biolmnet.data import (
     validate_and_align_omics,
 )
 from biolmnet.training import (
+    DEFAULT_VALIDATION_FRACTION,
     Hyperparameters,
     ModelBundle,
     pathway_importance,
@@ -45,7 +46,7 @@ from biolmnet.ui.styles import CSS
 
 
 GENEPT_OPTIONS = {
-    "Auto — bulk or single-cell based on source": "auto",
+    "Auto — lab-recommended single-cell context": "auto",
     "Bulk · large-3 context": "embedding_original_large_3.parquet",
     "Bulk · original ada-text context": "embedding_original_ada_text.parquet",
     "Single-cell · cell type, tissue, drug & pathway": (
@@ -101,9 +102,8 @@ def _resolve_embedding_file(option: str, source_mode: str, example_dataset: str)
     selected = GENEPT_OPTIONS[option]
     if selected != "auto":
         return selected
-    if source_mode == "BioLM-NET examples" and example_dataset == "scTrioseq2":
-        return "embedding_associations_cell_type_tissue_drug_pathway_openai_large.parquet"
-    return "embedding_original_large_3.parquet"
+    # Bozdag Lab recommendation (J. Rifat, 2026-10-17) for every source.
+    return "embedding_associations_cell_type_tissue_drug_pathway_openai_large.parquet"
 
 
 def _graph_preview_html(mask_density: str, gene_nodes: str, pathway_units: str, genept_dim: str) -> str:
@@ -514,6 +514,14 @@ def _loss_plot(history: list[dict[str, float]]) -> go.Figure:
     return figure
 
 
+def _validation_split_note(validation_fraction: float) -> str:
+    validation_percent = round(float(validation_fraction) * 100)
+    return (
+        f'<div class="field-hint">{validation_percent}% of samples are held out for validation '
+        f"(stratified by class). The remaining {100 - validation_percent}% are used for training.</div>"
+    )
+
+
 def estimate_training_duration(
     workspace: PreparedWorkspace | None,
     epochs: int,
@@ -949,6 +957,7 @@ def refresh_results(bundle: ModelBundle | None, run_meta: dict):
                 "Build the biological architecture, train a model, or run prediction to populate this page."
             ),
             gr.update(visible=False),
+            None,
         )
 
     macro_auc = run_meta.get("macro_auc")
@@ -965,7 +974,8 @@ def refresh_results(bundle: ModelBundle | None, run_meta: dict):
     )
         + "</div>"
     )
-    return stats, gr.update(visible=True), "", gr.update(visible=True)
+    loss_figure = _loss_plot(bundle.history) if bundle.history else None
+    return stats, gr.update(visible=True), "", gr.update(visible=True), loss_figure
 
 
 THEME = gr.themes.Base(primary_hue="gray", neutral_hue="gray")
@@ -986,18 +996,19 @@ with gr.Blocks(title="BioLM-NET Workbench", fill_width=True) as demo:
     session_id_state = gr.State("")
     intro_step_state = gr.State(0)
     intro_dismissed_state = gr.BrowserState(False, storage_key="biolmnet_intro_dismissed")
+    nav_ids = ["nav-data", "nav-train", "nav-export", "nav-predict", "nav-results"]
 
     with gr.Row(elem_classes=["app-shell"]):
         with gr.Column(elem_classes=["rail-col"]):
             with gr.Column(elem_classes=["rail"]):
-                gr.HTML(ui.brand_block_html())
+                gr.HTML(ui.brand_block_html(), elem_id="app-header")
                 rail_row_htmls = []
                 rail_row_buttons = []
                 with gr.Column(elem_classes=["rail-nav"]):
                     for index, label in enumerate(STAGE_LABELS):
                         with gr.Column(elem_classes=["rail-row-wrap"]):
                             row_html = gr.HTML(ui.rail_row_html(index + 1, label, "on" if index == 0 else "off", "Active" if index == 0 else "Locked"))
-                            row_button = gr.Button("", elem_classes=["rail-row-click"])
+                            row_button = gr.Button("", elem_id=nav_ids[index], elem_classes=["rail-row-click"])
                         rail_row_htmls.append(row_html)
                         rail_row_buttons.append(row_button)
                 with gr.Column(elem_classes=["rail-spacer", "rail-plate-wrap"]):
@@ -1011,7 +1022,7 @@ with gr.Blocks(title="BioLM-NET Workbench", fill_width=True) as demo:
                 intro_open_button = gr.Button("Introduction", size="sm", variant="secondary", elem_classes=["btn-ghost", "intro-open-btn"])
 
             # ── Data & Priors ────────────────────────────────────────
-            with gr.Column(visible=True) as data_page:
+            with gr.Column(visible=True, elem_id="page-data", elem_classes=["workspace-page", "data-workspace-page"]) as data_page:
                 with gr.Row(elem_classes=["pghd"]):
                     gr.HTML(
                         ui.title_block_html(
@@ -1028,7 +1039,7 @@ with gr.Blocks(title="BioLM-NET Workbench", fill_width=True) as demo:
                         source_mode = gr.Radio(
                             ["BioLM-NET examples", "GitHub folder", "Upload files"],
                             value="BioLM-NET examples", show_label=False, container=False,
-                            elem_classes=["seg-radio"],
+                            elem_id="dataset-source", elem_classes=["seg-radio"],
                         )
                         with gr.Row(elem_classes=["srow-row", "first"]):
                             with gr.Column(scale=1, min_width=0):
@@ -1038,7 +1049,7 @@ with gr.Blocks(title="BioLM-NET Workbench", fill_width=True) as demo:
                                             gr.HTML('<div class="field-label">Repository dataset</div>')
                                             example_dataset = gr.Dropdown(
                                                 ["BRCA", "COAD", "GBM", "scTrioseq2"], value="BRCA",
-                                                show_label=False, container=False,
+                                                show_label=False, container=False, elem_id="dataset-select",
                                             )
                                 with gr.Column(visible=False) as github_group:
                                     github_folder = gr.Textbox(
@@ -1144,12 +1155,13 @@ with gr.Blocks(title="BioLM-NET Workbench", fill_width=True) as demo:
                             "Re-validate inputs", size="sm", variant="secondary", elem_classes=["actbar-secondary-btn"]
                         )
                         prepare_button = gr.Button(
-                            "Build biological architecture", size="sm", variant="primary", elem_classes=["actbar-primary-btn"]
+                            "Build biological architecture", size="sm", variant="primary",
+                            elem_id="build-architecture-btn", elem_classes=["actbar-primary-btn"]
                         )
                 preparation_status = gr.HTML(visible=False)
 
             # ── Train Model ──────────────────────────────────────────
-            with gr.Column(visible=False) as train_page:
+            with gr.Column(visible=False, elem_id="page-train", elem_classes=["workspace-page"]) as train_page:
                 with gr.Row(elem_classes=["pghd"]):
                     gr.HTML(
                         ui.title_block_html(
@@ -1165,22 +1177,22 @@ with gr.Blocks(title="BioLM-NET Workbench", fill_width=True) as demo:
                             with gr.Column(min_width=140):
                                 gr.HTML(ui.esc("Epochs"))
                             with gr.Column():
-                                epochs = gr.Slider(5, 200, value=50, step=5, show_label=False, container=False)
+                                epochs = gr.Slider(5, 200, value=Hyperparameters.epochs, step=5, show_label=False, container=False)
                         with gr.Row(elem_classes=["srow-row"]):
                             with gr.Column(min_width=140):
                                 gr.HTML(ui.esc("Batch size"))
                             with gr.Column():
-                                batch_size = gr.Dropdown([8, 16, 32, 64, 128], value=16, show_label=False, container=False)
+                                batch_size = gr.Dropdown([8, 16, 32, 64, 128], value=Hyperparameters.batch_size, show_label=False, container=False)
                         with gr.Row(elem_classes=["srow-row"]):
                             with gr.Column(min_width=140):
                                 gr.HTML(ui.esc("Learning rate"))
                             with gr.Column():
-                                learning_rate = gr.Number(value=0.001, minimum=0.000001, show_label=False, container=False)
+                                learning_rate = gr.Number(value=Hyperparameters.learning_rate, minimum=0.000001, show_label=False, container=False)
                         with gr.Row(elem_classes=["srow-row"]):
                             with gr.Column(min_width=140):
                                 gr.HTML(ui.esc("L2 weight decay"))
                             with gr.Column():
-                                weight_decay = gr.Number(value=0.01, minimum=0, show_label=False, container=False)
+                                weight_decay = gr.Number(value=Hyperparameters.weight_decay, minimum=0, show_label=False, container=False)
                         with gr.Row(elem_classes=["srow-row", "last"]):
                             with gr.Column(min_width=140):
                                 gr.HTML(ui.esc("Dropout"))
@@ -1202,7 +1214,13 @@ with gr.Blocks(title="BioLM-NET Workbench", fill_width=True) as demo:
                             with gr.Column(min_width=140):
                                 gr.HTML(ui.esc("Validation fraction"))
                             with gr.Column():
-                                validation_fraction = gr.Slider(0.1, 0.4, value=0.2, step=0.05, show_label=False, container=False)
+                                validation_fraction = gr.Slider(
+                                    0.1, 0.4, value=DEFAULT_VALIDATION_FRACTION, step=0.05, show_label=False, container=False
+                                )
+                                validation_split_note = gr.HTML(_validation_split_note(DEFAULT_VALIDATION_FRACTION))
+                                validation_fraction.change(
+                                    _validation_split_note, inputs=[validation_fraction], outputs=[validation_split_note]
+                                )
                         with gr.Row(elem_classes=["srow-row"]):
                             with gr.Column(min_width=140):
                                 gr.HTML(ui.esc("Optimizer"))
@@ -1221,11 +1239,16 @@ with gr.Blocks(title="BioLM-NET Workbench", fill_width=True) as demo:
                         gr.HTML('<div class="mono" style="margin-bottom:9px">Training result</div>')
                         training_status = gr.HTML(_status("Prepare the architecture before training."))
                         loss_plot = gr.Plot(label="Loss by epoch", show_label=False)
-                        gr.HTML('<div class="mono" style="margin:20px 0 9px">Paper-faithful defaults</div>')
+                        gr.HTML('<div class="mono" style="margin:20px 0 9px">Architecture defaults</div>')
                         gr.HTML(
                             "".join(
                                 [
                                     ui.kv_plain_row("First layer", "trainable W ⊙ M"),
+                                    ui.kv_plain_row(
+                                        "Activations",
+                                        f"{Hyperparameters.biological_activation} / {Hyperparameters.projection_activation} / "
+                                        f"{Hyperparameters.fusion_activation}",
+                                    ),
                                     ui.kv_plain_row("PDI weights", "binary"),
                                     ui.kv_plain_row("PPI weights", "STRING norm."),
                                     ui.kv_plain_row("Attention", "GenePT query"),
@@ -1238,11 +1261,12 @@ with gr.Blocks(title="BioLM-NET Workbench", fill_width=True) as demo:
                     train_actionbar_note = gr.HTML('<div class="actbar-note">Set hyperparameters, then train.</div>')
                     with gr.Row(elem_classes=["actbar-buttons"]):
                         train_button = gr.Button(
-                            "Train BioLM-NET on ZeroGPU", size="sm", variant="primary", elem_classes=["actbar-primary-btn"]
+                            "Train BioLM-NET on ZeroGPU", size="sm", variant="primary",
+                            elem_id="train-btn", elem_classes=["actbar-primary-btn"]
                         )
 
             # ── Export Artifacts ─────────────────────────────────────
-            with gr.Column(visible=False) as export_page:
+            with gr.Column(visible=False, elem_id="page-export", elem_classes=["workspace-page"]) as export_page:
                 with gr.Row(elem_classes=["pghd"]):
                     gr.HTML(
                         ui.title_block_html(
@@ -1270,7 +1294,7 @@ with gr.Blocks(title="BioLM-NET Workbench", fill_width=True) as demo:
                     gr.HTML('<div class="actbar-note">Bundle written to a temp file on export · not persisted between Space restarts</div>')
 
             # ── Predict ──────────────────────────────────────────────
-            with gr.Column(visible=False) as predict_page:
+            with gr.Column(visible=False, elem_id="page-predict", elem_classes=["workspace-page"]) as predict_page:
                 with gr.Row(elem_classes=["pghd"]):
                     gr.HTML(
                         ui.title_block_html(
@@ -1334,12 +1358,12 @@ with gr.Blocks(title="BioLM-NET Workbench", fill_width=True) as demo:
                     with gr.Row(elem_classes=["actbar-buttons"]):
                         predict_button = gr.Button(
                             "Run inference", size="sm", variant="primary", interactive=False,
-                            elem_classes=["actbar-primary-btn"],
+                            elem_id="predict-btn", elem_classes=["actbar-primary-btn"],
                         )
                 prediction_status = gr.HTML(visible=False)
 
             # ── Results ──────────────────────────────────────────────
-            with gr.Column(visible=False) as results_page:
+            with gr.Column(visible=False, elem_id="page-results", elem_classes=["workspace-page"]) as results_page:
                 with gr.Row(elem_classes=["pghd"]):
                     gr.HTML(
                         ui.title_block_html(
@@ -1642,7 +1666,7 @@ with gr.Blocks(title="BioLM-NET Workbench", fill_width=True) as demo:
     ).then(
         refresh_results,
         inputs=[model_state, run_meta_state],
-        outputs=[results_stats, results_stats_row, results_empty_note, results_content],
+        outputs=[results_stats, results_stats_row, results_empty_note, results_content, results_loss_plot],
         show_progress="hidden",
     ).then(
         lambda df: gr.update(value=df), inputs=[validation_predictions_state], outputs=[results_validation_table],
@@ -1718,7 +1742,7 @@ with gr.Blocks(title="BioLM-NET Workbench", fill_width=True) as demo:
         show_progress="hidden",
     ).then(
         refresh_results, inputs=[model_state, run_meta_state],
-        outputs=[results_stats, results_stats_row, results_empty_note, results_content],
+        outputs=[results_stats, results_stats_row, results_empty_note, results_content, results_loss_plot],
         show_progress="hidden",
     )
 

biolmnet/training.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import copy
+import math
 import random
 from dataclasses import asdict, dataclass
 from typing import Callable
@@ -25,21 +26,36 @@ from .model import BioLMNet
 
 ProgressCallback = Callable[[float, str], None]
 
+# 75% training / 25% validation. The validation partition drives early
+# stopping and best-checkpoint selection, so it is not an untouched test set.
+DEFAULT_VALIDATION_FRACTION = 0.25
 
+# Activations BioLMNet accepts.
+ACTIVATIONS = ("relu", "sigmoid", "tanh")
+
+
+# Defaults follow the Bozdag Lab grid (J. Rifat, 2026-10-17): the configuration that stayed
+# stable across 8 seeds on BRCA (macro F1 0.749 +/- 0.018). The published activations
+# (relu/sigmoid/tanh) collapse at this learning rate and L2; see
+# diagnostics/2026-10-30-lab-hyperparameter-grid.
 @dataclass(frozen=True)
 class Hyperparameters:
-    epochs: int = 50
-    batch_size: int = 16
-    learning_rate: float = 0.001
+    epochs: int = 200
+    batch_size: int = 32
+    learning_rate: float = 0.01
     weight_decay: float = 0.01
     dropout: float = 0.3
     projection_dim: int = 64
     fusion_dim: int = 12
-    validation_fraction: float = 0.2
+    validation_fraction: float = DEFAULT_VALIDATION_FRACTION
     optimizer: str = "Adam"
     class_weighting: bool = True
     early_stopping_patience: int = 12
     seed: int = 42
+    biological_activation: str = "tanh"
+    projection_activation: str = "relu"
+    fusion_activation: str = "tanh"
+    sgd_momentum: float = 0.03
 
     def validate(self) -> None:
         if not 1 <= self.epochs <= 1000:
@@ -54,6 +70,11 @@ class Hyperparameters:
             raise ValueError("Validation fraction must be between 0.05 and 0.5.")
         if self.optimizer.lower() not in {"adam", "sgd"}:
             raise ValueError("Optimizer must be Adam or SGD.")
+        for name in ("biological_activation", "projection_activation", "fusion_activation"):
+            if getattr(self, name) not in ACTIVATIONS:
+                raise ValueError(f"{name.replace('_', ' ').capitalize()} must be one of {', '.join(ACTIVATIONS)}.")
+        if not 0 <= self.sgd_momentum < 1:
+            raise ValueError("SGD momentum must be in [0, 1).")
 
 
 @dataclass
@@ -86,6 +107,36 @@ def set_reproducible_seed(seed: int) -> None:
         torch.cuda.manual_seed_all(seed)
 
 
+def split_train_validation(
+    labels: np.ndarray, validation_fraction: float, seed: int
+) -> tuple[np.ndarray, np.ndarray]:
+    """Stratified sample-index split shared by both omics and the labels."""
+    classes, class_counts = np.unique(labels, return_counts=True)
+    if class_counts.min() < 2:
+        raise ValueError(
+            "Each class needs at least two samples for a stratified train/validation split."
+        )
+    sample_count = len(labels)
+    validation_count = math.ceil(validation_fraction * sample_count)
+    train_count = sample_count - validation_count
+    if min(validation_count, train_count) < len(classes):
+        raise ValueError(
+            f"A validation fraction of {validation_fraction:.2f} on {sample_count} samples "
+            f"gives {train_count} training and {validation_count} validation samples; "
+            f"each partition needs at least one sample per class ({len(classes)} classes) "
+            "for a stratified split. Adjust the validation fraction or add samples."
+        )
+
+    indices = np.arange(sample_count)
+    train_index, validation_index = train_test_split(
+        indices,
+        test_size=validation_fraction,
+        random_state=seed,
+        stratify=labels,
+    )
+    return train_index, validation_index
+
+
 def _fit_scaler(values: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     mean = values.mean(axis=0, dtype=np.float64).astype(np.float32)
     scale = values.std(axis=0, dtype=np.float64).astype(np.float32)
@@ -121,6 +172,9 @@ def _make_model(
         projection_dim=hyperparameters.projection_dim,
         fusion_dim=hyperparameters.fusion_dim,
         dropout=hyperparameters.dropout,
+        biological_activation=hyperparameters.biological_activation,
+        projection_activation=hyperparameters.projection_activation,
+        fusion_activation=hyperparameters.fusion_activation,
     )
 
 
@@ -153,18 +207,8 @@ def train(
     hyperparameters.validate()
     set_reproducible_seed(hyperparameters.seed)
     labels = workspace.labels
-    classes, class_counts = np.unique(labels, return_counts=True)
-    if class_counts.min() < 2:
-        raise ValueError(
-            "Each class needs at least two samples for a stratified train/validation split."
-        )
-
-    indices = np.arange(len(labels))
-    train_index, validation_index = train_test_split(
-        indices,
-        test_size=hyperparameters.validation_fraction,
-        random_state=hyperparameters.seed,
-        stratify=labels,
+    train_index, validation_index = split_train_validation(
+        labels, hyperparameters.validation_fraction, hyperparameters.seed
     )
     gene_mean, gene_scale = _fit_scaler(
         workspace.gene_expression[train_index]
@@ -204,6 +248,7 @@ def train(
         optimizer = torch.optim.SGD(
             model.parameters(),
             lr=hyperparameters.learning_rate,
+            momentum=hyperparameters.sgd_momentum,
             weight_decay=hyperparameters.weight_decay,
         )
     else:
@@ -357,9 +402,9 @@ def train(
             "projection_dim": hyperparameters.projection_dim,
             "fusion_dim": hyperparameters.fusion_dim,
             "dropout": hyperparameters.dropout,
-            "biological_activation": "relu",
-            "projection_activation": "sigmoid",
-            "fusion_activation": "tanh",
+            "biological_activation": hyperparameters.biological_activation,
+            "projection_activation": hyperparameters.projection_activation,
+            "fusion_activation": hyperparameters.fusion_activation,
         },
     }
     bundle = ModelBundle(

biolmnet/ui/styles.py

@@ -236,6 +236,12 @@ body, .gradio-container {
   flex: 1 1 auto !important;
   width: calc(100% - 236px) !important;
 }
+.workspace-page,
+.workspace-page > .form {
+  min-height: calc(100vh - 45px) !important;
+  display: flex !important;
+  flex-direction: column !important;
+}
 
 /* rail row (overlay: visible .rr under an invisible full-size button) */
 .rail-row-wrap { position: relative; }
@@ -424,6 +430,46 @@ body, .gradio-container {
 .strip.error .strip-kicker { color: var(--color-error); }
 .strip .strip-body { margin-top: 4px; font-size: 12.5px; line-height: 1.5; }
 .empty-note { padding: 14px 0; color: var(--color-muted); font-size: 13px; }
+.field-hint { margin-top: 6px; color: var(--color-muted); font-size: 12px; line-height: 1.4; }
+
+@media (min-width: 901px) and (max-height: 820px) {
+  .topbar {
+    padding-top: 10px;
+    padding-bottom: 10px;
+  }
+  .pghd {
+    padding-top: 18px;
+  }
+  .pghd h2 {
+    font-size: 31px;
+  }
+  .pghd .desc {
+    margin-top: 4px;
+    line-height: 1.35;
+  }
+  .stage-grid {
+    padding-top: 16px !important;
+    gap: 22px !important;
+  }
+  .data-workspace-page .section-heading-row {
+    margin-bottom: 8px !important;
+  }
+  .data-workspace-page .sect-spaced {
+    margin-top: 18px !important;
+  }
+  .data-workspace-page .srow-row {
+    min-height: 42px;
+  }
+  .data-workspace-page .dataframe,
+  .data-workspace-page .table-wrap {
+    max-height: 118px !important;
+    overflow: auto !important;
+  }
+  .actbar-row {
+    padding-top: 11px !important;
+    padding-bottom: 11px !important;
+  }
+}
 
 /* ── slots (one-line gr.File) ─────────────────────────────────────────── */
 .file-slot .file-preview { border: 1px solid var(--color-divider) !important; border-radius: var(--radius-md); background: #fff; font-size: 12.5px !important; }

tests/test_core.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 import pickle
+from dataclasses import replace
 
 import numpy as np
 import pandas as pd
+import pytest
 import torch
 
 from biolmnet.artifacts import load_bundle, save_bundle
@@ -15,7 +17,13 @@ from biolmnet.data import (
     deterministic_gene_embeddings,
 )
 from biolmnet.model import BioMaskedLinear, BioLMNet
-from biolmnet.training import Hyperparameters, predict, train
+from biolmnet.training import (
+    DEFAULT_VALIDATION_FRACTION,
+    Hyperparameters,
+    predict,
+    split_train_validation,
+    train,
+)
 
 
 def test_biological_mask_uses_pdi_and_undirected_ppi() -> None:
@@ -232,3 +240,217 @@ def test_zerogpu_duration_estimator_is_bounded_and_scales() -> None:
     long = estimate_training_duration(common[0], 200, *common[1:])
     assert 30 <= short <= 300
     assert short <= long <= 300
+
+
+_FAST = Hyperparameters(
+    epochs=3,
+    batch_size=8,
+    projection_dim=4,
+    fusion_dim=3,
+    dropout=0.0,
+    early_stopping_patience=3,
+)
+
+
+def test_default_split_is_75_training_25_validation() -> None:
+    assert DEFAULT_VALIDATION_FRACTION == 0.25
+    assert Hyperparameters().validation_fraction == 0.25
+    labels = np.repeat(np.array([0, 1]), 50)
+    train_index, validation_index = split_train_validation(
+        labels, Hyperparameters().validation_fraction, Hyperparameters().seed
+    )
+    assert len(train_index) == 75
+    assert len(validation_index) == 25
+
+
+def test_split_partitions_every_sample_exactly_once() -> None:
+    labels = np.repeat(np.array([0, 1, 2]), [40, 25, 18])
+    train_index, validation_index = split_train_validation(labels, 0.25, 42)
+    assert len(train_index) + len(validation_index) == len(labels)
+    assert set(train_index).isdisjoint(validation_index)
+    assert sorted(np.concatenate([train_index, validation_index])) == list(range(len(labels)))
+
+
+def test_split_preserves_class_proportions() -> None:
+    labels = np.repeat(np.array([0, 1, 2]), [120, 60, 20])  # 60 / 30 / 10
+    train_index, validation_index = split_train_validation(labels, 0.25, 42)
+    expected = np.array([0.6, 0.3, 0.1])
+    for partition in (train_index, validation_index):
+        observed = np.bincount(labels[partition], minlength=3) / len(partition)
+        np.testing.assert_allclose(observed, expected, atol=0.02)
+
+
+def test_split_is_reproducible_for_a_fixed_seed() -> None:
+    labels = np.repeat(np.array([0, 1]), [30, 34])
+    first = split_train_validation(labels, 0.25, 42)
+    second = split_train_validation(labels, 0.25, 42)
+    other_seed = split_train_validation(labels, 0.25, 7)
+    np.testing.assert_array_equal(first[0], second[0])
+    np.testing.assert_array_equal(first[1], second[1])
+    assert set(first[1]) != set(other_seed[1])
+
+
+def test_split_rejects_partitions_too_small_to_stratify() -> None:
+    with pytest.raises(ValueError, match="at least two samples"):
+        split_train_validation(np.array([0, 0, 0, 1]), 0.25, 42)
+    # 4 classes x 2 samples: 25% leaves 2 validation samples for 4 classes.
+    with pytest.raises(ValueError, match="one sample per class"):
+        split_train_validation(np.repeat(np.arange(4), 2), 0.25, 42)
+
+
+def test_training_uses_paired_rows_and_train_only_scaler() -> None:
+    workspace = _tiny_workspace()
+    hyperparameters = replace(_FAST, validation_fraction=0.25)
+    train_index, validation_index = split_train_validation(
+        workspace.labels, hyperparameters.validation_fraction, hyperparameters.seed
+    )
+    # Shift only the validation rows: a scaler fitted on all samples would move.
+    workspace.gene_expression[validation_index] += 100.0
+    workspace.dna_methylation[validation_index] -= 100.0
+
+    result = train(workspace, hyperparameters)
+    bundle = result.bundle
+
+    assert bundle.metrics["training_samples"] == 36
+    assert bundle.metrics["validation_samples"] == 12
+    np.testing.assert_array_equal(
+        result.validation_predictions["sample_row"].to_numpy(), validation_index
+    )
+    expected_labels = [workspace.label_names[value] for value in workspace.labels[validation_index]]
+    assert result.validation_predictions["observed"].tolist() == expected_labels
+
+    # Both omics scalers come from the same training rows only.
+    np.testing.assert_allclose(
+        bundle.gene_mean, workspace.gene_expression[train_index].mean(axis=0), rtol=1e-5
+    )
+    np.testing.assert_allclose(
+        bundle.gene_scale, workspace.gene_expression[train_index].std(axis=0), rtol=1e-5
+    )
+    np.testing.assert_allclose(
+        bundle.dna_mean, workspace.dna_methylation[train_index].mean(axis=0), rtol=1e-5
+    )
+    np.testing.assert_allclose(
+        bundle.dna_scale, workspace.dna_methylation[train_index].std(axis=0), rtol=1e-5
+    )
+
+    # Validation rows are transformed with that training-fitted scaler: the
+    # stored validation probabilities match predict(), which uses the bundle's scaler.
+    rescored = predict(
+        pd.DataFrame(workspace.gene_expression[validation_index], columns=bundle.gene_features),
+        pd.DataFrame(workspace.dna_methylation[validation_index], columns=bundle.dna_features),
+        bundle,
+    )
+    probability_columns = [f"P({label})" for label in workspace.label_names]
+    np.testing.assert_allclose(
+        result.validation_predictions[probability_columns].to_numpy(),
+        rescored[probability_columns].to_numpy(),
+        atol=1e-5,
+    )
+
+
+def test_class_weights_use_training_labels_only(monkeypatch) -> None:
+    import biolmnet.training as training
+
+    captured = {}
+    original = training.nn.CrossEntropyLoss
+
+    def spy(*args, weight=None, **kwargs):
+        captured["weight"] = None if weight is None else weight.cpu().numpy()
+        return original(*args, weight=weight, **kwargs)
+
+    monkeypatch.setattr(training.nn, "CrossEntropyLoss", spy)
+    workspace = _tiny_workspace()  # 24 control / 24 case: all-sample weights are [1, 1].
+    # Force a lopsided split so training-only weights clearly differ from all-sample ones.
+    train_index = np.concatenate([np.arange(0, 22), np.arange(24, 30)])
+    validation_index = np.setdiff1d(np.arange(48), train_index)
+    monkeypatch.setattr(
+        training, "split_train_validation", lambda *args: (train_index, validation_index)
+    )
+    train(workspace, _FAST)
+
+    counts = np.bincount(workspace.labels[train_index], minlength=2).astype(np.float32)
+    assert counts.tolist() == [22.0, 6.0]
+    np.testing.assert_allclose(captured["weight"], len(train_index) / (2 * counts), rtol=1e-6)
+
+
+def test_results_page_receives_training_history_plot() -> None:
+    from app import refresh_results
+
+    assert refresh_results(None, {})[-1] is None
+
+    result = train(_tiny_workspace(), _FAST)
+    figure = refresh_results(result.bundle, {})[-1]
+    traces = {trace.name: trace for trace in figure.data}
+    assert set(traces) == {"Training", "Validation"}
+    history = result.bundle.history
+    assert list(traces["Training"].y) == [row["training_loss"] for row in history]
+    assert list(traces["Validation"].y) == [row["validation_loss"] for row in history]
+
+    # Artifacts saved without history still render the rest of the page.
+    result.bundle.history = []
+    assert refresh_results(result.bundle, {})[-1] is None
+
+
+def test_activation_choices_reach_model_config_and_artifact(tmp_path) -> None:
+    hyperparameters = replace(
+        _FAST,
+        biological_activation="tanh",
+        projection_activation="relu",
+        fusion_activation="sigmoid",
+    )
+    result = train(_tiny_workspace(), hyperparameters)
+    branch = result.bundle.model.gene_branch
+    assert isinstance(branch.biological_activation, torch.nn.Tanh)
+    assert isinstance(branch.projection_activation, torch.nn.ReLU)
+    assert isinstance(result.bundle.model.fusion_activation, torch.nn.Sigmoid)
+    assert result.bundle.config["architecture"]["biological_activation"] == "tanh"
+
+    restored = load_bundle(save_bundle(result.bundle, tmp_path / "model.zip"))
+    assert isinstance(restored.model.gene_branch.biological_activation, torch.nn.Tanh)
+    assert isinstance(restored.model.fusion_activation, torch.nn.Sigmoid)
+
+
+def test_defaults_follow_lab_grid_configuration() -> None:
+    defaults = Hyperparameters()
+    assert (defaults.optimizer, defaults.learning_rate, defaults.weight_decay) == ("Adam", 0.01, 0.01)
+    assert (defaults.batch_size, defaults.projection_dim, defaults.dropout, defaults.epochs) == (32, 64, 0.3, 200)
+    assert (defaults.biological_activation, defaults.projection_activation, defaults.fusion_activation) == (
+        "tanh",
+        "relu",
+        "tanh",
+    )
+    assert defaults.sgd_momentum == 0.03
+
+
+def test_auto_genept_uses_lab_recommended_single_cell_file() -> None:
+    from app import GENEPT_OPTIONS, _resolve_embedding_file
+
+    auto = next(label for label, value in GENEPT_OPTIONS.items() if value == "auto")
+    for dataset in ("BRCA", "COAD", "GBM", "scTrioseq2"):
+        assert _resolve_embedding_file(auto, "BioLM-NET examples", dataset) == (
+            "embedding_associations_cell_type_tissue_drug_pathway_openai_large.parquet"
+        )
+
+
+def test_sgd_momentum_reaches_optimizer(monkeypatch) -> None:
+    import biolmnet.training as training
+
+    captured = {}
+    original = training.torch.optim.SGD
+
+    def spy(params, **kwargs):
+        captured.update(kwargs)
+        return original(params, **kwargs)
+
+    monkeypatch.setattr(training.torch.optim, "SGD", spy)
+    train(_tiny_workspace(), replace(_FAST, optimizer="SGD", learning_rate=0.01, sgd_momentum=0.03))
+    assert captured["momentum"] == 0.03
+
+
+@pytest.mark.parametrize(
+    "field, value",
+    [("biological_activation", "softmax"), ("fusion_activation", "gelu"), ("sgd_momentum", 1.0)],
+)
+def test_invalid_activation_or_momentum_is_rejected(field, value) -> None:
+    with pytest.raises(ValueError):
+        replace(Hyperparameters(), **{field: value}).validate()