diff --git a/antipasto/config.py b/antipasto/config.py index d036ae6..2db9036 100644 --- a/antipasto/config.py +++ b/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 diff --git a/antipasto/train/train_adapter.py b/antipasto/train/train_adapter.py index a8cb3c5..bd776a7 100644 --- a/antipasto/train/train_adapter.py +++ b/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())