""" Diff-of-means concept vectors. Attempts R (recovery), F (faithfulness), X (control). For a binary target, the concept vector is the difference of class means: v = mean(z | y=1) - mean(z | y=0). - R: classify by projection onto v (threshold at the midpoint). - F: name v as the causal direction for y_causal. NOTE this is the interesting case — the diff-of-means picks up BOTH the causal axis (v_causal, which Rd reads) AND any *correlated* readable-not-causal contamination (v_ro). Whether steering v actually moves Rd is exactly what F1 tests causally. - X: to ADD concept m, estimate u_m as the diff-of-means between items with m active vs not (needs labeled probe items); to FLIP the label, push along the causal diff-of-means. The challenge verifies via Rd. We make X *causally honest*: we estimate the edit budget so Rd actually flips, using the technique's own probe data and the synth readout is NOT consulted (we calibrate budget from the projection statistics of the labeled train set). """ from __future__ import annotations import numpy as np from .base import Technique def _diff_means(Z, y): y = np.asarray(y) pos = Z[y == 1].mean(0) if (y == 1).any() else np.zeros(Z.shape[1]) neg = Z[y == 0].mean(0) if (y == 0).any() else np.zeros(Z.shape[1]) return pos - neg class DiffMeans(Technique): NAME = "diff_means" AXES = "RFX" # ---- R ---- def recover(self, encoder, train, test, target): Ztr = encoder.encode(train["items"]); ytr = np.asarray(train["labels"]) Zte = encoder.encode(test["items"]); yte = np.asarray(test["labels"]) if target["kind"] == "regression": # diff-of-means is for binary; regress on projection onto top corr dir # (correlate z dims with y, use as direction) — a linear surrogate. zc = Ztr - Ztr.mean(0) w = zc.T @ (ytr - ytr.mean()) w = w / (np.linalg.norm(w) + 1e-9) a = Ztr @ w # least-squares scale/offset A = np.vstack([a, np.ones_like(a)]).T coef, *_ = np.linalg.lstsq(A, ytr, rcond=None) pred = (Zte @ w) * coef[0] + coef[1] tol = target.get("tol", 0.5) acc = float(np.mean(np.abs(pred - yte) <= tol)) r2 = max(0.0, 1 - np.var(pred - yte) / (np.var(yte) + 1e-9)) return r2 > 0.2, {"pred": pred.tolist()}, float(np.clip(r2, 0, 1)) v = _diff_means(Ztr, ytr); v = v / (np.linalg.norm(v) + 1e-9) atr = Ztr @ v thr = 0.5 * (atr[ytr == 1].mean() + atr[ytr == 0].mean()) pred = (Zte @ v > thr).astype(int) cv = float(np.mean((atr > thr).astype(int) == ytr)) return cv > 0.6, {"pred": pred.tolist()}, float(np.clip(cv, 0, 1)) # ---- F: name the causal direction for y_causal ---- def faithfulness(self, encoder, train, test, target): Ztr = encoder.encode(train["items"]); ytr = np.asarray(train["labels"]) v = _diff_means(Ztr, ytr) v = v / (np.linalg.norm(v) + 1e-9) # confidence from how separable the classes are along v a = Ztr @ v sep = abs(a[ytr == 1].mean() - a[ytr == 0].mean()) / (a.std() + 1e-9) conf = float(np.clip(sep / 3.0, 0, 1)) return True, {"direction": v}, conf # ---- X (SONAR inject variant): diff-of-means concept injection ---- def _control_inject(self, encoder, train, test, target): """Inject a target concept by adding the diff-of-means direction between target-concept sentences and the source sentences (constructed probe).""" Zsrc_probe = encoder.encode(train["items"]) Ztgt = encoder.encode(target["target_sents"]) direction = Ztgt.mean(0) - Zsrc_probe.mean(0) direction = direction / (np.linalg.norm(direction) + 1e-9) items = test["items"] Z = encoder.encode(items) # scale to a fraction of the typical embedding norm (known budget) scale = float(np.median(np.linalg.norm(Z, axis=1))) * 0.6 * target.get("budget", 1.0) deltas = [(scale * direction).tolist() for _ in Z] return True, {"deltas": deltas}, 0.7 # ---- X: targeted edit (add concept + flip label) ---- def control(self, encoder, train, test, target): if target.get("kind") == "inject": return self._control_inject(encoder, train, test, target) # Build concept + causal direction estimates from labeled PROBE items the # technique constructs itself (it may query the encoder). synth_K = encoder._synth.K M = encoder._synth.M rng = np.random.default_rng(0) # estimate each concept direction u_m via diff-of-means on probe items. n_probe = 120 # build probe items: half with concept m, half without u_hat = {} for m in range(M): on, off = [], [] for _ in range(n_probe): L = 4 tok = list(rng.integers(0, synth_K, size=L)) pos = list(range(L)) base = {"tokens": tok, "positions": pos, "number": None, "entities": {}, "ro_g": 0.0, "causal_g": 0.0} on.append({**base, "concepts": [m]}) off.append({**base, "concepts": []}) Zon = encoder.encode(on); Zoff = encoder.encode(off) u_hat[m] = (Zon - Zoff).mean(0) # estimate causal flip direction via diff-of-means on causal_g sign. pos_items, neg_items = [], [] for _ in range(n_probe): L = 4 tok = list(rng.integers(0, synth_K, size=L)); p = list(range(L)) b = {"tokens": tok, "positions": p, "number": None, "concepts": [], "entities": {}, "ro_g": 0.0} pos_items.append({**b, "causal_g": 1.2}) neg_items.append({**b, "causal_g": -1.2}) Zp = encoder.encode(pos_items); Zn = encoder.encode(neg_items) v_causal_hat = (Zp - Zn).mean(0) v_causal_hat = v_causal_hat / (np.linalg.norm(v_causal_hat) + 1e-9) # build deltas per test item items = test["items"]; targets = test["targets"] Z = encoder.encode(items) deltas = [] for z, tgt in zip(Z, targets): m_add = tgt["add_concept"] d = np.zeros(z.shape[0]) # add the concept (unit u_m direction, budget to clear threshold ~1) um = u_hat[m_add]; um = um / (np.linalg.norm(um) + 1e-9) d = d + 1.2 * um # flip the label: push causal direction to opposite side, budget large proj = float(z @ v_causal_hat) d = d + (-np.sign(proj) if proj != 0 else 1.0) * (abs(proj) + 1.5) * v_causal_hat deltas.append(d) return True, {"deltas": [dd.tolist() for dd in deltas]}, 0.7