From 06ac1b7e1a68982749540e85cacc9a0c6cc68a0e Mon Sep 17 00:00:00 2001
From: wassname <1103714+wassname@users.noreply.github.com>
Date: Fri, 16 Jan 2026 02:16:33 +0000
Subject: [PATCH 1/7] early stop after all warmups
---
antipasto/train/train_adapter.py | 65 +++++++++++++++++++++++---------
1 file changed, 47 insertions(+), 18 deletions(-)
diff --git a/antipasto/train/train_adapter.py b/antipasto/train/train_adapter.py
index b683186..0121f93 100644
--- a/antipasto/train/train_adapter.py
+++ b/antipasto/train/train_adapter.py
@@ -1030,23 +1030,43 @@ def train_epoch(
)
wandb_run.log(val_metrics, step=step)
- # Early stopping with min_delta (relative improvement threshold)
- # Early stopping (disabled when patience=0, e.g., with one-cycle scheduler)
- # Skip early stopping during warmup - LR is still ramping up
+ # Early stopping with min_delta (relative improvement threshold).
+ # Enable ONLY when coherence + monotonic + focus are ON, and only after
+ # all warmups are finished (LR warmup + coh/mono warmups).
warmup_steps = int(total_steps * config.warmup_pct) if total_steps else 0
- in_warmup = opt_step < warmup_steps
- # Detect first validation AFTER warmup (best_val_loss still at inf means we haven't started tracking)
- first_post_warmup = (not in_warmup) and (best_val_loss[0] == float("inf"))
-
- if in_warmup:
- logger.debug(f"Warmup: opt_step {opt_step}/{warmup_steps}, skipping early stopping check")
- elif first_post_warmup:
- # First validation after warmup - reset best_val_loss to current
+ mono_warmup_frac = config.mono_warmup_frac if config.mono_warmup_frac >= 0 else config.warmup_pct
+ coh_warmup_frac = config.coh_warmup_frac if config.coh_warmup_frac >= 0 else config.warmup_pct
+ mono_warmup_steps = int(total_steps * mono_warmup_frac) if total_steps else 0
+ coh_warmup_steps = int(total_steps * coh_warmup_frac) if total_steps else 0
+
+ # "focus" is considered enabled unless we explicitly ignore it.
+ focus_enabled = config.focus_softness < 1.0
+ early_stop_enabled = (
+ config.early_stop_patience > 0
+ and config.coh
+ and config.mono
+ and focus_enabled
+ and best_val_loss is not None
+ and patience_counter is not None
+ )
+ early_stop_ready_step = max(warmup_steps, mono_warmup_steps, coh_warmup_steps)
+ early_stop_ready = opt_step >= early_stop_ready_step
+ first_post_ready = early_stop_enabled and early_stop_ready and (best_val_loss[0] == float("inf"))
+
+ if early_stop_enabled and not early_stop_ready:
+ logger.debug(
+ f"Early stop gated: opt_step {opt_step}/{early_stop_ready_step} "
+ f"(warmup={warmup_steps}, mono_warmup={mono_warmup_steps}, coh_warmup={coh_warmup_steps})"
+ )
+ elif first_post_ready:
best_val_loss[0] = val_loss
patience_counter[0] = 0
- logger.info(f"Warmup complete at opt_step {opt_step}/{warmup_steps}. Starting early stopping with val_loss={val_loss:.4f}")
-
- if config.early_stop_patience > 0 and best_val_loss is not None and patience_counter is not None and not in_warmup and not first_post_warmup:
+ logger.info(
+ f"Early stopping enabled at opt_step {opt_step}/{early_stop_ready_step}. "
+ f"Starting tracking with val_loss={val_loss:.4f}"
+ )
+
+ if early_stop_enabled and early_stop_ready and not first_post_ready:
# Require relative improvement > min_delta to count as "better"
improved = val_loss < best_val_loss[0] * (1 - config.early_stop_min_delta)
@@ -1826,6 +1846,7 @@ def train_model(config: TrainingConfig):
num_workers=0 if config.quick else 8,
pin_memory=True,
persistent_workers=False if config.quick else True,
+ drop_last=True, # need full batch for fisher
)
val_dataloader = DataLoader(
val_dataset_pt,
@@ -1835,17 +1856,25 @@ def train_model(config: TrainingConfig):
num_workers=0 if config.quick else 8,
pin_memory=True,
persistent_workers=False if config.quick else True,
+ drop_last=True, # need full batch for fisher
)
total_steps = config.n_epochs * len(train_dataloader) // config.grad_accum_steps
opt = torch.optim.AdamW(
model.parameters(), lr=config.lr, weight_decay=config.wd
)
+ focus_enabled = config.focus_softness < 1.0
+ early_stop_enabled = (
+ config.early_stop_patience > 0 and config.coh and config.mono and focus_enabled
+ )
scheduler = torch.optim.lr_scheduler.OneCycleLR(
- opt, max_lr=config.lr, total_steps=total_steps, pct_start=config.warmup_pct,
-
- # Early stopping and one cycle are not usually combined, this setting effectively turns it into constant LR with warmup
- final_div_factor=1.0 if (config.early_stop_patience > 0) else 1e5
+ opt,
+ max_lr=config.lr,
+ total_steps=total_steps,
+ pct_start=config.warmup_pct,
+ # Early stopping and one-cycle are not usually combined; when early stopping
+ # is enabled, use effectively-constant LR with warmup.
+ final_div_factor=1.0 if early_stop_enabled else 1e5,
)
logger.info(f"Training: {config.n_epochs} epochs, {total_steps} steps")
From 4de4166d4ee355e0966ae68e7bd4a3eb36ad69f1 Mon Sep 17 00:00:00 2001
From: wassname <1103714+wassname@users.noreply.github.com>
Date: Fri, 16 Jan 2026 02:57:19 +0000
Subject: [PATCH 2/7] refine approx_intersection_bases: increase min_overlap to
0.5 and add diagnostic logging
---
antipasto/peft_utils/subspaces.py | 18 +++++++++++++-----
1 file changed, 13 insertions(+), 5 deletions(-)
diff --git a/antipasto/peft_utils/subspaces.py b/antipasto/peft_utils/subspaces.py
index 360a1e1..711685c 100644
--- a/antipasto/peft_utils/subspaces.py
+++ b/antipasto/peft_utils/subspaces.py
@@ -42,6 +42,7 @@ from einops import einsum
from jaxtyping import Float
from loguru import logger
from tqdm import tqdm
+import numpy as np
def get_hidden_size(model: nn.Module) -> int:
@@ -213,7 +214,7 @@ def approx_intersection_bases(
V_a: Float[Tensor, "d r_a"],
V_b: Float[Tensor, "d r_b"],
top_k: int = 256,
- min_overlap: float = 0.1,
+ min_overlap: float = 0.5,
) -> tuple[Float[Tensor, "d k"], Float[Tensor, "k"]]:
"""Intersection of two subspaces via principal angles.
@@ -229,9 +230,9 @@ def approx_intersection_bases(
Args:
V_a, V_b: Orthonormal bases [d_model, rank]
top_k: Maximum number of intersection directions to return
- min_overlap: Minimum cos(principal_angle) to include (default 0.1).
- S=1 means perfect overlap, S=0 means orthogonal.
- Directions with S < min_overlap are excluded as "not truly shared".
+ min_overlap: Minimum cos(principal_angle) to include (default 0.5 = 60°).
+ S=1 means perfect overlap (0°), S=0 means orthogonal (90°).
+ 0.1 (cos 84°) was too permissive; 0.5 (cos 60°) ensures actual alignment.
Returns:
V_shared: [d_model, k] orthonormal basis of shared directions
@@ -253,8 +254,15 @@ def approx_intersection_bases(
n_high_overlap = high_overlap_mask.sum().item()
k = min(top_k, n_high_overlap, S.shape[0])
+ # Log overlap quality for diagnostics
+ mean_overlap = S[:min(10, len(S))].mean().item()
+ logger.debug(f"Intersection overlap: max={S[0]:.3f}, top10_mean={mean_overlap:.3f}, n>{min_overlap}={n_high_overlap}/{len(S)}")
+
if k == 0:
- logger.warning(f"intersect_bases: no directions with overlap > {min_overlap} (max S={S[0]:.3f}). Returning top-1 anyway.")
+ logger.warning(
+ f"Intersection: no directions with overlap > {min_overlap:.2f} (cos {min_overlap:.2f} = {np.arccos(min_overlap)*180/np.pi:.0f}°). "
+ f"Max overlap={S[0]:.3f} ({np.arccos(S[0].item())*180/np.pi:.0f}°). Returning top-1 fallback."
+ )
k = 1
# Directions in original space:
From 08d54c70953d6f9c0a85492c49490041a8add099 Mon Sep 17 00:00:00 2001
From: wassname <1103714+wassname@users.noreply.github.com>
Date: Fri, 16 Jan 2026 06:09:28 +0000
Subject: [PATCH 3/7] bug fix
---
antipasto/train/daily_dilemas.py | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/antipasto/train/daily_dilemas.py b/antipasto/train/daily_dilemas.py
index 38c53ea..ee64421 100644
--- a/antipasto/train/daily_dilemas.py
+++ b/antipasto/train/daily_dilemas.py
@@ -1187,7 +1187,8 @@ def _compute_steering_f1_for_method(
common_idx = df_neg.index.intersection(df_0.index).intersection(df_pos.index)
if len(common_idx) == 0:
return {"steering_f1": np.nan, "net_correct": np.nan, "correct_w": np.nan,
- "wrong_w": np.nan, "arb_w": np.nan, "precision": np.nan,
+ "wrong_w": np.nan, "arb_w": np.nan, "correct_rate": np.nan,
+ "wrong_rate": np.nan, "arb_rate": np.nan, "precision": np.nan,
"recall": np.nan, "pmass_ratio": np.nan}
df_neg = df_neg.loc[common_idx]
From 43bca499426b9f0b4690f0b92cedf371fbb327b4 Mon Sep 17 00:00:00 2001
From: "wassname (Michael J Clark)" <1103714+wassname@users.noreply.github.com>
Date: Fri, 16 Jan 2026 15:17:21 +0800
Subject: [PATCH 4/7] Update README.md
---
README.md | 2 ++
1 file changed, 2 insertions(+)
diff --git a/README.md b/README.md
index 1a8c3d6..7c6d79b 100644
--- a/README.md
+++ b/README.md
@@ -14,6 +14,8 @@
**Why use it?** Prompting is fragile. System prompts get ignored. Jailbreaks work. AntiPaSTO trains directly on the model's internal representations, measuring and modifying what the model actually computes rather than what it says it will do. On the DailyDilemmas benchmark, it outperforms prompting on small models (≤4B) and complements arithmetic steering methods on larger ones.
+
+
From 226fd408b831b9ac9bc1b24e9ce464a07489219f Mon Sep 17 00:00:00 2001
From: "wassname (Michael J Clark)" <1103714+wassname@users.noreply.github.com>
Date: Fri, 16 Jan 2026 15:18:32 +0800
Subject: [PATCH 5/7] Replace image with new loss landscape visualization
Updated the image in the README to illustrate the loss landscape with a new visual representation.
---
README.md | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/README.md b/README.md
index 7c6d79b..95e8b3f 100644
--- a/README.md
+++ b/README.md
@@ -14,7 +14,7 @@
**Why use it?** Prompting is fragile. System prompts get ignored. Jailbreaks work. AntiPaSTO trains directly on the model's internal representations, measuring and modifying what the model actually computes rather than what it says it will do. On the DailyDilemmas benchmark, it outperforms prompting on small models (≤4B) and complements arithmetic steering methods on larger ones.
-
+
From 081ca634c2072ba94401e89dd1272a3383428644 Mon Sep 17 00:00:00 2001
From: "wassname (Michael J Clark)" <1103714+wassname@users.noreply.github.com>
Date: Fri, 16 Jan 2026 15:23:35 +0800
Subject: [PATCH 6/7] Update README.md by modifying image placement
Removed an image of the loss landscape and added it back later in the document.
---
README.md | 7 +++++--
1 file changed, 5 insertions(+), 2 deletions(-)
diff --git a/README.md b/README.md
index 95e8b3f..3bd8558 100644
--- a/README.md
+++ b/README.md
@@ -14,8 +14,6 @@
**Why use it?** Prompting is fragile. System prompts get ignored. Jailbreaks work. AntiPaSTO trains directly on the model's internal representations, measuring and modifying what the model actually computes rather than what it says it will do. On the DailyDilemmas benchmark, it outperforms prompting on small models (≤4B) and complements arithmetic steering methods on larger ones.
-
-
@@ -144,3 +142,8 @@ Built on the shoulders of other chefs:
url = {https://arxiv.org/abs/2601.07473}
}
```
+
+
+
+
+
From ddcd830a4b1a19c8de09ccf45634368a1b5ddb9b Mon Sep 17 00:00:00 2001
From: wassname <1103714+wassname@users.noreply.github.com>
Date: Fri, 16 Jan 2026 09:47:35 +0000
Subject: [PATCH 7/7] gradual warmup
---
antipasto/config.py | 141 ++++++++++++++++++++++++-------
antipasto/train/train_adapter.py | 35 +++++---
2 files changed, 133 insertions(+), 43 deletions(-)
diff --git a/antipasto/config.py b/antipasto/config.py
index 2db9036..b2e2ceb 100644
--- a/antipasto/config.py
+++ b/antipasto/config.py
@@ -40,7 +40,7 @@ class TrainingConfig:
seed: int = 42
"""Random seed for reproducibility (layer selection, dim selection, training dynamics)."""
- init_n_samples: int = 1000
+ init_n_samples: int = 2000
"""Number of samples for WANDA-style dimension selection and subspace computation.
Higher = more stable activation statistics, but slower init.
@@ -52,10 +52,10 @@ class TrainingConfig:
data_seed: int = 42
"""Fixed seed for data selection (which suffixes are used)."""
- model_name: str = "Qwen/Qwen3-4B-Instruct-2507"
+ model_name: str = "google/gemma-3-12b-it"
quantization_type: Literal["4bit", "8bit", "none"] = "none"
- n_modules: int = 512
+ n_modules: int = 256
"""Total number of layer×module combinations to select (by gradient importance).
Examples with n_modules=5:
@@ -64,7 +64,7 @@ class TrainingConfig:
The selection is sparse: each layer×module is ranked by gradient, top-k are selected.
Not Cartesian product - can have multiple modules at one layer, none at another.
- Default 42 (≈14 layers × 3 modules for typical 36-layer model).
+ Default 256.
"""
target_modules: List[str] = ["residual-writers"]
@@ -80,23 +80,25 @@ class TrainingConfig:
Explicit list: ["down_proj", "o_proj"] - only these module suffixes are candidates.
"""
- bs: int = 8
+ bs: int = 14
"""Batch size"""
- n_epochs: int = 20
+ n_epochs: int = 30
- lr: float = 5e-4
- """Learning rate. Sweep findings: 1e-4 too low (-18 F1), 3e-5 too low (-17 F1).
- For Cayley: 4e-4 to 6e-4. For LoRA/DoRA: ~10x lower (3e-5 to 6e-5).
+ lr: float = 0.002
+ """Learning rate.
+
+ Empirically, Cayley rotations tend to tolerate higher LR than LoRA/DoRA.
+ This repo's default matches a strong run on Qwen3-14B.
"""
- wd: float = 1e-8
+ wd: float = 1e-4
"""Weight decay"""
n_logs: int = 10
"""Log this many times per training"""
- val_every_n_samples: int = 512
+ val_every_n_samples: int = 1024
"""Validate every N training samples (independent of logging)."""
effective_bs: int = 32
@@ -108,7 +110,7 @@ class TrainingConfig:
val_split: float = 0.15
"""Fraction of data for validation"""
- early_stop_patience: int = 11
+ early_stop_patience: int = 16
"""Stop if val loss doesn't improve for N validation checks. 0 = disabled (recommended with one-cycle scheduler)."""
early_stop_min_delta: float = 0.00001
@@ -171,7 +173,7 @@ class TrainingConfig:
- taskdiff_x_write_x_notlogits: Task ∩ write ∩ (lm_head^⊥)
"""
- loss_subspace_rank: Optional[int] = 8
+ loss_subspace_rank: Optional[int] = 4
"""Rank (top-k) for loss subspace.
If None (default), select rank automatically via `loss_subspace_energy_frac`
@@ -187,12 +189,12 @@ class TrainingConfig:
k such that cumulative energy >= this fraction. 60% was used in MSRS paper
"""
- loss_layer_frac: float = 0.9
+ loss_layer_frac: float = 0.5
"""Depth fraction (0-1) at which to apply representation loss.
The loss is computed at a single layer: int(loss_layer_frac * num_hidden_layers).
- Default 0.8 (80% depth) is in the "planning zone" where Fisher ratio and
+ Default 0.5 (50% depth) is a simple mid-depth choice; prior sweeps often found a "planning zone" where Fisher ratio and
cross-sample consistency peak across tested architectures (Qwen, Gemma). Also suppurted by supported by e.g 2024-Gurnee-Universal-Neurons-in-GPT2-Language-Models.md
Rationale: gradient-based layer selection was circular (gradients flow FROM
@@ -220,7 +222,7 @@ class TrainingConfig:
dataset_name: str = "honest"
- max_samples: Optional[int] = 800
+ max_samples: Optional[int] = 3000
"""Max training samples (None = all)"""
n_last_tokens: int = 3
@@ -274,15 +276,15 @@ class TrainingConfig:
Projection loss naturally creates ordering; mono is a safety rail, not driver.
"""
- mono_margin: float = 0.4
+ mono_margin: float = 0.5
"""Monotonic threshold_frac: fraction of √H_ref for minimum separation.
Threshold = threshold_frac × √H_ref + threshold_floor.
- With H_ref=4 nats (typical), threshold_frac=0.4, floor=0.04: threshold ≈ 0.84 nats.
+ With H_ref=4 nats (typical), threshold_frac=0.5, floor=0.04: threshold ≈ 1.04 nats.
Sweep findings (2026-01-07, gemma1b):
| margin | F1 |
- | 0.4 | 23.6 | ← current default
+ | 0.5 | (default)
| 0.2 | 17.2 |
| 0.25 | 0.0 | (collapsed)
@@ -300,7 +302,7 @@ class TrainingConfig:
Prevents division issues and provides small stable deadzone.
"""
- mono_weight: float = 20.0
+ mono_weight: float = 30.0
"""Monotonicity loss scaling.
WARNING: Values ≥100 trap adapters in bad init - can't learn "no change" at c=0.
@@ -310,35 +312,36 @@ class TrainingConfig:
"""
mono_warmup_frac: float = -2
- """Constraint warmup using -N syntax (binary: off during warmup, on after).
+ """Constraint warmup using -N syntax (gradual linear ramp from 0 to full weight).
- Negative (-N): N × warmup_pct (e.g., -2 = 2× LR warmup = 20% at default)
- - Zero: No warmup, active from start
+ - Zero: No warmup, full weight from start
- Positive: Explicit fraction (e.g., 0.3 = 30% of training)
- Default -2: constraints activate at 2× LR warmup. Lets projection establish
- direction before constraints kick in.
+ Default -2: mono ramps up over 2× LR warmup. Lets projection establish
+ direction before constraints reach full strength.
"""
coh_warmup_frac: float = -2
- """Constraint warmup using -N syntax (binary: off during warmup, on after).
+ """Constraint warmup using -N syntax (gradual linear ramp from 0 to full weight).
- Negative (-N): N × warmup_pct (e.g., -2 = 2× LR warmup = 20% at default)
- - Zero: No warmup, active from start
+ - Zero: No warmup, full weight from start
- Positive: Explicit fraction (e.g., 0.3 = 30% of training)
- Default -2: synchronized with mono_warmup_frac.
+ Default -2: coh ramps up over 2× LR warmup. Prevents coh from fighting
+ projection loss early when deltas are large.
"""
conc_warmup_frac: float = -2
- """Concentration weighting warmup using -N syntax (binary: off during warmup, on after).
+ """Concentration (focus) weighting warmup using -N syntax (gradual linear ramp).
- Negative (-N): N × warmup_pct (e.g., -2 = 2× LR warmup = 20% at default)
- - Zero: No warmup, active from start
+ - Zero: No warmup, full weight from start
- Positive: Explicit fraction (e.g., 0.3 = 30% of training)
- During warmup, delta_*_norm_full=None disables subspace focus weighting.
- Default -2: synchronized with other constraints.
+ During ramp-up, focus weighting interpolates from 1.0 (no penalty for out-of-subspace)
+ to the configured focus_softness. Default -2: synchronized with other constraints.
"""
orth_weight: float = 0
@@ -647,12 +650,86 @@ default_configs = {
bs=64,
),
),
+ # Qwen/Qwen3-32B
+ "q32b-80gb": (
+ "Qwen 32B on 80GB GPU (maximum size)",
+ TrainingConfig(
+ model_name="Qwen/Qwen3-32B",
+ bs=4,
+ ),
+ ),
+ "q14b-80gb": (
+ "Qwen 14B on 80GB GPU (production quality)",
+ TrainingConfig(
+ model_name="Qwen/Qwen3-14B",
+ bs=12,
+ ),
+ ),
+
+ "q14b-goodrun": (
+ "Qwen 14B best run (fisher, coh+mono, r64)",
+ TrainingConfig(
+ model_name="Qwen/Qwen3-14B",
+ quantization_type="none",
+ n_modules=256,
+ target_modules=["residual-writers"],
+ bs=12,
+ n_epochs=20,
+ lr=0.002,
+ wd=1e-4,
+ n_logs=10,
+ val_every_n_samples=512,
+ effective_bs=32,
+ quick=False,
+ val_split=0.15,
+ early_stop_patience=14,
+ early_stop_min_delta=1e-5,
+ warmup_pct=0.1,
+ r=64,
+ rot_u=False,
+ rot_v=True,
+ dim_select_method="wanda_svd_l1_trip",
+ max_rotation_angle=pi / 4,
+ loss_subspace="taskdiff_x_suppressed_x_write",
+ loss_subspace_rank=4,
+ loss_subspace_energy_frac=0.6,
+ loss_layer_frac=0.5,
+ min_adapter_layer_frac=0.1,
+ dataset_name="honest",
+ max_samples=3000,
+ n_last_tokens=3,
+ coh=True,
+ coh_weight=10.0,
+ coh_thresh=0.9,
+ coh_barrier_mode="log1p_squared",
+ coh_lse_temperature=3.0,
+ mono=True,
+ mono_margin=0.5,
+ mono_threshold_floor=0.04,
+ mono_weight=30.0,
+ mono_warmup_frac=-2,
+ coh_warmup_frac=-2,
+ orth_weight=0,
+ antisym_margin=0.0,
+ fisher_var_floor_frac=0.1,
+ fisher_abs_std_floor=0.05,
+ fisher_detach_std=True,
+ eval_max_dilemmas=None,
+ eval_max_tokens=288,
+ use_wandb=True,
+ wandb_project="AntiPaSTO",
+ wandb_tags=None,
+ verbose=1,
+ PROMPT=PROMPT,
+ PERSONAS=PERSONAS,
+ ),
+ ),
# add gemma4b
"gemma12b-80gb": (
"Gemma 3 12B on 80GB GPU",
TrainingConfig(
model_name="google/gemma-3-12b-it",
- bs=4,
+ bs=14,
),
),
diff --git a/antipasto/train/train_adapter.py b/antipasto/train/train_adapter.py
index 0121f93..5a6dda3 100644
--- a/antipasto/train/train_adapter.py
+++ b/antipasto/train/train_adapter.py
@@ -126,10 +126,19 @@ def compute_batch_loss(
coh_warmup_steps = resolve_warmup(config.coh_warmup_frac)
conc_warmup_steps = resolve_warmup(config.conc_warmup_frac)
- # Binary switch: constraints off during warmup, on after
- effective_mono_weight = config.mono_weight if step >= mono_warmup_steps else 0.0
- enable_coherence_effective = config.coh and (step >= coh_warmup_steps)
- enable_concentration = step >= conc_warmup_steps
+ # Gradual linear ramp: 0 → full weight over warmup period
+ def ramp_weight(weight: float, warmup_steps: int) -> float:
+ if warmup_steps <= 0:
+ return weight
+ progress = min(1.0, step / warmup_steps)
+ return weight * progress
+
+ effective_mono_weight = ramp_weight(config.mono_weight, mono_warmup_steps)
+ effective_coh_weight = ramp_weight(config.coh_weight, coh_warmup_steps) if config.coh else 0.0
+ enable_coherence_effective = config.coh # Always enabled if config.coh=True, weight controls strength
+ # Focus warmup: interpolate focus_softness from 1.0 (disabled) to configured value
+ focus_warmup_progress = min(1.0, step / conc_warmup_steps) if conc_warmup_steps > 0 else 1.0
+ effective_focus_softness = 1.0 + (config.focus_softness - 1.0) * focus_warmup_progress
attention_mask = batch["attention_mask"]
mask_cho = attention_mask[::2]
@@ -258,7 +267,7 @@ def compute_batch_loss(
delta_neg_norm_full = delta_neg_agg.norm(dim=-1) # [b]
# Antisymmetric loss (Fisher + align + delta_full)
- # Disable concentration during warmup (delta_norm_full=None)
+ # Focus now uses gradual ramp via effective_focus_softness (1.0 = disabled, <1 = enabled)
loss_dict = contrastive_steering_loss_with_ref(
s_ref_cho=s_ref_cho,
s_ref_rej=s_ref_rej,
@@ -270,9 +279,9 @@ def compute_batch_loss(
last_n_tokens=config.n_last_tokens,
orth_weight=config.orth_weight,
antisym_margin=config.antisym_margin,
- focus_softness=config.focus_softness,
- delta_pos_norm_full=delta_pos_norm_full if enable_concentration else None,
- delta_neg_norm_full=delta_neg_norm_full if enable_concentration else None,
+ focus_softness=effective_focus_softness,
+ delta_pos_norm_full=delta_pos_norm_full,
+ delta_neg_norm_full=delta_neg_norm_full,
fisher_var_floor_frac=config.fisher_var_floor_frac,
fisher_abs_std_floor=config.fisher_abs_std_floor,
fisher_detach_std=config.fisher_detach_std,
@@ -324,7 +333,7 @@ def compute_batch_loss(
ref_label_logp=ref_coherence,
pi_label_logp=pi_coherence,
mask=mask_logp,
- scale=config.coh_weight,
+ scale=effective_coh_weight,
ref_logits=ref_logits,
pi_logits=pi_logits,
coh_thresh_frac=config.coh_thresh,
@@ -1034,8 +1043,12 @@ def train_epoch(
# Enable ONLY when coherence + monotonic + focus are ON, and only after
# all warmups are finished (LR warmup + coh/mono warmups).
warmup_steps = int(total_steps * config.warmup_pct) if total_steps else 0
- mono_warmup_frac = config.mono_warmup_frac if config.mono_warmup_frac >= 0 else config.warmup_pct
- coh_warmup_frac = config.coh_warmup_frac if config.coh_warmup_frac >= 0 else config.warmup_pct
+
+ # resolve_warmup: -N → N × warmup_pct, else explicit
+ def resolve_warmup_frac(frac: float) -> float:
+ return (-frac) * config.warmup_pct if frac < 0 else frac
+ mono_warmup_frac = resolve_warmup_frac(config.mono_warmup_frac)
+ coh_warmup_frac = resolve_warmup_frac(config.coh_warmup_frac)
mono_warmup_steps = int(total_steps * mono_warmup_frac) if total_steps else 0
coh_warmup_steps = int(total_steps * coh_warmup_frac) if total_steps else 0