More flexibility in Saluki model

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`. |
@@ -224,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,
@@ -239,12 +242,53 @@ 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. |
 | `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`.
+
+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`
 

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
 }

src/transcriptml/models/reproduce.py

@@ -15,11 +15,13 @@ 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
     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``."""
@@ -89,11 +91,13 @@ 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,
         keras_bn_momentum: float = 0.90,
         bn_eps: float = 1e-3,
+        head_layernorm: bool = False,
     ):
         """Create the Saluki architecture reproduction.
 
@@ -102,18 +106,31 @@ 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.
             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__()
         self.seq_depth = int(seq_depth)
         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)
         self.conv0 = nn.Conv1d(seq_depth, filters, kernel_size=kernel_size, padding=0, bias=False)
@@ -126,17 +143,21 @@ 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),
                     }
                 )
             )
         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()
 

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,

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},
         }

tests/test_cli_analysis.py

@@ -18,6 +18,9 @@ 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
 
 
 def test_init_run_cli_writes_templates(tmp_path):
@@ -37,6 +40,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()

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}})
@@ -87,6 +100,67 @@ 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)
+    assert saluki_exact.pooling == "max"
+    assert all(isinstance(block["pool"], torch.nn.MaxPool1d) for block in saluki_exact.blocks)
+
+    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,)
+
+    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",