Nibmeister/auto-antislop
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix chosen-token regularization

Browse Source
This commit is contained in:
sam-paech
2026-07-23 15:42:54 -07:00
parent 6299030455
commit bc9e75fdec
10 changed files with 273 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
auto_antislop_config.yaml
View File

@@ -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 ─────────────────────────────────────────

3
configs/GLM-4-32B-0414.yaml
View File

@@ -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

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

@@ -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

3
configs/default.yaml
View File

@@ -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 ─────────────────────────────────────────

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

@@ -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

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

@@ -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

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

@@ -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

113
utils/dataset_helpers.py
View File

@@ -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
if len(chosen_cts_orig) >= 10:
top_counts = sorted(chosen_cts_orig.values(), reverse=True)
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()}
tgt_chosen = _chosen_target_quotas(
chosen_cts_orig,
chosen_reg_strength,
)
# 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()\

154
utils/test_dataset_helpers.py Normal file
View File

@@ -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()

8
utils/trainer_dataloaders.py
View File

@@ -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"),
)