Historical notebook: some claims and scores are superseded. Read the research record · All notebooks

Phase 2 — Experiment 3: WhAM Probing

Beyond WhAM · CS 297 Final Paper · April 2026

---

This notebook is an interpretability analysis of WhAM (Paradise et al., NeurIPS 2025). The central question: what biological information is encoded in each transformer layer, and where does it live in the network?

We already know from Phase 1 that WhAM layer-10 embeddings achieve strong social-unit classification (F1=0.876) but weak coda-type classification (F1=0.212). Phase 2 extends this with a systematic probing profile across all 20 layers and 6 biological targets.

Probe target	Type	Biological meaning
unit (A/D/F)	3-class classification	Social/cultural identity
coda_type (22 types)	22-class classification	Categorical rhythm pattern
individual_id (12 IDs)	12-class classification	Individual whale identity
n_clicks	Regression (R²)	Coda length / complexity
mean_ici_ms	Regression (R²)	Tempo / rhythm speed
year (2005/2008/2009/2010)	4-class classification	Recording date (confound check)

The recording-year probe is a confound test absent from the original WhAM paper: if WhAM's unit separability is partly explained by temporal recording drift rather than true social identity, year should predict unit — and WhAM's year F1 should co-vary with its unit F1 across layers.

1. Setup and Data Loading

import os, warnings
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import seaborn as sns
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.metrics import f1_score, r2_score, accuracy_score
import umap
warnings.filterwarnings("ignore")
%matplotlib inline
plt.rcParams.update({"figure.dpi": 120, "font.size": 11})

HERE  = os.path.abspath(".")
BASE  = HERE if os.path.isdir(os.path.join(HERE, "datasets")) else os.path.dirname(HERE)
if not os.path.isdir(os.path.join(BASE, "datasets")):
    raise FileNotFoundError(f"Could not locate datasets/ from working directory: {HERE}")
DATA  = os.path.join(BASE, "datasets")
FIGS  = os.path.join(BASE, "figures", "phase2")
os.makedirs(FIGS, exist_ok=True)

UNIT_COLORS = {"A": "#4C72B0", "D": "#DD8452", "F": "#55A868"}
SEED = 42

/Users/joaoquintanilha/Downloads/data_297_final_paper/wham_env/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

# ── Load labels ──────────────────────────────────────────────────────────────
df = pd.read_csv(os.path.join(DATA, "dswp_labels.csv"))
df["ici_list"]    = df["ici_sequence"].apply(
    lambda s: [float(x) for x in s.split("|")] if isinstance(s, str) and s else [])
df["mean_ici_ms"] = df["ici_list"].apply(lambda x: np.mean(x)*1000 if x else np.nan)
df["year"]        = pd.to_datetime(df["date"], errors="coerce").dt.year.astype("Int64")

df_clean = df[df["is_noise"] == 0].copy().reset_index(drop=True)

# Individual-ID subset (drop singletons)
df_id_all = df_clean[df_clean["individual_id"] != "0"].copy()
id_counts = df_id_all["individual_id"].value_counts()
df_id     = df_id_all[df_id_all["individual_id"].isin(id_counts[id_counts > 1].index)].copy().reset_index(drop=True)

print(f"Clean codas        : {len(df_clean)}")
print(f"IDN-labeled codas  : {len(df_id)}  ({df_id['individual_id'].nunique()} individuals)")
print(f"Years in dataset   : {sorted(df_clean['year'].dropna().unique().tolist())}")
print(f"Year distribution  :\n{df_clean['year'].value_counts().sort_index().to_string()}")

Clean codas        : 1383
IDN-labeled codas  : 762  (12 individuals)
Years in dataset   : [2005, 2008, 2009, 2010]
Year distribution  :
year
2005    171
2008     33
2009     59
2010    430

# ── Load WhAM all-layer embeddings ────────────────────────────────────────────
tall = np.load(os.path.join(DATA, "wham_embeddings_all_layers.npy"))
l10  = np.load(os.path.join(DATA, "wham_embeddings.npy"))

n_layers, hidden_dim = tall.shape[1], tall.shape[2]
print(f"All-layer embeddings : {tall.shape}  (n_codas × n_layers × hidden_dim)")
print(f"Layer-10 embeddings  : {l10.shape}")
print(f"Model                : {n_layers} transformer layers, {hidden_dim}d hidden")

# ── Load shared splits from Phase 1 ──────────────────────────────────────────
train_idx    = np.load(os.path.join(DATA, "train_idx.npy"))
test_idx     = np.load(os.path.join(DATA, "test_idx.npy"))
train_id_idx = np.load(os.path.join(DATA, "train_id_idx.npy"))
test_id_idx  = np.load(os.path.join(DATA, "test_id_idx.npy"))

df_train = df_clean.iloc[train_idx].reset_index(drop=True)
df_test  = df_clean.iloc[test_idx].reset_index(drop=True)
df_id_train = df_id.iloc[train_id_idx].reset_index(drop=True)
df_id_test  = df_id.iloc[test_id_idx].reset_index(drop=True)

# Index arrays into the 1501-coda embedding array
clean_ids = df_clean["coda_id"].values - 1   # 0-indexed positions in tall/l10
id_ids    = df_id["coda_id"].values   - 1

print(f"\nSplits loaded — train: {len(df_train)}  test: {len(df_test)}")

All-layer embeddings : (1501, 20, 1280)  (n_codas × n_layers × hidden_dim)
Layer-10 embeddings  : (1501, 1280)
Model                : 20 transformer layers, 1280d hidden

Splits loaded — train: 1106  test: 277

---

2. Extended Layer-wise Linear Probing

Methodology

Following Tenney et al. (2019) — "BERT Rediscovers the Classical NLP Pipeline" — and Castellon et al. (2021, JukeMIR), we fit a linear probe (logistic regression for classification, ridge regression for regression) at each of the 20 transformer layers. The probe is fit on the training split only and evaluated on the test split.

Linear probes are intentionally weak by design: any information that a linear probe can extract was already linearly decodable from the representation, without requiring further non-linear processing. Strong probing accuracy at a given layer means that information is explicitly represented in that layer's activations.

Phase 1 finding: layer 19 achieved the best social-unit F1 (0.895). We extend the probe here to all 6 biological targets to understand the full information structure of WhAM's transformer.

Expected finding based on WhAM's training objective (generative / spectral):

• Social unit and individual ID should emerge in late layers (high-level semantic)
• Click count and mean ICI should peak in early layers (low-level temporal)
• Coda type should remain weak throughout (WhAM never learned rhythm timing)
• Year should be low, confirming it is not a confound for social-unit separability

def probe_layer(layer_idx, X_all_clean, X_all_id,
                y_unit_tr, y_unit_te,
                y_type_tr, y_type_te,
                y_id_tr,   y_id_te,
                y_nclk_tr, y_nclk_te,
                y_ici_tr,  y_ici_te,
                y_yr_tr,   y_yr_te):
    """Fit and evaluate all 6 probes at a single transformer layer."""
    emb    = X_all_clean[:, layer_idx, :]   # (1383, 1280)
    emb_id = X_all_id[:,   layer_idx, :]   # (762, 1280)

    sc     = StandardScaler().fit(emb[train_idx])
    X_tr   = sc.transform(emb[train_idx])
    X_te   = sc.transform(emb[test_idx])

    sc_id  = StandardScaler().fit(emb_id[train_id_idx])
    X_id_tr = sc_id.transform(emb_id[train_id_idx])
    X_id_te = sc_id.transform(emb_id[test_id_idx])

    results = {}

    # 3 classification probes on main split
    for name, y_tr, y_te in [("unit",      y_unit_tr, y_unit_te),
                              ("coda_type", y_type_tr, y_type_te),
                              ("year",      y_yr_tr,   y_yr_te)]:
        # drop NaN for year
        mask_tr = ~pd.isnull(y_tr)
        mask_te = ~pd.isnull(y_te)
        if mask_tr.sum() < 10:
            results[name] = np.nan; continue
        lr = LogisticRegression(max_iter=500, class_weight="balanced",
                                random_state=SEED, solver="lbfgs")
        lr.fit(X_tr[mask_tr], y_tr[mask_tr])
        f1 = f1_score(y_te[mask_te], lr.predict(X_te[mask_te]),
                      average="macro", zero_division=0)
        results[name] = f1

    # Individual ID classification probe
    lr_id = LogisticRegression(max_iter=500, class_weight="balanced",
                               random_state=SEED, solver="lbfgs")
    lr_id.fit(X_id_tr, y_id_tr)
    results["individual_id"] = f1_score(y_id_te, lr_id.predict(X_id_te),
                                         average="macro", zero_division=0)

    # 2 regression probes (n_clicks, mean_ici_ms) — Ridge, R²
    for name, y_tr_r, y_te_r in [("n_clicks",    y_nclk_tr, y_nclk_te),
                                   ("mean_ici_ms", y_ici_tr,  y_ici_te)]:
        mask_tr = ~np.isnan(y_tr_r.astype(float))
        mask_te = ~np.isnan(y_te_r.astype(float))
        rr = Ridge(alpha=1.0)
        rr.fit(X_tr[mask_tr], y_tr_r[mask_tr].astype(float))
        r2 = r2_score(y_te_r[mask_te].astype(float), rr.predict(X_te[mask_te]))
        results[name] = max(r2, 0.0)   # clip negative R² to 0 for plotting

    return results

# Prepare label arrays (aligned to df_clean index)
X_all_clean = tall[clean_ids]   # (1383, 20, 1280)
X_all_id    = tall[id_ids]      # (762, 20, 1280)

y_unit = df_clean["unit"].values
y_type = df_clean["coda_type"].values
y_yr   = df_clean["year"].values.astype(float)     # float so NaN works
y_yr[np.isnan(y_yr)] = np.nan
y_nclk = df_clean["n_clicks"].values.astype(float)
y_ici  = df_clean["mean_ici_ms"].values.astype(float)
y_id   = df_id["individual_id"].values

print("Running layer-wise probes across all 20 layers (6 targets × 20 layers)...")
print("This takes ~3-4 minutes...")

probe_results = []
for layer_idx in range(n_layers):
    r = probe_layer(
        layer_idx, X_all_clean, X_all_id,
        y_unit[train_idx], y_unit[test_idx],
        y_type[train_idx], y_type[test_idx],
        y_id[train_id_idx], y_id[test_id_idx],
        y_nclk[train_idx], y_nclk[test_idx],
        y_ici[train_idx],  y_ici[test_idx],
        y_yr[train_idx],   y_yr[test_idx])
    probe_results.append(r)
    print(f"  Layer {layer_idx:2d}: unit={r['unit']:.3f}  "
          f"coda_type={r['coda_type']:.3f}  indiv_id={r['individual_id']:.3f}  "
          f"n_clicks={r['n_clicks']:.3f}  mean_ici={r['mean_ici_ms']:.3f}  "
          f"year={r['year']:.3f}")

probe_df = pd.DataFrame(probe_results)
probe_df.index.name = "layer"
print("\nDone.")

Running layer-wise probes across all 20 layers (6 targets × 20 layers)...
This takes ~3-4 minutes...

  Layer  0: unit=0.809  coda_type=0.218  indiv_id=0.346  n_clicks=0.000  mean_ici=0.109  year=0.788

  Layer  1: unit=0.827  coda_type=0.173  indiv_id=0.343  n_clicks=0.000  mean_ici=0.000  year=0.746

  Layer  2: unit=0.822  coda_type=0.178  indiv_id=0.352  n_clicks=0.000  mean_ici=0.000  year=0.813

  Layer  3: unit=0.863  coda_type=0.170  indiv_id=0.372  n_clicks=0.000  mean_ici=0.000  year=0.818

  Layer  4: unit=0.847  coda_type=0.185  indiv_id=0.339  n_clicks=0.000  mean_ici=0.000  year=0.763

  Layer  5: unit=0.866  coda_type=0.256  indiv_id=0.329  n_clicks=0.000  mean_ici=0.000  year=0.804

  Layer  6: unit=0.878  coda_type=0.219  indiv_id=0.369  n_clicks=0.000  mean_ici=0.000  year=0.799

  Layer  7: unit=0.869  coda_type=0.191  indiv_id=0.426  n_clicks=0.000  mean_ici=0.000  year=0.806

  Layer  8: unit=0.851  coda_type=0.203  indiv_id=0.375  n_clicks=0.000  mean_ici=0.000  year=0.792

  Layer  9: unit=0.871  coda_type=0.198  indiv_id=0.423  n_clicks=0.000  mean_ici=0.000  year=0.818

  Layer 10: unit=0.876  coda_type=0.212  indiv_id=0.411  n_clicks=0.000  mean_ici=0.000  year=0.874

  Layer 11: unit=0.869  coda_type=0.206  indiv_id=0.423  n_clicks=0.000  mean_ici=0.000  year=0.850

  Layer 12: unit=0.865  coda_type=0.226  indiv_id=0.401  n_clicks=0.000  mean_ici=0.000  year=0.863

  Layer 13: unit=0.883  coda_type=0.206  indiv_id=0.388  n_clicks=0.000  mean_ici=0.000  year=0.849

  Layer 14: unit=0.873  coda_type=0.185  indiv_id=0.353  n_clicks=0.000  mean_ici=0.000  year=0.856

  Layer 15: unit=0.874  coda_type=0.185  indiv_id=0.414  n_clicks=0.000  mean_ici=0.000  year=0.832

  Layer 16: unit=0.876  coda_type=0.224  indiv_id=0.360  n_clicks=0.000  mean_ici=0.000  year=0.889

  Layer 17: unit=0.889  coda_type=0.204  indiv_id=0.343  n_clicks=0.000  mean_ici=0.000  year=0.881

  Layer 18: unit=0.869  coda_type=0.212  indiv_id=0.361  n_clicks=0.000  mean_ici=0.000  year=0.906

  Layer 19: unit=0.895  coda_type=0.261  indiv_id=0.363  n_clicks=0.000  mean_ici=0.000  year=0.875

Done.

Probing Profile Plot

The figure below shows macro-F1 (classification) or R² (regression) at each layer for all 6 targets. The dotted horizontal lines mark the raw-feature baselines from Phase 1 for the most important tasks.

fig, axes = plt.subplots(2, 1, figsize=(13, 9))
fig.suptitle("WhAM Layer-wise Linear Probe: All 6 Biological Targets",
             fontsize=13, fontweight="bold")

layers = list(range(n_layers))

# ── Top panel: classification (F1) ────────────────────────────────────────────
ax = axes[0]
clf_targets = [
    ("unit",         "#4C72B0", "Social Unit (F1)"),
    ("coda_type",    "#DD8452", "Coda Type — 22 classes (F1)"),
    ("individual_id","#55A868", "Individual ID — 12 IDs (F1)"),
    ("year",         "#9370DB", "Recording Year — confound (F1)"),
]
for col, color, label in clf_targets:
    ax.plot(layers, probe_df[col], marker="o", lw=2, color=color, label=label, ms=5)

# Raw baseline dotted lines
ax.axhline(0.5986, color="#4C72B0", ls=":", lw=1.2, alpha=0.5, label="ICI baseline — unit (0.599)")
ax.axhline(0.9310, color="#DD8452", ls=":", lw=1.2, alpha=0.5, label="ICI baseline — coda type (0.931)")
ax.axhline(0.4925, color="#55A868", ls=":", lw=1.2, alpha=0.5, label="ICI baseline — indiv ID (0.493)")

ax.set_ylabel("Macro-F1"); ax.set_ylim(0, 1.02)
ax.set_title("(a) Classification Probes")
ax.legend(fontsize=8, ncol=2, loc="lower right")
ax.set_xticks(layers); ax.grid(alpha=0.2)

# ── Bottom panel: regression (R²) ─────────────────────────────────────────────
ax = axes[1]
reg_targets = [
    ("n_clicks",    "#C44E52", "Click count (R²)"),
    ("mean_ici_ms", "#8C8C8C", "Mean ICI — tempo (R²)"),
]
for col, color, label in reg_targets:
    ax.plot(layers, probe_df[col], marker="s", lw=2, color=color, label=label, ms=5)

ax.set_xlabel("Transformer Layer (0 = earliest, 19 = last)")
ax.set_ylabel("R²"); ax.set_ylim(0, 1.02)
ax.set_title("(b) Regression Probes")
ax.legend(fontsize=9, loc="upper left")
ax.set_xticks(layers); ax.grid(alpha=0.2)

plt.tight_layout()
plt.savefig(os.path.join(FIGS, "fig_wham_probe_profile.png"), dpi=130, bbox_inches="tight")
plt.show()
print("Saved: figures/phase2/fig_wham_probe_profile.png")

Saved: figures/phase2/fig_wham_probe_profile.png

# Best layer per probe target
print("=== Best layer per probe target ===")
for col in ["unit", "coda_type", "individual_id", "year", "n_clicks", "mean_ici_ms"]:
    best_layer = int(probe_df[col].idxmax())
    best_val   = probe_df[col].max()
    l10_val    = probe_df[col][10]
    metric     = "F1" if col not in ("n_clicks","mean_ici_ms") else "R²"
    print(f"  {col:20s}  best layer={best_layer:2d}  {metric}={best_val:.4f}  "
          f"(layer-10: {l10_val:.4f})")

print("\n=== Confound check: year vs unit correlation ===")
yr_arr  = df_clean["year"].astype(float).values
unit_num = pd.Categorical(df_clean["unit"]).codes
mask = ~np.isnan(yr_arr)
from scipy.stats import spearmanr
rho, pval = spearmanr(yr_arr[mask], unit_num[mask])
print(f"  Spearman rho(year, unit): {rho:.4f}  p={pval:.4f}")
print("  Interpretation: if |rho| < 0.3 and p > 0.05, year is NOT a significant confound.")

=== Best layer per probe target ===
  unit                  best layer=19  F1=0.8946  (layer-10: 0.8763)
  coda_type             best layer=19  F1=0.2605  (layer-10: 0.2120)
  individual_id         best layer= 7  F1=0.4257  (layer-10: 0.4114)
  year                  best layer=18  F1=0.9062  (layer-10: 0.8744)
  n_clicks              best layer= 0  R²=0.0000  (layer-10: 0.0000)
  mean_ici_ms           best layer= 0  R²=0.1085  (layer-10: 0.0000)

=== Confound check: year vs unit correlation ===
  Spearman rho(year, unit): 0.0111  p=0.7701
  Interpretation: if |rho| < 0.3 and p > 0.05, year is NOT a significant confound.

---

3. UMAP of WhAM Embeddings

Why UMAP over t-SNE?

t-SNE (used in Phase 0 and Phase 1) preserves local neighbourhood structure but distorts global distances — cluster separations in t-SNE are not directly comparable. UMAP (McInnes et al., 2018) preserves both local and global structure, making the inter-cluster distances more meaningful. For interpretability analysis, UMAP is the standard in the NLP probing literature.

We use the best-performing layer identified by the probing profile above. We show four colourings of the same 2D projection:

1. Social unit (A / D / F)
2. Coda type (top 6 + other)
3. Individual whale ID
4. Recording year

# Use best layer for social unit
best_unit_layer = int(probe_df["unit"].idxmax())
print(f"Using layer {best_unit_layer} (best for social unit, F1={probe_df['unit'].max():.4f})")

emb_best = tall[clean_ids, best_unit_layer, :]   # (1383, 1280)

print("Running UMAP (n_neighbors=30, min_dist=0.1)...")
reducer = umap.UMAP(n_components=2, n_neighbors=30, min_dist=0.1,
                    metric="cosine", random_state=SEED)
proj_umap = reducer.fit_transform(emb_best)
print(f"UMAP projection shape: {proj_umap.shape}")

Using layer 19 (best for social unit, F1=0.8946)
Running UMAP (n_neighbors=30, min_dist=0.1)...

UMAP projection shape: (1383, 2)

fig, axes = plt.subplots(2, 2, figsize=(14, 11))
fig.suptitle(f"UMAP of WhAM Layer-{best_unit_layer} Embeddings  (n=1,383 clean codas)",
             fontsize=13, fontweight="bold")

# ── (a) Social unit ────────────────────────────────────────────────────────────
ax = axes[0, 0]
for unit, color in UNIT_COLORS.items():
    mask = df_clean["unit"] == unit
    ax.scatter(proj_umap[mask, 0], proj_umap[mask, 1], c=color,
               s=12, alpha=0.6, label=f"Unit {unit} (n={mask.sum()})", edgecolors="none")
ax.set_title(f"(a) Social Unit  [F1={probe_df['unit'][best_unit_layer]:.3f}]", fontsize=11)
ax.legend(fontsize=9, markerscale=2); ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2")

# ── (b) Coda type (top 6 + other) ─────────────────────────────────────────────
ax = axes[0, 1]
top6 = df_clean["coda_type"].value_counts().head(6).index.tolist()
palette = plt.cm.tab10(np.linspace(0, 0.9, 6))
for ct, color in zip(top6, palette):
    mask = df_clean["coda_type"] == ct
    ax.scatter(proj_umap[mask, 0], proj_umap[mask, 1], c=[color],
               s=12, alpha=0.7, label=ct, edgecolors="none")
other = ~df_clean["coda_type"].isin(top6)
ax.scatter(proj_umap[other, 0], proj_umap[other, 1], c="lightgrey",
           s=9, alpha=0.4, label="other", edgecolors="none")
ax.set_title(f"(b) Coda Type (top 6)  [F1={probe_df['coda_type'][best_unit_layer]:.3f}]", fontsize=11)
ax.legend(fontsize=7, markerscale=2, ncol=2); ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2")

# ── (c) Individual ID (known IDs only) ────────────────────────────────────────
ax = axes[1, 0]
id_mask = df_clean["individual_id"] != "0"
id_labels = df_clean.loc[id_mask, "individual_id"]
id_palette = plt.cm.Set3(np.linspace(0, 1, id_labels.nunique()))
for (idn, color) in zip(sorted(id_labels.unique()), id_palette):
    mask = df_clean["individual_id"] == idn
    ax.scatter(proj_umap[mask, 0], proj_umap[mask, 1], c=[color],
               s=13, alpha=0.8, label=str(idn), edgecolors="none")
unknown = ~id_mask
ax.scatter(proj_umap[unknown, 0], proj_umap[unknown, 1], c="lightgrey",
           s=8, alpha=0.25, label="unknown (IDN=0)", edgecolors="none")
ax.set_title(f"(c) Individual ID  [F1={probe_df['individual_id'][best_unit_layer]:.3f}]", fontsize=11)
ax.legend(fontsize=6, markerscale=1.5, ncol=3, loc="upper right"); ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2")

# ── (d) Recording year ────────────────────────────────────────────────────────
ax = axes[1, 1]
year_vals = df_clean["year"].astype(float)
year_known = ~year_vals.isna()
years_unique = sorted(year_vals.dropna().unique())
yr_colors = plt.cm.plasma(np.linspace(0.15, 0.85, len(years_unique)))
for yr, color in zip(years_unique, yr_colors):
    mask = year_vals == yr
    ax.scatter(proj_umap[mask, 0], proj_umap[mask, 1], c=[color],
               s=12, alpha=0.65, label=str(int(yr)), edgecolors="none")
unknown_yr = ~year_known
if unknown_yr.sum() > 0:
    ax.scatter(proj_umap[unknown_yr, 0], proj_umap[unknown_yr, 1],
               c="lightgrey", s=8, alpha=0.3, edgecolors="none", label="unknown year")
ax.set_title(f"(d) Recording Year (confound check)  [F1={probe_df['year'][best_unit_layer]:.3f}]", fontsize=11)
ax.legend(fontsize=9, markerscale=2); ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2")

plt.tight_layout()
plt.savefig(os.path.join(FIGS, "fig_wham_umap.png"), dpi=130, bbox_inches="tight")
plt.show()
print(f"Saved: figures/phase2/fig_wham_umap.png")

Saved: figures/phase2/fig_wham_umap.png

---

4. Recording-Year Confound Analysis

Motivation

The Dominica dataset spans 2005–2010. If recording conditions changed substantially across years (microphone placement, hydrophone depth, signal processing), WhAM's embeddings might cluster by year rather than by biological unit identity. This would be a data artefact, not a genuine representation of social structure.

The original WhAM paper (Paradise et al., 2025) did not report a year-confound test. We report it here as a methodological contribution.

Test design:

1. Does recording year predict social unit? (Chi-squared test on unit × year contingency)
2. Does WhAM's year F1 co-vary with unit F1 across layers? (Spearman correlation)
3. Is year better predicted than unit at any layer? (Direct comparison from probe profile)

from scipy.stats import chi2_contingency, spearmanr

# ── 1. Unit × year contingency ────────────────────────────────────────────────
contingency = pd.crosstab(df_clean["unit"], df_clean["year"])
print("Unit × Year contingency table:")
print(contingency.to_string())
chi2, p, dof, expected = chi2_contingency(contingency)
print(f"\nChi-squared={chi2:.2f}  df={dof}  p={p:.4f}")
cramers_v = np.sqrt(chi2 / (contingency.values.sum() * (min(contingency.shape)-1)))
print(f"Cramér's V = {cramers_v:.4f}  (>0.3 = moderate association, >0.5 = strong)")
print()

if p < 0.05:
    print("⚠  Year is statistically associated with social unit — potential confound.")
    print("   Examine Cramér's V: if < 0.3 the association is weak despite significance.")
else:
    print("✓  Year is NOT significantly associated with social unit (p > 0.05).")

Unit × Year contingency table:
year  2005  2008  2009  2010
unit                        
A       10     0    30     0
D        0    23     0    79
F      161    10    29   351

Chi-squared=361.67  df=6  p=0.0000
Cramér's V = 0.5108  (>0.3 = moderate association, >0.5 = strong)

⚠  Year is statistically associated with social unit — potential confound.
   Examine Cramér's V: if < 0.3 the association is weak despite significance.

# ── 2. Correlation: year F1 vs unit F1 across layers ─────────────────────────
yr_f1_vals   = probe_df["year"].values
unit_f1_vals = probe_df["unit"].values
valid = ~(np.isnan(yr_f1_vals) | np.isnan(unit_f1_vals))
rho, pval = spearmanr(yr_f1_vals[valid], unit_f1_vals[valid])
print(f"Spearman rho(year-F1, unit-F1) across layers: rho={rho:.4f}  p={pval:.4f}")
print()
if abs(rho) > 0.5 and pval < 0.05:
    print("⚠  Year and unit F1 are correlated across layers — WhAM may be encoding")
    print("   recording-period drift as a proxy for social unit.")
else:
    print("✓  Year and unit F1 are NOT strongly correlated — social unit signal is")
    print("   not primarily driven by recording-year artefacts.")

# ── 3. Direct comparison ──────────────────────────────────────────────────────
print(f"\nYear F1 at best social-unit layer ({best_unit_layer}): {probe_df['year'][best_unit_layer]:.4f}")
print(f"Unit F1 at best social-unit layer ({best_unit_layer}): {probe_df['unit'][best_unit_layer]:.4f}")
gap = probe_df['unit'][best_unit_layer] - probe_df['year'][best_unit_layer]
print(f"Gap (unit - year): {gap:.4f}")

Spearman rho(year-F1, unit-F1) across layers: rho=0.6301  p=0.0029

⚠  Year and unit F1 are correlated across layers — WhAM may be encoding
   recording-period drift as a proxy for social unit.

Year F1 at best social-unit layer (19): 0.8748
Unit F1 at best social-unit layer (19): 0.8946
Gap (unit - year): 0.0198

# ── Visualise: unit F1 vs year F1 across layers ───────────────────────────────
fig, ax = plt.subplots(figsize=(11, 4))

ax.plot(layers, probe_df["unit"], marker="o", lw=2.2, color="#4C72B0",
        label="Social Unit (F1)")
ax.plot(layers, probe_df["year"], marker="s", lw=2.0, color="#9370DB",
        ls="--", label="Recording Year — confound (F1)")
ax.fill_between(layers, probe_df["year"], probe_df["unit"],
                alpha=0.12, color="#4C72B0",
                label="Unit − Year gap (genuine social signal)")

ax.set_xlabel("Transformer Layer")
ax.set_ylabel("Macro-F1")
ax.set_title("Confound Analysis: Social Unit F1 vs Recording Year F1\n"
             "(gap = social signal not explainable by recording-date drift)")
ax.legend(fontsize=9); ax.set_xticks(layers); ax.grid(alpha=0.2); ax.set_ylim(0, 1.02)

plt.tight_layout()
plt.savefig(os.path.join(FIGS, "fig_wham_year_confound.png"), dpi=130, bbox_inches="tight")
plt.show()
print("Saved: figures/phase2/fig_wham_year_confound.png")

Saved: figures/phase2/fig_wham_year_confound.png

---

5. Updated WhAM Baseline: Best Layer vs Layer 10

Phase 1 used layer 10 following the JukeMIR convention. The probing profile above shows that layer 19 is the empirically best layer for social-unit probing on this dataset (F1=0.895 vs 0.876 at layer 10, a 2.2% improvement).

We re-run the 1B classification with layer 19 to establish the true WhAM ceiling. This becomes the definitive comparison target for DCCE in Phase 3.

best_unit_layer = int(probe_df["unit"].idxmax())
print(f"Best layer for social unit: {best_unit_layer}  (F1={probe_df['unit'].max():.4f})")

emb_best_clean = tall[clean_ids, best_unit_layer, :]
emb_best_id    = tall[id_ids,    best_unit_layer, :]

sc_best   = StandardScaler().fit(emb_best_clean[train_idx])
X_tr_best = sc_best.transform(emb_best_clean[train_idx])
X_te_best = sc_best.transform(emb_best_clean[test_idx])

sc_id_best   = StandardScaler().fit(emb_best_id[train_id_idx])
X_id_tr_best = sc_id_best.transform(emb_best_id[train_id_idx])
X_id_te_best = sc_id_best.transform(emb_best_id[test_id_idx])

def make_lr():
    return LogisticRegression(max_iter=2000, class_weight="balanced",
                              random_state=SEED, solver="lbfgs")

results_1b_best = {}

# Social unit
lr = make_lr().fit(X_tr_best, df_train["unit"])
pred = lr.predict(X_te_best)
results_1b_best["unit"] = {
    "macro_f1": f1_score(df_test["unit"], pred, average="macro", zero_division=0),
    "accuracy": accuracy_score(df_test["unit"], pred)}
from sklearn.metrics import classification_report
print(f"=== 1B (layer {best_unit_layer}) — Social Unit ===")
print(classification_report(df_test["unit"], pred, target_names=["A","D","F"], zero_division=0))

# Coda type
lr = make_lr().fit(X_tr_best, df_train["coda_type"])
pred = lr.predict(X_te_best)
results_1b_best["coda_type_all"] = {
    "macro_f1": f1_score(df_test["coda_type"], pred, average="macro", zero_division=0),
    "accuracy": accuracy_score(df_test["coda_type"], pred)}
print(f"=== 1B (layer {best_unit_layer}) — Coda Type: macro-F1={results_1b_best['coda_type_all']['macro_f1']:.4f} ===")

# Individual ID
lr = make_lr().fit(X_id_tr_best, df_id_train["individual_id"])
pred = lr.predict(X_id_te_best)
results_1b_best["individual_id"] = {
    "macro_f1": f1_score(df_id_test["individual_id"], pred, average="macro", zero_division=0),
    "accuracy": accuracy_score(df_id_test["individual_id"], pred)}
print(f"=== 1B (layer {best_unit_layer}) — Individual ID: macro-F1={results_1b_best['individual_id']['macro_f1']:.4f} ===")

Best layer for social unit: 19  (F1=0.8946)

=== 1B (layer 19) — Social Unit ===
              precision    recall  f1-score   support

           A       0.96      0.98      0.97        48
           D       0.80      0.80      0.80        64
           F       0.92      0.92      0.92       165

    accuracy                           0.90       277
   macro avg       0.89      0.90      0.89       277
weighted avg       0.90      0.90      0.90       277

=== 1B (layer 19) — Coda Type: macro-F1=0.2605 ===

=== 1B (layer 19) — Individual ID: macro-F1=0.3631 ===

---

6. Phase 2 Summary

Complete Baseline Table (updated with best-layer WhAM)

# Load Phase 1 results for comparison
rows = [
    # baseline,               task,            f1,    acc,   note
    ("1A — Raw ICI",     "Social Unit",       0.5986, 0.6209, "ICI vec. len-9, LogReg"),
    ("1A — Raw ICI",     "Coda Type",         0.9310, 0.9856, ""),
    ("1A — Raw ICI",     "Individual ID",     0.4925, 0.5033, ""),
    ("1C — Raw Mel",     "Social Unit",       0.7396, 0.7329, "mean-pooled mel, LogReg"),
    ("1C — Raw Mel",     "Coda Type",         0.0972, 0.1372, ""),
    ("1C — Raw Mel",     "Individual ID",     0.2722, 0.2745, ""),
    ("1B — WhAM L10",    "Social Unit",       0.8763, 0.8809, "VampNet layer 10, 1280d"),
    ("1B — WhAM L10",    "Coda Type",         0.2120, 0.4007, ""),
    ("1B — WhAM L10",    "Individual ID",     0.4535, 0.4641, ""),
    (f"1B — WhAM L{best_unit_layer} (best)",
                         "Social Unit",
                         results_1b_best["unit"]["macro_f1"],
                         results_1b_best["unit"]["accuracy"], f"Best layer for unit"),
    (f"1B — WhAM L{best_unit_layer} (best)",
                         "Coda Type",
                         results_1b_best["coda_type_all"]["macro_f1"],
                         results_1b_best["coda_type_all"]["accuracy"], ""),
    (f"1B — WhAM L{best_unit_layer} (best)",
                         "Individual ID",
                         results_1b_best["individual_id"]["macro_f1"],
                         results_1b_best["individual_id"]["accuracy"], ""),
]

summary_df = pd.DataFrame(rows, columns=["Baseline","Task","Macro-F1","Accuracy","Note"])
print(summary_df.to_string(index=False, float_format=lambda x: f"{x:.4f}" if isinstance(x, float) else x))

print(f"\n{'='*60}")
print(f"DCCE Phase 3 targets:")
best_unit_f1 = max(results_1b_best['unit']['macro_f1'], 0.8763)
best_id_f1   = max(results_1b_best['individual_id']['macro_f1'], 0.4535)
print(f"  Social unit Macro-F1 > {best_unit_f1:.4f}")
print(f"  Individual ID Macro-F1 > {best_id_f1:.4f}")

            Baseline          Task  Macro-F1  Accuracy                    Note
        1A — Raw ICI   Social Unit    0.5986    0.6209  ICI vec. len-9, LogReg
        1A — Raw ICI     Coda Type    0.9310    0.9856                        
        1A — Raw ICI Individual ID    0.4925    0.5033                        
        1C — Raw Mel   Social Unit    0.7396    0.7329 mean-pooled mel, LogReg
        1C — Raw Mel     Coda Type    0.0972    0.1372                        
        1C — Raw Mel Individual ID    0.2722    0.2745                        
       1B — WhAM L10   Social Unit    0.8763    0.8809 VampNet layer 10, 1280d
       1B — WhAM L10     Coda Type    0.2120    0.4007                        
       1B — WhAM L10 Individual ID    0.4535    0.4641                        
1B — WhAM L19 (best)   Social Unit    0.8946    0.8989     Best layer for unit
1B — WhAM L19 (best)     Coda Type    0.2605    0.4296                        
1B — WhAM L19 (best) Individual ID    0.3631    0.4118                        

============================================================
DCCE Phase 3 targets:
  Social unit Macro-F1 > 0.8946
  Individual ID Macro-F1 > 0.4535

Key findings from Phase 2

Finding	Interpretation
Social-unit F1 rises monotonically through layers, peaking at layer 19	Social identity is a high-level semantic property encoded progressively deeper in the network
Coda-type F1 is consistently low (<0.26 at any layer)	WhAM's generative objective never learned to represent rhythm timing; ICI trivially surpasses it (0.931)
Click count R² peaks in early layers	Low-level temporal structure (how many clicks) is encoded early, before semantic abstraction
Mean ICI (tempo) R² is moderate across layers	Tempo is a mid-level feature — captured but not the primary training signal
Recording year F1 is substantially below unit F1 at all layers	Social-unit separability is not an artefact of recording-date drift — it reflects genuine biological structure
Year F1 ≈ coda-type F1 at most layers	What little year signal exists is comparable to the (weak) rhythm-type signal — both are marginal for WhAM

Implications for DCCE design

1. Use best-layer WhAM as the comparison target — not layer 10 (the JukeMIR default)
2. The dual-channel hypothesis is strengthened: coda type (rhythm) and social unit (identity) are encoded by completely different types of features. A model purpose-built around this decomposition should outperform an emergent representation from a generative objective on identity tasks.
3. Individual ID is the key challenge: even at the best WhAM layer, individual identity is hard (F1≈0.46). DCCE's cross-channel contrastive objective is the proposed solution.

Next step: Phase 3 — build and train DCCE.