From ffd8890bc63a9c725dc20cdd03415e4e2ea23a59 Mon Sep 17 00:00:00 2001 From: isvock Date: Thu, 30 Jul 2026 15:06:05 -0700 Subject: [PATCH 1/2] Add option to use LayerNorm in head --- docs/training_configuration.md | 22 +++++++- scripts/example_train_config.json | 3 +- src/transcriptml/models/reproduce.py | 16 ++++-- src/transcriptml/training/trainer.py | 18 ++++++- src/transcriptml/workflows/init_run.py | 1 + tests/test_cli_analysis.py | 2 + tests/test_splits_models.py | 35 +++++++++++++ tests/test_training_losses.py | 72 ++++++++++++++++++++++++++ 8 files changed, 163 insertions(+), 6 deletions(-) diff --git a/docs/training_configuration.md b/docs/training_configuration.md index bce2d93..748f400 100644 --- a/docs/training_configuration.md +++ b/docs/training_configuration.md @@ -49,6 +49,7 @@ defaults: "num_workers": 0, "mmap_mode": "r", "seed": 42, + "head_layernorm": false, "split_source": "auto", "split": { "method": "random", @@ -130,6 +131,7 @@ above. | `seed` | integer | `123` | Seeds Python, NumPy, and PyTorch. It also seeds a config-defined random split unless `split.seed` is set. | | `progress` | boolean | `true` | Whether to print data-processing, batch, epoch, and evaluation progress. | | `debug_epoch_predictions` | boolean | `false` | Save deterministic end-of-epoch train and validation predictions to `debug_epoch_predictions.csv`. This adds one evaluation pass over the training split per epoch. | +| `head_layernorm` | boolean | `false` | For `saluki_exact`, replace both dense-head BatchNorm layers with per-example LayerNorm. Other model types reject `true`. | | `sequence_controls` | mapping, list, or `null` | `null` | Optional sequence ablations applied before split selection. | | `split_source` | string | `"auto"` | Whether splits come from the bundle or from the `split` block. | | `split` | mapping | random 80/10/10 | Config-defined split settings, used according to `split_source`. | @@ -244,7 +246,25 @@ This is the model used by the standard Saluki workflow. | `augment_shift` | `3` | Maximum random right shift applied during training. The sampled shift is between zero and this value; set to zero to disable it. | | `ln_epsilon` | `0.007` | Numerical epsilon used by channel layer normalization. | | `keras_bn_momentum` | `0.9` | Batch-normalization momentum expressed using the Keras convention reproduced by this model. | -| `bn_eps` | `0.001` | Numerical epsilon used by batch-normalization layers in the head. | +| `bn_eps` | `0.001` | Numerical epsilon used by normalization layers in the head. | +| `head_layernorm` | `false` | Checkpoint-level record of whether the head uses LayerNorm. During training, set the top-level `head_layernorm` field instead. | + +Set the top-level training option to enable the experimental head: + +```json +{ + "model": {"name": "saluki_exact", "params": {}}, + "head_layernorm": true +} +``` + +This preserves the original `Normalization → ReLU → Linear → Dropout → +Normalization → ReLU → Linear` ordering, but makes both head normalization +layers independent of batch and running statistics. LayerNorm uses `bn_eps` +so enabling the option changes the normalization behavior without also +changing its numerical epsilon. The resolved value is saved in checkpoint +`model_config.params`, allowing the checkpoint loader to reconstruct the +correct head, and is also recorded in `summary.json`. ### `saluki_like` diff --git a/scripts/example_train_config.json b/scripts/example_train_config.json index d2aa9b9..8a78cb4 100644 --- a/scripts/example_train_config.json +++ b/scripts/example_train_config.json @@ -23,5 +23,6 @@ "num_workers": 0, "mmap_mode": "r", "seed": 42, - "debug_epoch_predictions": false + "debug_epoch_predictions": false, + "head_layernorm": false } diff --git a/src/transcriptml/models/reproduce.py b/src/transcriptml/models/reproduce.py index 3986237..cc91e74 100644 --- a/src/transcriptml/models/reproduce.py +++ b/src/transcriptml/models/reproduce.py @@ -20,6 +20,7 @@ class SalukiExactConfig: ln_epsilon: float = 0.007 keras_bn_momentum: float = 0.90 bn_eps: float = 1e-3 + head_layernorm: bool = False def to_kwargs(self) -> dict[str, object]: """Return constructor keyword arguments for ``SalukiExact``.""" @@ -94,6 +95,7 @@ def __init__( ln_epsilon: float = 0.007, keras_bn_momentum: float = 0.90, bn_eps: float = 1e-3, + head_layernorm: bool = False, ): """Create the Saluki architecture reproduction. @@ -106,7 +108,10 @@ def __init__( augment_shift: Maximum stochastic right shift during training. ln_epsilon: Epsilon used by channel layer normalization. keras_bn_momentum: Keras-style batch-normalization momentum value. - bn_eps: Epsilon used by batch-normalization layers. + bn_eps: Epsilon used by normalization layers in the prediction + head. + head_layernorm: Whether to replace the two head batch-normalization + layers with per-example layer normalization. """ super().__init__() @@ -114,6 +119,7 @@ def __init__( self.filters = int(filters) self.kernel_size = int(kernel_size) self.num_layers = int(num_layers) + self.head_layernorm = bool(head_layernorm) bn_momentum_pt = 1.0 - float(keras_bn_momentum) self.shift = StochasticShift(augment_shift) self.conv0 = nn.Conv1d(seq_depth, filters, kernel_size=kernel_size, padding=0, bias=False) @@ -133,10 +139,14 @@ def __init__( self.pre_rnn_ln = ChannelLayerNorm(filters, eps=ln_epsilon) self.pre_rnn_act = nn.ReLU() self.gru = nn.GRU(input_size=filters, hidden_size=filters, batch_first=True) - self.bn1 = nn.BatchNorm1d(filters, eps=bn_eps, momentum=bn_momentum_pt) + if self.head_layernorm: + self.bn1 = nn.LayerNorm(filters, eps=bn_eps) + self.bn2 = nn.LayerNorm(filters, eps=bn_eps) + else: + self.bn1 = nn.BatchNorm1d(filters, eps=bn_eps, momentum=bn_momentum_pt) + self.bn2 = nn.BatchNorm1d(filters, eps=bn_eps, momentum=bn_momentum_pt) self.fc1 = nn.Linear(filters, filters) self.drop1 = nn.Dropout(dropout) - self.bn2 = nn.BatchNorm1d(filters, eps=bn_eps, momentum=bn_momentum_pt) self.fc2 = nn.Linear(filters, 1) self.reset_parameters() diff --git a/src/transcriptml/training/trainer.py b/src/transcriptml/training/trainer.py index 00ac296..819fc54 100644 --- a/src/transcriptml/training/trainer.py +++ b/src/transcriptml/training/trainer.py @@ -43,6 +43,7 @@ class TrainConfig: seed: int = 123 progress: bool = True debug_epoch_predictions: bool = False + head_layernorm: bool = False sequence_controls: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None = None split_source: str = "auto" split: Mapping[str, Any] = field( @@ -247,6 +248,12 @@ def _would_create_singleton_batch(n_examples: int, batch_size: int) -> bool: return n % bs == 1 +def _has_batch_normalization(model: nn.Module) -> bool: + """Return whether a model contains a PyTorch batch-normalization module.""" + + return any(isinstance(module, nn.modules.batchnorm._BatchNorm) for module in model.modules()) + + def _run_loader( model: nn.Module, loader: DataLoader | None, @@ -526,10 +533,16 @@ def train_model(bundle: DatasetBundle, config: TrainConfig | Mapping[str, Any]) splits, split_source_used = _select_splits(bundle, cfg) split_counts = {name: len(splits.get(name, [])) for name in ("train", "val", "test")} model_config = normalize_model_config(cfg.model) + if cfg.head_layernorm and model_config.name != "saluki_exact": + raise ValueError("head_layernorm is only supported for model 'saluki_exact'") + if model_config.name == "saluki_exact": + model_config.params = dict(model_config.params or {}) + model_config.params["head_layernorm"] = bool(cfg.head_layernorm) log_progress( ( "training: " f"device={device}, output={out}, loss={normalized_loss_config['name']}, " + f"head_layernorm={cfg.head_layernorm}, " f"split_source={split_source_used} requested={cfg.split_source}, " f"train={split_counts['train']}, val={split_counts['val']}, " f"test={split_counts['test']}" @@ -540,7 +553,9 @@ def train_model(bundle: DatasetBundle, config: TrainConfig | Mapping[str, Any]) optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.learning_rate, weight_decay=cfg.weight_decay) dataset = _ArrayRegressionDataset(bundle.X, y_train, aux_arrays) pin_memory = device.type == "cuda" - drop_last_train = _would_create_singleton_batch(len(splits["train"]), cfg.batch_size) + drop_last_train = _has_batch_normalization(model) and _would_create_singleton_batch( + len(splits["train"]), cfg.batch_size + ) if drop_last_train: log_progress( ( @@ -748,6 +763,7 @@ def train_model(bundle: DatasetBundle, config: TrainConfig | Mapping[str, Any]) "best_monitor_values": best_metrics, "epochs_run": len(history), "loss": normalized_loss_config, + "head_layernorm": bool(cfg.head_layernorm), "split_source_requested": cfg.split_source, "split_source_used": split_source_used, "split_counts": split_counts, diff --git a/src/transcriptml/workflows/init_run.py b/src/transcriptml/workflows/init_run.py index ec94489..f278457 100644 --- a/src/transcriptml/workflows/init_run.py +++ b/src/transcriptml/workflows/init_run.py @@ -26,6 +26,7 @@ def _train_config(workflow: str) -> dict[str, Any]: "num_workers": 0, "mmap_mode": "r", "seed": 42, + "head_layernorm": False, "split_source": "auto", "split": {"method": "random", "val_frac": 0.1, "test_frac": 0.1}, } diff --git a/tests/test_cli_analysis.py b/tests/test_cli_analysis.py index 2eb6ab6..d1bb9c6 100644 --- a/tests/test_cli_analysis.py +++ b/tests/test_cli_analysis.py @@ -18,6 +18,7 @@ def test_models_cli_list_and_show_json(capsys): payload = json.loads(capsys.readouterr().out) assert payload["name"] == "saluki_exact" assert payload["params"]["filters"] == 64 + assert payload["params"]["head_layernorm"] is False def test_init_run_cli_writes_templates(tmp_path): @@ -37,6 +38,7 @@ def test_init_run_cli_writes_templates(tmp_path): assert train_config["num_workers"] == 0 assert train_config["mmap_mode"] == "r" assert train_config["seed"] == 42 + assert train_config["head_layernorm"] is False assert train_config["split_source"] == "auto" assert not (out_dir / "run_config.json").exists() assert (out_dir / "README.md").exists() diff --git a/tests/test_splits_models.py b/tests/test_splits_models.py index 90f7f19..3036289 100644 --- a/tests/test_splits_models.py +++ b/tests/test_splits_models.py @@ -87,6 +87,41 @@ def test_model_registry_dummy_forward(): ) assert saluki(x6).shape == (2,) + saluki_exact = build_model( + { + "name": "saluki_exact", + "params": { + "filters": 8, + "kernel_size": 3, + "num_layers": 1, + "dropout": 0.0, + "augment_shift": 0, + }, + } + ) + assert isinstance(saluki_exact.bn1, torch.nn.BatchNorm1d) + assert isinstance(saluki_exact.bn2, torch.nn.BatchNorm1d) + + saluki_exact_ln = build_model( + { + "name": "saluki_exact", + "params": { + "filters": 8, + "kernel_size": 3, + "num_layers": 1, + "dropout": 0.0, + "augment_shift": 0, + "bn_eps": 0.002, + "head_layernorm": True, + }, + } + ) + assert isinstance(saluki_exact_ln.bn1, torch.nn.LayerNorm) + assert isinstance(saluki_exact_ln.bn2, torch.nn.LayerNorm) + assert saluki_exact_ln.bn1.eps == pytest.approx(0.002) + saluki_exact_ln.train() + assert saluki_exact_ln(x6[:1]).shape == (1,) + legnet = build_model( { "name": "legnet", diff --git a/tests/test_training_losses.py b/tests/test_training_losses.py index 5a733e3..a2981b3 100644 --- a/tests/test_training_losses.py +++ b/tests/test_training_losses.py @@ -7,6 +7,7 @@ import transcriptml.training.trainer as trainer from transcriptml.data.bundle import DatasetBundle +from transcriptml.models.registry import load_checkpoint from transcriptml.training.losses import build_training_loss from transcriptml.training.trainer import train_model @@ -83,6 +84,13 @@ def _tiny_x(n: int) -> np.ndarray: return x +def _tiny_saluki_x(n: int) -> np.ndarray: + x = np.zeros((n, 6, 32), dtype=np.uint8) + for i in range(n): + x[i, i % 4, :] = 1 + return x + + def _tiny_model_config() -> dict[str, object]: return { "name": "small_cnn", @@ -138,6 +146,70 @@ def test_train_model_default_mse_remains_compatible(tmp_path): assert not (tmp_path / "debug_epoch_predictions.csv").exists() +def test_train_model_saluki_head_layernorm_checkpoint_roundtrip(tmp_path): + bundle = DatasetBundle( + X=_tiny_saluki_x(4), + y=np.linspace(-1.0, 1.0, 4, dtype=np.float32), + schema="saluki6", + splits={"train": [0], "val": [1, 2], "test": [3]}, + ) + + result = train_model( + bundle, + { + "dataset": "unused", + "output_dir": str(tmp_path), + "model": { + "name": "saluki_exact", + "params": { + "filters": 4, + "kernel_size": 3, + "num_layers": 1, + "dropout": 0.0, + "augment_shift": 0, + }, + }, + "head_layernorm": True, + "batch_size": 1, + "epochs": 1, + "progress": False, + }, + ) + + assert isinstance(result["model"].bn1, torch.nn.LayerNorm) + assert isinstance(result["model"].bn2, torch.nn.LayerNorm) + assert np.isfinite(result["history"][0]["train_loss"]) + assert result["summary"]["head_layernorm"] is True + + loaded_model, checkpoint = load_checkpoint(tmp_path / "best.pt") + assert isinstance(loaded_model.bn1, torch.nn.LayerNorm) + assert isinstance(loaded_model.bn2, torch.nn.LayerNorm) + assert checkpoint["model_config"]["params"]["head_layernorm"] is True + assert checkpoint["train_config"]["head_layernorm"] is True + + +def test_train_model_rejects_head_layernorm_for_other_models(tmp_path): + bundle = DatasetBundle( + X=_tiny_x(4), + y=np.linspace(-1.0, 1.0, 4, dtype=np.float32), + schema="rna4", + splits={"train": [0, 1], "val": [2], "test": [3]}, + ) + + with pytest.raises(ValueError, match="only supported.*saluki_exact"): + train_model( + bundle, + { + "dataset": "unused", + "output_dir": str(tmp_path), + "model": _tiny_model_config(), + "head_layernorm": True, + "epochs": 1, + "progress": False, + }, + ) + + def test_train_model_debug_epoch_predictions_csv(tmp_path): y = np.linspace(-1.0, 1.0, 8, dtype=np.float32) bundle = DatasetBundle( From f92d6b62838bc3f2ebf040c16be2a1ab5d32d0e6 Mon Sep 17 00:00:00 2001 From: isvock Date: Thu, 30 Jul 2026 16:31:12 -0700 Subject: [PATCH 2/2] Allow more flexibility in Saluki model --- docs/training_configuration.md | 26 ++++++++++++++++++- src/transcriptml/models/reproduce.py | 13 +++++++++- tests/test_cli_analysis.py | 2 ++ tests/test_splits_models.py | 39 ++++++++++++++++++++++++++++ 4 files changed, 78 insertions(+), 2 deletions(-) diff --git a/docs/training_configuration.md b/docs/training_configuration.md index 748f400..4ba0772 100644 --- a/docs/training_configuration.md +++ b/docs/training_configuration.md @@ -226,6 +226,7 @@ This is the model used by the standard Saluki workflow. "filters": 64, "kernel_size": 5, "num_layers": 6, + "pooling": "max", "dropout": 0.3, "augment_shift": 3, "ln_epsilon": 0.007, @@ -241,7 +242,8 @@ This is the model used by the standard Saluki workflow. | `seq_depth` | `6` | Number of input channels. Keep this at six for an ordinary Saluki bundle. | | `filters` | `64` | Width of the convolutional stack, GRU, and dense hidden layer. | | `kernel_size` | `5` | Width of the initial and repeated one-dimensional convolutions. | -| `num_layers` | `6` | Number of convolution, dropout, and max-pooling blocks after the initial convolution. | +| `num_layers` | `6` | Number of convolution, dropout, and pooling blocks after the initial convolution. | +| `pooling` | `"max"` | Downsampling operation in each convolutional block. Use `"max"` or `"average"`. | | `dropout` | `0.3` | Dropout probability in convolutional blocks and the dense head. | | `augment_shift` | `3` | Maximum random right shift applied during training. The sampled shift is between zero and this value; set to zero to disable it. | | `ln_epsilon` | `0.007` | Numerical epsilon used by channel layer normalization. | @@ -266,6 +268,28 @@ changing its numerical epsilon. The resolved value is saved in checkpoint `model_config.params`, allowing the checkpoint loader to reconstruct the correct head, and is also recorded in `summary.json`. +To use average pooling and disable stochastic shift augmentation, set the +corresponding model parameters: + +```json +{ + "model": { + "name": "saluki_exact", + "params": { + "pooling": "average", + "augment_shift": 0 + } + } +} +``` + +During each training forward pass, a positive `augment_shift` samples one +integer offset from zero through the configured maximum. All channels are +shifted right together, zeros are inserted at the left boundary, and the same +number of positions are removed from the right boundary. Evaluation mode never +applies the shift. Setting `augment_shift` to zero bypasses the operation in +training mode as well. + ### `saluki_like` `saluki_like` is a compact convolutional/GRU model inspired by Saluki rather diff --git a/src/transcriptml/models/reproduce.py b/src/transcriptml/models/reproduce.py index cc91e74..c9238b2 100644 --- a/src/transcriptml/models/reproduce.py +++ b/src/transcriptml/models/reproduce.py @@ -15,6 +15,7 @@ class SalukiExactConfig: filters: int = 64 kernel_size: int = 5 num_layers: int = 6 + pooling: str = "max" dropout: float = 0.3 augment_shift: int = 3 ln_epsilon: float = 0.007 @@ -90,6 +91,7 @@ def __init__( filters: int = 64, kernel_size: int = 5, num_layers: int = 6, + pooling: str = "max", dropout: float = 0.3, augment_shift: int = 3, ln_epsilon: float = 0.007, @@ -104,6 +106,8 @@ def __init__( filters: Number of convolutional and recurrent feature channels. kernel_size: Width of the convolution kernels. num_layers: Number of repeated convolution/pooling blocks. + pooling: Downsampling operation used by each convolutional block. + Supported values are ``"max"`` and ``"average"``. dropout: Dropout probability used in convolutional and dense layers. augment_shift: Maximum stochastic right shift during training. ln_epsilon: Epsilon used by channel layer normalization. @@ -119,6 +123,13 @@ def __init__( self.filters = int(filters) self.kernel_size = int(kernel_size) self.num_layers = int(num_layers) + self.pooling = str(pooling).strip().lower() + if self.pooling == "max": + pool_cls = nn.MaxPool1d + elif self.pooling == "average": + pool_cls = nn.AvgPool1d + else: + raise ValueError("pooling must be either 'max' or 'average'") self.head_layernorm = bool(head_layernorm) bn_momentum_pt = 1.0 - float(keras_bn_momentum) self.shift = StochasticShift(augment_shift) @@ -132,7 +143,7 @@ def __init__( "act": nn.ReLU(), "conv": nn.Conv1d(filters, filters, kernel_size=kernel_size, padding=0), "drop": nn.Dropout(dropout), - "pool": nn.MaxPool1d(kernel_size=2, stride=2), + "pool": pool_cls(kernel_size=2, stride=2), } ) ) diff --git a/tests/test_cli_analysis.py b/tests/test_cli_analysis.py index d1bb9c6..a3fbd44 100644 --- a/tests/test_cli_analysis.py +++ b/tests/test_cli_analysis.py @@ -18,6 +18,8 @@ def test_models_cli_list_and_show_json(capsys): payload = json.loads(capsys.readouterr().out) assert payload["name"] == "saluki_exact" assert payload["params"]["filters"] == 64 + assert payload["params"]["pooling"] == "max" + assert payload["params"]["augment_shift"] == 3 assert payload["params"]["head_layernorm"] is False diff --git a/tests/test_splits_models.py b/tests/test_splits_models.py index 3036289..29b412b 100644 --- a/tests/test_splits_models.py +++ b/tests/test_splits_models.py @@ -3,6 +3,7 @@ import torch from transcriptml.data.bundle import DatasetBundle +from transcriptml.models.reproduce import StochasticShift from transcriptml.models.registry import build_model from transcriptml.training.trainer import TrainConfig, _monitor_improved, _monitor_names, _select_splits from transcriptml.training.splits import predefined_split_indices, random_split_indices @@ -73,6 +74,18 @@ def test_multiple_monitor_metrics_use_or_improvement(): assert not improved +def test_stochastic_shift_can_be_disabled(): + x = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4) + + disabled = StochasticShift(shift_max=0) + disabled.train() + assert torch.equal(disabled(x), x) + + enabled = StochasticShift(shift_max=3) + enabled.eval() + assert torch.equal(enabled(x), x) + + def test_model_registry_dummy_forward(): x4 = torch.randn(2, 4, 32) small = build_model({"name": "small_cnn", "params": {"in_ch": 4, "n_filters": 8, "head_hidden": 8}}) @@ -101,6 +114,8 @@ def test_model_registry_dummy_forward(): ) assert isinstance(saluki_exact.bn1, torch.nn.BatchNorm1d) assert isinstance(saluki_exact.bn2, torch.nn.BatchNorm1d) + assert saluki_exact.pooling == "max" + assert all(isinstance(block["pool"], torch.nn.MaxPool1d) for block in saluki_exact.blocks) saluki_exact_ln = build_model( { @@ -122,6 +137,30 @@ def test_model_registry_dummy_forward(): saluki_exact_ln.train() assert saluki_exact_ln(x6[:1]).shape == (1,) + saluki_exact_average_pool = build_model( + { + "name": "saluki_exact", + "params": { + "filters": 8, + "kernel_size": 3, + "num_layers": 1, + "pooling": "average", + "dropout": 0.0, + "augment_shift": 0, + "head_layernorm": True, + }, + } + ) + assert saluki_exact_average_pool.pooling == "average" + assert all( + isinstance(block["pool"], torch.nn.AvgPool1d) + for block in saluki_exact_average_pool.blocks + ) + assert saluki_exact_average_pool(x6).shape == (2,) + + with pytest.raises(ValueError, match="pooling must be either 'max' or 'average'"): + build_model({"name": "saluki_exact", "params": {"pooling": "median"}}) + legnet = build_model( { "name": "legnet",