add & simplify warmup for mono, coh, and now focus

antipasto/config.py

@@ -309,29 +309,36 @@ class TrainingConfig:
     Dec 2024 variance analysis: mono_weight=100 gave 50% CV on flip%, mono_weight=10-20 is stable.
     """
     
-    mono_warmup_frac: float = 0.5
-    """Fraction of training before mono loss kicks in. During warmup, mono_weight=0.
+    mono_warmup_frac: float = -2
+    """Constraint warmup using -N syntax (binary: off during warmup, on after).
     
-    - -1.0: Follow LR warmup (use warmup_pct)
-    - 0.0: No warmup, mono active from start
-    - 0.5 (default): Mono kicks in after 50% of training
+    - Negative (-N): N × warmup_pct (e.g., -2 = 2× LR warmup = 20% at default)
+    - Zero: No warmup, active from start
+    - Positive: Explicit fraction (e.g., 0.3 = 30% of training)
     
-    At symmetric init, mono is satisfied (both endpoints at baseline) → zero gradient
-    OR mono fights projection before direction established → saddle trap.
-    Long warmup (50%) lets projection find direction first, then mono enforces.
+    Default -2: constraints activate at 2× LR warmup. Lets projection establish
+    direction before constraints kick in.
     """
 
-    coh_warmup_frac: float = -1
-    """Fraction of training before coherence loss kicks in. During warmup, coh disabled.
+    coh_warmup_frac: float = -2
+    """Constraint warmup using -N syntax (binary: off during warmup, on after).
     
-    - -1.0: Follow LR warmup (use warmup_pct)
-    - 0.0: No warmup, coh active from start
-    - 0.2: Coh kicks in after 20% of training
+    - Negative (-N): N × warmup_pct (e.g., -2 = 2× LR warmup = 20% at default)
+    - Zero: No warmup, active from start  
+    - Positive: Explicit fraction (e.g., 0.3 = 30% of training)
     
-    2026-01-05 sweep finding: coh=False outperforms coh=True by +5-14 F1.
-    Likely the same great-wall problem as mono: coherence fights projection early
-    before the adapter has found its steering direction. Warmup lets projection
-    establish antisymmetry first, then coherence provides soft guardrails.
+    Default -2: synchronized with mono_warmup_frac.
+    """
+
+    conc_warmup_frac: float = -2
+    """Concentration weighting warmup using -N syntax (binary: off during warmup, on after).
+    
+    - Negative (-N): N × warmup_pct (e.g., -2 = 2× LR warmup = 20% at default)
+    - Zero: No warmup, active 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.
     """
 
     orth_weight: float = 0

antipasto/train/train_adapter.py

@@ -114,6 +114,23 @@ def compute_batch_loss(
     if scale_adapter_fn is None:
         scale_adapter_fn = lambda coeff: ScaleAdapter(model, coeff=coeff)
     
+    # Constraint warmup: -N syntax → N × warmup_pct, 0 → no warmup, >0 → explicit frac
+    def resolve_warmup(frac: float) -> int:
+        if frac < 0:
+            effective = (-frac) * config.warmup_pct  # -2 → 2× LR warmup
+        else:
+            effective = frac
+        return int(effective * total_steps) if total_steps else 0
+    
+    mono_warmup_steps = resolve_warmup(config.mono_warmup_frac)
+    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
+    
     attention_mask = batch["attention_mask"]
     mask_cho = attention_mask[::2]
     mask_rej = attention_mask[1::2]
@@ -241,6 +258,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)
     loss_dict = contrastive_steering_loss_with_ref(
         s_ref_cho=s_ref_cho,
         s_ref_rej=s_ref_rej,
@@ -253,8 +271,8 @@ def compute_batch_loss(
         orth_weight=config.orth_weight,
         antisym_margin=config.antisym_margin,
         focus_softness=config.focus_softness,
-        delta_pos_norm_full=delta_pos_norm_full,
-        delta_neg_norm_full=delta_neg_norm_full,
+        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,
         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,
@@ -370,25 +388,6 @@ def compute_batch_loss(
         },
     }
     
-    # Compute effective mono_weight with warmup (follows LR warmup by default)
-    mono_warmup_frac = config.mono_warmup_frac if config.mono_warmup_frac >= 0 else config.warmup_pct
-    if config.mono and mono_warmup_frac > 0 and total_steps is None:
-        raise ValueError(
-            "compute_batch_loss: mono warmup requires total_steps, but got total_steps=None. "
-            "Pass total_steps through (train + val) so mono_weight warmup behaves as intended."
-        )
-    warmup_steps = int(mono_warmup_frac * total_steps) if total_steps else 0
-    if step < warmup_steps:
-        effective_mono_weight = 0.0
-    else:
-        effective_mono_weight = config.mono_weight
-    
-    # Compute effective coherence with warmup (same pattern as mono)
-    # 2026-01-05: coh=False >> coh=True by +5-14 F1. Warmup avoids great-wall problem.
-    coh_warmup_frac = config.coh_warmup_frac if config.coh_warmup_frac >= 0 else config.warmup_pct
-    coh_warmup_steps = int(coh_warmup_frac * total_steps) if total_steps else 0
-    enable_coherence_effective = config.coh and (step >= coh_warmup_steps)
-    
     total_loss, loss_components_dict, meta_pos, meta_neg, meta_shared = combine_dual_coef_losses(
         loss_pos=loss_results[+1.0],
         loss_neg=loss_results[-1.0],
@@ -1356,7 +1355,7 @@ def evaluate_model(
 
 
 @torch.no_grad()
-def generate_example_output(model, tokenizer, choice_ids, max_new_tokens=64, instructions=""):
+def generate_example_output(model, tokenizer, choice_ids, max_new_tokens=64, instructions="", skip_special_tokens=False):
     """Generate example outputs at different steering coefficients to show training progress.
 
     Args:
@@ -1404,8 +1403,8 @@ Action: Tell a white lie"""
     pmass = logp_choices.exp().sum(-1) 
 
     N = input_ids.shape[1]
-    q = tokenizer.decode(outputs.sequences[0][:N], skip_special_tokens=False)
-    a = tokenizer.decode(outputs.sequences[0][N:], skip_special_tokens=False)
+    q = tokenizer.decode(outputs.sequences[0][:N], skip_special_tokens=skip_special_tokens)
+    a = tokenizer.decode(outputs.sequences[0][N:], skip_special_tokens=skip_special_tokens)
     score = torch.mean(logratios).item()
 
     return (q, a, score, seq_nll[0].item(), pmass[0].item())