Fix chosen-token regularization

auto_antislop_config.yaml

@@ -259,7 +259,8 @@ finetune_shuffle_seed: 666
 # 0 = off; 0.9 strongly downsamples overrepresented rule violations
 # (this is useful because the raw generated dataset is typically very skewed)
 ftpo_sample_rejected_regularisation_strength: 0.8
-ftpo_sample_chosen_regularisation_strength: 0.2
+# 0 = off; positive values trim globally overrepresented chosen-token slots
+ftpo_sample_chosen_regularisation_strength: 0.0
 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
 
 # ── FTPO-specific hyper-parameters ─────────────────────────────────────────

configs/GLM-4-32B-0414.yaml

@@ -265,7 +265,8 @@ finetune_shuffle_seed: 666
 # 0 = off; 0.9 strongly downsamples overrepresented rule violations
 # (this is useful because the raw generated dataset is typically very skewed)
 ftpo_sample_rejected_regularisation_strength: 0.8
-ftpo_sample_chosen_regularisation_strength: 0.2
+# 0 = off; positive values trim globally overrepresented chosen-token slots
+ftpo_sample_chosen_regularisation_strength: 0.0
 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
 
 

configs/Mistral-Small-3.2-24B-Instruct-2506.yaml

@@ -265,7 +265,8 @@ finetune_shuffle_seed: 666
 # 0 = off; 0.9 strongly downsamples overrepresented rule violations
 # (this is useful because the raw generated dataset is typically very skewed)
 ftpo_sample_rejected_regularisation_strength: 0.7
-ftpo_sample_chosen_regularisation_strength: 0.2
+# 0 = off; positive values trim globally overrepresented chosen-token slots
+ftpo_sample_chosen_regularisation_strength: 0.0
 ftpo_sample_min_chosen_tokens: 3 # filter out ftpo samples that have fewer than this number in the chosen tokens list
 
 

configs/default.yaml

@@ -259,7 +259,8 @@ finetune_shuffle_seed: 666
 # 0 = off; 0.9 strongly downsamples overrepresented rule violations
 # (this is useful because the raw generated dataset is typically very skewed)
 ftpo_sample_rejected_regularisation_strength: 0.8
-ftpo_sample_chosen_regularisation_strength: 0.2
+# 0 = off; positive values trim globally overrepresented chosen-token slots
+ftpo_sample_chosen_regularisation_strength: 0.0
 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
 
 # ── FTPO-specific hyper-parameters ─────────────────────────────────────────

configs/gemma-3-12b-it.yaml

@@ -258,7 +258,8 @@ finetune_shuffle_seed: 666
 # 0 = off; 0.9 strongly downsamples overrepresented rule violations
 # (this is useful because the raw generated dataset is typically very skewed)
 ftpo_sample_rejected_regularisation_strength: 0.8
-ftpo_sample_chosen_regularisation_strength: 0.2
+# 0 = off; positive values trim globally overrepresented chosen-token slots
+ftpo_sample_chosen_regularisation_strength: 0.0
 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
 
 

configs/gemma-3-27b-it.yaml

@@ -258,7 +258,8 @@ finetune_shuffle_seed: 666
 # 0 = off; 0.9 strongly downsamples overrepresented rule violations
 # (this is useful because the raw generated dataset is typically very skewed)
 ftpo_sample_rejected_regularisation_strength: 0.6
-ftpo_sample_chosen_regularisation_strength: 0.2
+# 0 = off; positive values trim globally overrepresented chosen-token slots
+ftpo_sample_chosen_regularisation_strength: 0.0
 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
 
 

configs/gemma-3-4b-it.yaml

@@ -258,7 +258,8 @@ finetune_shuffle_seed: 666
 # 0 = off; 0.9 strongly downsamples overrepresented rule violations
 # (this is useful because the raw generated dataset is typically very skewed)
 ftpo_sample_rejected_regularisation_strength: 0.8
-ftpo_sample_chosen_regularisation_strength: 0.2
+# 0 = off; positive values trim globally overrepresented chosen-token slots
+ftpo_sample_chosen_regularisation_strength: 0.0
 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
 
 

utils/dataset_helpers.py

@@ -16,6 +16,82 @@ logger = logging.getLogger(__name__)
 _WATCH = [" nodded", " leaned"]
 
 
+def _chosen_target_quotas(
+    chosen_counts: Counter[str],
+    strength: float,
+) -> dict[str, int]:
+    """Return per-token occurrence caps for chosen-token regularisation."""
+    if not chosen_counts or strength <= 0:
+        return dict(chosen_counts)
+
+    # Prevent the largest outliers from dominating the distribution before
+    # applying the smoother median-threshold regularisation.
+    if len(chosen_counts) >= 10:
+        cap_value = sorted(chosen_counts.values(), reverse=True)[9]
+        capped = {
+            token: min(count, cap_value)
+            for token, count in chosen_counts.items()
+        }
+    else:
+        capped = dict(chosen_counts)
+
+    median = float(np.median(list(capped.values())))
+    return {
+        token: int(round(
+            count
+            if count <= median
+            else count * (median / count) ** strength
+        ))
+        for token, count in capped.items()
+    }
+
+
+def _trim_chosen_to_quotas(
+    rows: list[dict],
+    quotas: dict[str, int],
+    rng: np.random.Generator,
+) -> list[dict]:
+    """Uniformly retain chosen-token occurrences up to their global quotas."""
+    counts = Counter(
+        token
+        for row in rows
+        for token in (row.get("multi_chosen_decoded") or [])
+    )
+    if all(quotas.get(token, count) >= count for token, count in counts.items()):
+        return [dict(row) for row in rows]
+
+    # Randomise both row and slot traversal so quota allocation does not
+    # systematically favour early examples or a token's probability rank.
+    seen: Counter[str] = Counter()
+    trimmed: list[dict | None] = [None] * len(rows)
+    for row_idx_raw in rng.permutation(len(rows)):
+        row_idx = int(row_idx_raw)
+        row = rows[row_idx]
+        decoded = row.get("multi_chosen_decoded") or []
+        keep: set[int] = set()
+        for slot_idx_raw in rng.permutation(len(decoded)):
+            slot_idx = int(slot_idx_raw)
+            token = decoded[slot_idx]
+            if seen[token] < max(0, quotas.get(token, counts[token])):
+                keep.add(slot_idx)
+                seen[token] += 1
+
+        new_row = dict(row)
+        new_row["multi_chosen_decoded"] = [
+            token for slot_idx, token in enumerate(decoded) if slot_idx in keep
+        ]
+
+        raw = row.get("multi_chosen_raw")
+        if isinstance(raw, list) and len(raw) == len(decoded):
+            new_row["multi_chosen_raw"] = [
+                token for slot_idx, token in enumerate(raw) if slot_idx in keep
+            ]
+
+        trimmed[row_idx] = new_row
+
+    return [row for row in trimmed if row is not None]
+
+
 def load_ftpo_multi_dataset(
     path: Path,
     tokenizer,
@@ -108,24 +184,10 @@ def load_ftpo_multi_dataset(
 
     _log_top(chosen_cts_orig, "ORIGINAL CHOSEN TOKENS")
 
-    # Trim the peak: cap top tokens to match the 10th highest count
+    tgt_chosen = _chosen_target_quotas(
-    if len(chosen_cts_orig) >= 10:
+        chosen_cts_orig,
-        top_counts = sorted(chosen_cts_orig.values(), reverse=True)
+        chosen_reg_strength,
-        cap_value = top_counts[9]  # 10th highest count
+    )
-        chosen_cts_capped = Counter()
-        for tok, cnt in chosen_cts_orig.items():
-            chosen_cts_capped[tok] = min(cnt, cap_value)
-    else:
-        chosen_cts_capped = chosen_cts_orig.copy()
-
-    # Now calculate regularization on the capped distribution
-    med_chosen = float(np.median(list(chosen_cts_capped.values())))
-    w_chosen = {tok: 1.0 if c <= med_chosen
-            else (med_chosen / c) ** chosen_reg_strength
-            for tok, c in chosen_cts_capped.items()}
-
-    tgt_chosen = {tok: int(round(c * w_chosen.get(tok, 1.0)))
-                for tok, c in chosen_cts_capped.items()}
 
     # Log the target quotas
     quota_items = sorted(tgt_chosen.items(), key=lambda x: x[1], reverse=True)[:20]
@@ -137,7 +199,13 @@ def load_ftpo_multi_dataset(
         _WATCH[1], tgt_chosen.get(_WATCH[1], 0), chosen_cts_orig.get(_WATCH[1], 0),
     )
 
-    _log_top(Counter(r["rejected_decoded"] for r in rows), "POST-CHOSEN")
+    rows = _trim_chosen_to_quotas(rows, tgt_chosen, rng)
+    chosen_cts_trimmed = Counter(
+        token
+        for row in rows
+        for token in (row["multi_chosen_decoded"] or [])
+    )
+    _log_top(chosen_cts_trimmed, "POST-CHOSEN")
 
     # ────────────────────────────────────────────────────────────────
     # 3️⃣  Apply min_chosen_tokens row filter
@@ -206,6 +274,13 @@ def load_ftpo_multi_dataset(
     
     rows = selected
 
+    final_chosen_counts = Counter(
+        token
+        for row in rows
+        for token in (row["multi_chosen_decoded"] or [])
+    )
+    _log_top(final_chosen_counts, "FINAL CHOSEN TOKENS")
+
     # ── Dump the final row subset exactly as it was read (no tokenisation) ──
     if experiment_run_dir is not None:
         ts = datetime.now(timezone.utc).astimezone()\

utils/test_dataset_helpers.py

@@ -0,0 +1,154 @@
+import json
+import tempfile
+import unittest
+from collections import Counter
+from pathlib import Path
+
+import numpy as np
+
+from utils.dataset_helpers import (
+    _chosen_target_quotas,
+    _trim_chosen_to_quotas,
+    load_ftpo_multi_dataset,
+)
+
+
+def _counts(rows):
+    return Counter(
+        token
+        for row in rows
+        for token in row["multi_chosen_decoded"]
+    )
+
+
+class _Tokenized:
+    def __init__(self, input_ids):
+        self.input_ids = input_ids
+
+
+class _FakeTokenizer:
+    truncation_side = "right"
+
+    def __call__(self, text, **_kwargs):
+        def encode(value):
+            if value.startswith("context-"):
+                return [1, 2]
+            return [100 + sum(value.encode("utf-8"))]
+
+        if isinstance(text, list):
+            return _Tokenized([encode(value) for value in text])
+        return _Tokenized(encode(text))
+
+
+class ChosenRegularisationTests(unittest.TestCase):
+    def test_zero_strength_disables_all_trimming(self):
+        counts = Counter({f"token-{i}": 20 - i for i in range(12)})
+
+        self.assertEqual(_chosen_target_quotas(counts, 0), dict(counts))
+
+    def test_positive_strength_caps_and_regularises_outliers(self):
+        counts = Counter({f"token-{i}": 100 - 5 * i for i in range(12)})
+
+        quotas = _chosen_target_quotas(counts, 0.2)
+
+        tenth_highest = sorted(counts.values(), reverse=True)[9]
+        self.assertLess(quotas["token-0"], counts["token-0"])
+        self.assertLessEqual(quotas["token-0"], tenth_highest)
+        self.assertEqual(quotas["token-11"], counts["token-11"])
+
+    def test_trimming_enforces_quotas_and_keeps_raw_fields_aligned(self):
+        rows = [
+            {
+                "multi_chosen_decoded": [" common", " rare-a", " common"],
+                "multi_chosen_raw": ["raw-common-1", "raw-rare-a", "raw-common-2"],
+            },
+            {
+                "multi_chosen_decoded": [" common", " rare-b"],
+                "multi_chosen_raw": ["raw-common-3", "raw-rare-b"],
+            },
+        ]
+        quotas = {" common": 2, " rare-a": 1, " rare-b": 1}
+
+        trimmed = _trim_chosen_to_quotas(
+            rows, quotas, np.random.default_rng(3407)
+        )
+
+        self.assertEqual(_counts(trimmed), Counter(quotas))
+        for row in trimmed:
+            self.assertEqual(
+                len(row["multi_chosen_decoded"]),
+                len(row["multi_chosen_raw"]),
+            )
+        # The helper does not mutate the source rows.
+        self.assertEqual(rows[0]["multi_chosen_decoded"].count(" common"), 2)
+
+    def test_trimming_is_reproducible(self):
+        rows = [
+            {
+                "multi_chosen_decoded": [" common", f" unique-{i}"],
+                "multi_chosen_raw": [" common", f" unique-{i}"],
+            }
+            for i in range(20)
+        ]
+        quotas = {" common": 5, **{f" unique-{i}": 1 for i in range(20)}}
+
+        first = _trim_chosen_to_quotas(
+            rows, quotas, np.random.default_rng(123)
+        )
+        second = _trim_chosen_to_quotas(
+            rows, quotas, np.random.default_rng(123)
+        )
+
+        self.assertEqual(first, second)
+        self.assertEqual(_counts(first)[" common"], 5)
+
+    def test_minimum_filter_applies_after_trimming(self):
+        rows = [
+            {"multi_chosen_decoded": [" common", " keep"]},
+            {"multi_chosen_decoded": [" common", " other"]},
+        ]
+
+        trimmed = _trim_chosen_to_quotas(
+            rows,
+            {" common": 1, " keep": 1, " other": 1},
+            np.random.default_rng(7),
+        )
+        surviving = [
+            row for row in trimmed if len(row["multi_chosen_decoded"]) >= 2
+        ]
+
+        self.assertEqual(len(surviving), 1)
+
+    def test_loader_filters_rows_after_applying_chosen_quotas(self):
+        rows = [
+            {
+                "context_with_chat_template": f"context-{i}",
+                "rejected_decoded": " reject",
+                "multi_chosen_decoded": [" common", f" unique-{i}"],
+                "multi_chosen_raw": ["raw-common", f"raw-unique-{i}"],
+            }
+            for i in range(4)
+        ]
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            path = Path(tmp_dir) / "ftpo.jsonl"
+            with path.open("w", encoding="utf-8") as handle:
+                for row in rows:
+                    handle.write(json.dumps(row) + "\n")
+
+            dataset = load_ftpo_multi_dataset(
+                path,
+                _FakeTokenizer(),
+                chosen_reg_strength=1.0,
+                min_chosen_tokens=2,
+                num_proc=1,
+            )
+
+        # " common" is trimmed from four occurrences to one, so only its
+        # containing row still meets the two-chosen-token minimum.
+        self.assertEqual(len(dataset), 1)
+        self.assertEqual(len(dataset[0]["chosen_ids"]), 2)
+
+
+if __name__ == "__main__":
+    unittest.main()

utils/trainer_dataloaders.py

@@ -110,7 +110,7 @@ def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer:
             # balance *rejected* tokens
             rejected_reg_strength = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
             # balance *chosen* tokens
-            chosen_reg_strength   = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
+            chosen_reg_strength   = config.get("ftpo_sample_chosen_regularisation_strength", 0.0),
             # hard floor on |chosen|
             min_chosen_tokens     = config.get("ftpo_sample_min_chosen_tokens", 3),
             # overall training-set cap (used for per-token quotas too)
@@ -158,7 +158,7 @@ def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer:
             experiment_run_dir = experiment_run_dir,
             max_seq_len                  = max_seq_length,
             rejected_reg_strength        = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
-            chosen_reg_strength          = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
+            chosen_reg_strength          = config.get("ftpo_sample_chosen_regularisation_strength", 0.0),
             min_chosen_tokens            = config.get("ftpo_sample_min_chosen_tokens", 3),
             max_train_examples           = config.get("finetune_max_train_examples"),
         )
@@ -247,7 +247,7 @@ def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer:
             experiment_run_dir = experiment_run_dir,
             max_seq_len               = max_seq_length,
             rejected_reg_strength     = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
-            chosen_reg_strength       = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
+            chosen_reg_strength       = config.get("ftpo_sample_chosen_regularisation_strength", 0.0),
             min_chosen_tokens         = config.get("ftpo_sample_min_chosen_tokens", 3),
             max_train_examples        = config.get("finetune_max_train_examples"),
         )
@@ -299,7 +299,7 @@ def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer:
             experiment_run_dir = experiment_run_dir,
             max_seq_len               = max_seq_length,
             rejected_reg_strength     = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
-            chosen_reg_strength       = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
+            chosen_reg_strength       = config.get("ftpo_sample_chosen_regularisation_strength", 0.0),
             min_chosen_tokens         = config.get("ftpo_sample_min_chosen_tokens", 3),
             max_train_examples        = config.get("finetune_max_train_examples"),
         )