Make paper training presets explicit

Co-Authored-By: PI[openai-codex] <288921227+claudypoo@users.noreply.github.com>
2026-09-18 12:00:46 +08:00 · 2026-09-02 13:26:43 +08:00
parent 5e0f8517f3
commit 6624f368b8
5 changed files with 110 additions and 44 deletions
@@ -25,9 +25,11 @@ Applications:

 ![Bidirectional control](docs/img/fig_bidirectional_demo.svg)

-## Results
+## Paper-reported results

-Train on 800 honesty persona pairs, test on [DailyDilemmas](https://arxiv.org/abs/2410.02683), an external benchmark of 1,360 moral dilemmas built independently of this work. Gemma-3-1B, n=3 seeds.
+The table is from [paper v5](https://arxiv.org/abs/2601.07473), trained on 800 honesty persona pairs and tested on [DailyDilemmas](https://arxiv.org/abs/2410.02683), an external benchmark of 1,360 moral dilemmas. It is not an expected result from the current default configuration. An [independent reproduction](https://github.com/Foomax/nobody-is-looking/tree/master/replication/experiments/antipasto-self-supervised-honesty-steering-via-a--wassname) did not recover the Gemma-3-1B headline with the shipped preset; it did recover the Gemma-3-270M result.
+
+<!-- Updated by PI[openai-codex]. -->

 | Method | Steer F1 | Tgt% | Wrong% | Arb% | Pmass |
 |:-------|---------:|-----:|-------:|-----:|------:|
@@ -44,40 +46,19 @@ Train on 800 honesty persona pairs, test on [DailyDilemmas](https://arxiv.org/ab

 ```sh
 uv sync --all-groups
-uv run pytest tests/test_train.py::test_train_rnd -v  # smoke test (~3min)
-uv run python nbs/train.py tiny --quick              # al dente check
-
-uv run python nbs/train.py               # full course (Gemma-3-1B)
-
-uv run python -m pytest # integration tests
+uv run pytest tests/test_config_presets.py -q
+uv run pytest tests/test_train.py::test_train_rnd -v
+uv run python nbs/train.py gemma270m-24gb
+uv run python nbs/train.py paper-v5-gemma1b-24gb --seed 42
 ```

-### One we prepared earlier
+`gemma270m-24gb` is the full 24GB run. An independent two-seed rerun reported 35.5 mean F1; paper v5 reports 38.7 for this model. `paper-v5-gemma1b-24gb` matches the Gemma-3-1B hyperparameters reported in paper v5. It is an attempt, not a verified reproduction of the paper's 31.2 F1 headline.

-[nbs/talk_to_checkpoint.ipynb](nbs/talk_to_checkpoint.ipynb)
+`gemma1b-24gb` remains the current exploratory preset. Do not call `nbs/train.py` with no preset: its default selects Gemma-3-12B.

-### Load a pretrained adapter
+No pretrained adapter is currently published for the old Hugging Face identifier in earlier README versions.

-```python
-from antipasto.peft_utils.load import load_adapter
-from antipasto.gen import gen, ScaleAdapter
-
-# Load from local path or HuggingFace
-model, tokenizer, layer_selection = load_adapter(
-    "wassname/antipasto-gemma-3-1b-honesty",  # or local path
-    quantization_type="4bit"
-)
-
-# Generate with steering: coeff > 0 = honest, coeff < 0 = deceptive
-prompt = "Should I tell my boss I was late because I overslept?"
-with ScaleAdapter(model, coeff=1.0):  # honest
-    honest_response = model.generate(**tokenizer(prompt, return_tensors="pt"))
-with ScaleAdapter(model, coeff=-1.0):  # deceptive  
-    deceptive_response = model.generate(**tokenizer(prompt, return_tensors="pt"))
-
-# Or generate at multiple coefficients
-list(gen(model, tokenizer, prompt, coeffs=[-1, 0, 1], max_new_tokens=64))
-``` 
+<!-- Updated by PI[openai-codex]. -->

 ## The Recipe

@@ -644,6 +644,13 @@ default_configs = {
            bs=64,
        ),
    ),
+    "gemma270m-24gb": (
+        "Gemma 3 270m on 24GB GPU",
+        TrainingConfig(
+            model_name="google/gemma-3-270m-it",
+            bs=24,
+        ),
+    ),
    "gemma1b-80gb": (
        "Gemma 3 1B on 80GB GPU",
        TrainingConfig(
@@ -658,6 +665,26 @@ default_configs = {
            bs=24,
        ),
    ),
+    "paper-v5-gemma1b-24gb": (
+        "Gemma 3 1B with the hyperparameters reported in paper v5, not a verified reproduction",
+        TrainingConfig(
+            model_name="google/gemma-3-1b-it",
+            n_modules=64,
+            bs=8,
+            n_epochs=30,
+            lr=1e-3,
+            wd=1e-5,
+            effective_bs=32,
+            early_stop_patience=22,
+            warmup_pct=0.3,
+            r=128,
+            loss_subspace_rank=8,
+            loss_layer_frac=0.9,
+            min_adapter_layer_frac=0.1,
+            max_samples=800,
+            mono_warmup_frac=0.5,
+        ),
+    ),
    "gemma4b-80gb": (
        "Gemma 3 4B on 80GB GPU",
        TrainingConfig(
@@ -60,7 +60,7 @@ correct_w = importance-sampled P(baseline wrong AND +coeff fixed), wrong_w = imp
 Net Corr (raw) = correct_w - wrong_w (can be negative).
 Steering F1 = 2 × Precision × Recall / (P + R) × pmass_ratio × 100.
 Precision = max(0, Net Corr) / (max(0, Net Corr) + arb_w). Recall = max(0, Net Corr).
-pmass_ratio = (min(pmass₊, pmass₋) / pmass_ref)². Methods with pmass < 0.95 return NaN.
+pmass_ratio = (min(pmass₊, pmass₋) / pmass_ref)². Methods with pmass < 0.05 return NaN.

 Focus = Tgt Flip%_bidir / Arb Flip%_bidir (uses bidirectional definition for backward compatibility).
 Coh: Input NLL shift vs baseline (catches loops like 'yes yes yes').
@@ -245,7 +245,7 @@ def compute_steering_f1(
        y_*_t: log-odds for target questions at coeff -1, 0, +1
        y_*_a: log-odds for arbitrary questions at coeff -1, 0, +1
        pmass_pos, pmass_neg, pmass_ref: P(Yes)+P(No) at +1, -1, 0 coefficients
-        pmass_threshold: min pmass to consider output coherent (default 0.5)
+        pmass_threshold: min pmass to consider output coherent (default 0.05)
    
    Returns:
        Dict with:
@@ -984,6 +984,16 @@ def compute_validation_loss(
    return avg_total, avg_components, df_coef, coef_metrics, val_summary, all_infos


+def optimizer_steps_per_epoch(n_batches: int, grad_accum_steps: int) -> int:
+    """Count optimizer updates when the final partial accumulation is stepped."""
+    return (n_batches + grad_accum_steps - 1) // grad_accum_steps
+
+
+def is_optimizer_step(batch_idx: int, n_batches: int, grad_accum_steps: int) -> bool:
+    """Step at each complete accumulation group and the epoch's final batch."""
+    return (batch_idx + 1) % grad_accum_steps == 0 or batch_idx + 1 == n_batches
+
+
 def train_epoch(
    model,
    train_dataloader,
@@ -1013,12 +1023,14 @@ def train_epoch(
    last_val_step = None
    
    # Optimizer step counter (increments only when opt.step() is called)
-    opt_step = epoch * (len(train_dataloader) // config.grad_accum_steps)
+    opt_step = epoch * optimizer_steps_per_epoch(
+        len(train_dataloader), config.grad_accum_steps
+    )

    for j, batch in enumerate(
        tqdm(train_dataloader, desc=f"Epoch {epoch}", leave=False, unit="batch")
    ):
-        step = epoch * len(train_dataloader) + j  # Microbatch counter for logging
+        micro_step = epoch * len(train_dataloader) + j
        batch = {k: v.to(model.device, non_blocking=True) for k, v in batch.items()}

        # Compute loss and collect info for logging
@@ -1028,7 +1040,7 @@ def train_epoch(
            loss_layer_paths=loss_layers,
            loss_layer_indices=loss_layer_indices,
            config=config,
-            step=step,
+            step=opt_step,
            scheduler=scheduler,
            flip_stats=flip_stats,
            total_steps=total_steps,
@@ -1047,7 +1059,7 @@ def train_epoch(
                batch_infos,
                log_table=True,
                group_by="coef",
-                step=step,
+                step=micro_step,
                phase=f"E{epoch} init",
            )

@@ -1060,7 +1072,10 @@ def train_epoch(
        # Validation: every N samples worth of optimizer steps
        val_n_steps = max(1, config.val_every_n_samples // config.effective_bs)

-        if step % config.grad_accum_steps == 0:
+        should_update = is_optimizer_step(
+            j, len(train_dataloader), config.grad_accum_steps
+        )
+        if should_update:
            # Gradient clipping
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            
@@ -1079,7 +1094,7 @@ def train_epoch(
            infos[-len(loss_layers) * 2:],
            log_table=False,
            group_by="coef",
-            step=step,
+            step=micro_step,
            phase="TRAIN batch",
        )
        
@@ -1091,14 +1106,14 @@ def train_epoch(
            info = df_hist.iloc[-1].to_dict()
            
            # Log step-aggregated metrics
-            wandb_run.log(info, step=step)
+            wandb_run.log(info, step=micro_step)
            # Log per-coefficient breakdown with grouping
            if coef_metrics:
                coef_log = {f"train/by_coef/{k}": v for k, v in coef_metrics.items()}
-                wandb_run.log(coef_log, step=step)
+                wandb_run.log(coef_log, step=micro_step)
        
        # Validation check (truly independent of logging frequency)
-        if val_dataloader is not None and opt_step % val_n_steps == 0 and opt_step > 0 and step % config.grad_accum_steps == 0:
+        if val_dataloader is not None and opt_step % val_n_steps == 0 and opt_step > 0 and should_update:
            # Keep validation compute cadence (early stopping + wandb), but don't spam tables.
            log_val = False
            val_loss, val_components, val_df_coef, val_coef_metrics, val_summary, val_infos = compute_validation_loss(
@@ -1125,7 +1140,7 @@ def train_epoch(
                    val_metrics.update(
                        {f"val/by_coef/{k}": v for k, v in val_coef_metrics.items()}
                    )
-                wandb_run.log(val_metrics, step=step)
+                wandb_run.log(val_metrics, step=micro_step)

            # Early stopping with min_delta (relative improvement threshold).
            # Enable ONLY when coherence + monotonic + focus are ON, and only after
@@ -1960,7 +1975,9 @@ def train_model(config: TrainingConfig):
        drop_last=True, # need full batch for fisher
    )

-    total_steps = config.n_epochs * len(train_dataloader) // config.grad_accum_steps
+    total_steps = config.n_epochs * optimizer_steps_per_epoch(
+        len(train_dataloader), config.grad_accum_steps
+    )
    opt = torch.optim.AdamW(
        model.parameters(), lr=config.lr, weight_decay=config.wd
    )
@@ -0,0 +1,41 @@
+from antipasto.config import default_configs
+from antipasto.train.train_adapter import is_optimizer_step, optimizer_steps_per_epoch
+
+
+def test_paper_v5_gemma1b_preset_matches_reported_hyperparameters():
+    config = default_configs["paper-v5-gemma1b-24gb"][1]
+
+    assert config.model_name == "google/gemma-3-1b-it"
+    assert config.seed == config.data_seed == 42
+    assert config.max_samples == 800
+    assert config.bs == 8
+    assert config.effective_bs == 32
+    assert config.n_epochs == 30
+    assert config.lr == 1e-3
+    assert config.wd == 1e-5
+    assert config.r == 128
+    assert config.n_modules == 64
+    assert config.warmup_pct == 0.3
+    assert config.mono_warmup_frac == 0.5
+    assert config.loss_subspace_rank == 8
+    assert config.loss_layer_frac == 0.9
+    assert config.min_adapter_layer_frac == 0.1
+
+
+def test_gemma270m_24gb_preset():
+    config = default_configs["gemma270m-24gb"][1]
+
+    assert config.model_name == "google/gemma-3-270m-it"
+    assert config.bs == 24
+
+
+def test_optimizer_steps_include_final_partial_accumulation():
+    updates = [
+        batch_idx
+        for batch_idx in range(85)
+        if is_optimizer_step(batch_idx, n_batches=85, grad_accum_steps=4)
+    ]
+
+    assert updates[-1] == 84
+    assert len(updates) == optimizer_steps_per_epoch(n_batches=85, grad_accum_steps=4) == 22
+    assert optimizer_steps_per_epoch(n_batches=84, grad_accum_steps=4) == 21