compute focus in raw space, improving numerical stability

This commit is contained in:
wassname
2026-01-15 11:53:28 +08:00
parent acaa60d556
commit b29addb24d
+8 -5
View File
@@ -372,16 +372,19 @@ def contrastive_steering_loss_with_ref(
# Concentration-aware: weight by subspace focus (||proj|| / ||full||)
# focus ∈ [0, 1] = fraction of delta energy in loss subspace
# Clamped to [0, 1] because numerically can exceed 1 due to Fisher weighting mismatch
# Compute in RAW space (undo Fisher) so focus is apples-to-apples with delta_full
cos_pos_ref_used = cos_pos_ref
cos_neg_ref_used = cos_neg_ref
focus_pos = None
focus_neg = None
if delta_pos_norm_full is not None and delta_neg_norm_full is not None:
proj_norm_pos = antisym_pos_agg.norm(dim=-1) # [b]
proj_norm_neg = antisym_neg_agg.norm(dim=-1) # [b]
focus_pos = (proj_norm_pos / delta_pos_norm_full.clamp(min=eps)).clamp(max=1.0)
focus_neg = (proj_norm_neg / delta_neg_norm_full.clamp(min=eps)).clamp(max=1.0)
# Undo Fisher normalization to get raw-space norms for concentration
proj_raw_pos = antisym_pos_agg * std_for_div.unsqueeze(0) # [b, r]
proj_raw_neg = antisym_neg_agg * std_for_div.unsqueeze(0) # [b, r]
proj_norm_pos = proj_raw_pos.norm(dim=-1) # [b]
proj_norm_neg = proj_raw_neg.norm(dim=-1) # [b]
focus_pos = proj_norm_pos / delta_pos_norm_full.clamp(min=eps)
focus_neg = proj_norm_neg / delta_neg_norm_full.clamp(min=eps)
if focus_softness > 0:
focus_pos = focus_pos.pow(1.0 - focus_softness)
focus_neg = focus_neg.pow(1.0 - focus_softness)