wassname/Unsupervised-Elicitation
1
0
Fork 0
mirror of https://github.com/wassname/Unsupervised-Elicitation.git synced 2026-09-09 11:16:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

wip

Browse Source
This commit is contained in:
wassname
2025-10-18 17:09:00 +08:00
parent 84afdd3e7e
commit 0bcf2f50d8
2 changed files with 388 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/mjc_research_journal.md
+95
View File
@@ -383,3 +383,98 @@ So first lets plan, if you took at `nbs/05_backprop.py` you see it loads the LLM
Q: why is it so flow when it's not much data? It 10 epochs enougth when the final loss is nan?
## 2025-10-17 16:03:17
.venv ❯ uv run python nbs/05_backprop.py --lr 1 --opt lbfgs --epochs 1
Num labels: 1580, Num tuples: 31546
Sample prior: tensor([ 0.8857, -0.8857])
Sample tuple: {'context_uids': [43, 100, 135, 229, 234, 248], 'uid': 1344, 'llm_pred_diff': -4.999999865889549}
Labels shape: torch.Size([1580, 2])
Loss weights: {'mutual': 1.0, 'ranking': 0.5, 'prior': 0.1, 'direct': 0.5, 'entropy': 0.1}
Sample loss: 0.176172137260437
Running backprop experiment with config: Config(epochs=1, device='cuda', lr=1.0, weight_decay=0.0001, opt='lbfgs', test_mode=False, subsample_size=500)
--- Testing mutual = 1 ---
Optimizing: 100%|███████████████████████████████████████████████████████| 1/1 [00:53<00:00, 53.86s/it]
Losses for mutual: [5.44526710655191e-06]
Saved plot to outputs/backprop/loss_mutual.png
mutual: Acc 0.7854, ICM Corr 0.0000, Final Loss 0.0000, LLM Acc 0.8196
--- Testing ranking = 1 ---
Optimizing: 100%|█████████████████████████████████████████████████████| 1/1 [17:12<00:00, 1032.43s/it]
Losses for ranking: [0.6582178473472595]
Saved plot to outputs/backprop/loss_ranking.png
ranking: Acc 0.7785, ICM Corr 0.0000, Final Loss 0.6582, LLM Acc 0.8196
--- Testing prior = 1 ---
Optimizing: 100%|███████████████████████████████████████████████████████| 1/1 [01:12<00:00, 72.24s/it]
Losses for prior: [2.2664652377102357e-09]
Saved plot to outputs/backprop/loss_prior.png
prior: Acc 0.7854, ICM Corr 0.0000, Final Loss 0.0000, LLM Acc 0.8196
# 2025-10-18 08:15:58
kprop* ⇡
.venv ❯ uv run python nbs/05_backprop.py --lr 0.6 --epochs 30
Num labels: 1580, Num tuples: 31546
Sample prior: tensor([ 0.8857, -0.8857])
Sample tuple: {'context_uids': [43, 100, 135, 229, 234, 248], 'uid': 1344, 'llm_pred_diff': -4.999999865889549}
Labels shape: torch.Size([1580, 2])
Loss weights: {'mutual': 1.0, 'ranking': 0.5, 'prior': 0.1, 'direct': 0.5, 'entropy': 0.1}
Sample loss: 0.176172137260437
Running backprop experiment with config: Config(epochs=30, device='cuda', lr=0.6, weight_decay=0.0001, opt='adamw', test_mode=False, subsample_size=500)
--- Testing mutual = 1 ---
Optimizing: 100%|█████████████████████████████████████████████████████| 30/30 [36:53<00:00, 73.78s/it]
Losses for mutual: [5.44526710655191e-06, 2.9732079838140635e-06, 2.1852715690329205e-06, 9.025415579344553e-07, 6.105115062382538e-07, 7.952484111228841e-07, 9.283519943892315e-07, 8.908560289455636e-07, 7.671465596104099e-07, 6.811345087953669e-07, 6.456870664806047e-07, 6.183820460137213e-07, 5.555203301810252e-07, 4.691419519531337e-07, 3.902587479842623e-07, 3.493263136533642e-07, 3.4498489753787e-07, 3.607521250614809e-07, 3.7419397358462447e-07, 3.6516246382234385e-07, 3.349568089561217e-07, 2.991528162965551e-07, 2.700797381294251e-07, 2.546181292473193e-07, 2.4628150185890263e-07, 2.3550394701032928e-07, 2.1443301534418424e-07, 1.9170951759406307e-07, 1.7880944369608187e-07, 1.780347105295732e-07]
Saved plot to outputs/backprop/loss_mutual.png
mutual: Acc 0.4810, ICM Corr 0.0000, Final Loss 0.0000, LLM Acc 0.8196
--- Testing ranking = 1 ---
Optimizing: 100%|█████████████████████████████████████████████████████| 30/30 [34:28<00:00, 68.95s/it]
Losses for ranking: [0.6582178473472595, 0.3679029941558838, 0.24453502893447876, 0.20176899433135986, 0.18665046989917755, 0.18075571954250336, 0.17820386588573456, 0.17699174582958221, 0.17636868357658386, 0.17602629959583282, 0.17582713067531586, 0.17570547759532928, 0.17562797665596008, 0.17557671666145325, 0.1755417287349701, 0.17551715672016144, 0.17549940943717957, 0.17548635601997375, 0.175476536154747, 0.17546899616718292, 0.17546315491199493, 0.17545852065086365, 0.17545484006404877, 0.1754518300294876, 0.17544937133789062, 0.17544734477996826, 0.17544566094875336, 0.17544424533843994, 0.17544303834438324, 0.17544202506542206]
Saved plot to outputs/backprop/loss_ranking.png
ranking: Acc 0.7785, ICM Corr 0.0000, Final Loss 0.1754, LLM Acc 0.8196
--- Testing prior = 1 ---
Optimizing: 100%|█████████████████████████████████████████████████████| 30/30 [34:23<00:00, 68.78s/it]
Losses for prior: [2.2664652377102357e-09, 2.654478237218427e-07, 0.05624896660447121, 0.0022053676657378674, 0.03513054549694061, 0.0355745404958725, 0.012270934879779816, 0.005214049015194178, 0.015497569926083088, 0.021049635484814644, 0.013052523136138916, 0.00378213357180357, 0.004644844681024551, 0.010396131314337254, 0.010932975448668003, 0.006458289921283722, 0.003136074636131525, 0.003716099541634321, 0.005757021717727184, 0.0058898585848510265, 0.0038308403454720974, 0.002110244007781148, 0.0023815245367586613, 0.003376308362931013, 0.0032705573830753565, 0.002220130292698741, 0.0015421019634231925, 0.0017008127178996801, 0.0019950990099459887, 0.0017811475554481149]
Saved plot to outputs/backprop/loss_prior.png
prior: Acc 0.7899, ICM Corr 0.0000, Final Loss 0.0018, LLM Acc 0.8196
--- Testing direct = 1 ---
Optimizing: 100%|█████████████████████████████████████████████████████| 30/30 [33:30<00:00, 67.01s/it]
Losses for direct: [0.0775514468550682, 0.046122197061777115, 0.03277236223220825, 0.026906674727797508, 0.025021158158779144, 0.016025632619857788, 0.008202997036278248, 0.007869631983339787, 0.010838741436600685, 0.01052842941135168, 0.00679533276706934, 0.003636465175077319, 0.0032942702528089285, 0.004085233435034752, 0.0038681041914969683, 0.002642488107085228, 0.0017067184671759605, 0.0016403202898800373, 0.0018180719343945384, 0.0016242492711171508, 0.0011892315233126283, 0.0009423168376088142, 0.0009403791627846658, 0.0009179338812828064, 0.0007548658177256584, 0.000589850649703294, 0.0005384557880461216, 0.0005269444081932306, 0.0004550317826215178, 0.0003538132878020406]
Saved plot to outputs/backprop/loss_direct.png
direct: Acc 0.5032, ICM Corr 0.0000, Final Loss 0.0004, LLM Acc 0.8196
--- Testing entropy = 1 ---
Optimizing: 100%|█████████████████████████████████████████████████████| 30/30 [33:10<00:00, 66.33s/it]
Losses for entropy: [0.5828092694282532, 0.3595152199268341, 0.17632251977920532, 0.07604636251926422, 0.03231503441929817, 0.014426507987082005, 0.006915734149515629, 0.0035665021277964115, 0.0019684401340782642, 0.0011546822497621179, 0.0007149464217945933, 0.0004643737047445029, 0.00031468566157855093, 0.00022143956448417157, 0.00016114392201416194, 0.00012083828187314793, 9.308385051554069e-05, 7.346301572397351e-05, 5.924811557633802e-05, 4.873540092376061e-05, 4.080458529642783e-05, 3.472351818345487e-05, 2.997669616888743e-05, 2.622308602440171e-05, 2.3212656742543913e-05, 2.0782714273082092e-05, 1.8795792129822075e-05, 1.7143045624834485e-05, 1.576120121171698e-05, 1.4614166502724402e-05]
Saved plot to outputs/backprop/loss_entropy.png
entropy: Acc 0.7854, ICM Corr 0.0000, Final Loss 0.0000, LLM Acc 0.8196
--- Testing reward = 1 ---
Optimizing: 100%|█████████████████████████████████████████████████████| 30/30 [30:00<00:00, 60.03s/it]
Losses for reward: [-1.4354013204574585, -2.8484108448028564, -4.755040168762207, -8.751503944396973, -16.39519500732422, -32.38448715209961, -65.13714599609375, -139.09461975097656, -292.2000427246094, -624.2462158203125, -1316.2877197265625, -2716.56103515625, -5299.6962890625, -9705.220703125, -15951.375, -22737.826171875, -28519.1171875, -31865.12890625, -33273.56640625, -33850.38671875, -33984.828125, -34111.6953125, -34113.6484375, -34153.8828125, -34138.08203125, -34178.19140625, -34185.4140625, -34177.91796875, -34191.97265625, -34179.22265625]
Saved plot to outputs/backprop/loss_reward.png
reward: Acc 0.7829, ICM Corr 0.0000, Final Loss -34179.2227, LLM Acc 0.8196
`uv run python nbs/05_backprop.py --lr 0.4 --epochs 30 --test_mode`
If Acc >0.75, try mutual + prior combined
Share results - I can suggest further tweaks (e.g., fix direct loss weighting by consistency_key)
Results Summary (lr=0.1, 30 epochs, test_mode):
Individual losses:
mutual: 0.7095 (converged smoothly)
ranking: 0.7456 (stable)
prior: 0.7880 (best individual, close to baseline!)
direct: 0.5108 (broken - over-smoothing)
entropy: 0.7854 (stable)
reward: 0.7823 (stable but unbounded loss)
nbs/05_backprop.py
+293 -107
View File
@@ -7,7 +7,7 @@ import random
from pathlib import Path
from tqdm.auto import tqdm
from dataclasses import dataclass, field
from typing import List, Dict, Any, Tuple
from typing import List, Dict, Any, Tuple, Literal
import numpy as np
import pandas as pd
from loguru import logger
@@ -16,7 +16,19 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import simple_parsing # Add this import after other imports
from loguru import logger
import sys
try:
import wandb
WANDB_AVAILABLE = True
except ImportError:
WANDB_AVAILABLE = False
logger.warning("wandb not installed, skipping logging. Install with: uv pip install wandb")
# Configure loguru for JSON output without prefix
logger.remove() # Remove default sink
logger.add(sys.stdout, format="{message}")
# %% [code]
# Load TruthfulQA data with labels
@@ -100,15 +112,50 @@ for pred in preds: # Subsample
'llm_pred_diff': pred['score'],
})
# Precompute for speed: targets, llm_diffs, context_index_lists, group_index_lists
targets_list = []
llm_diffs_list = []
context_index_lists = []
valid_tuples = [] # Filtered tuples with valid context
for t in tuples:
target_uid = t['uid']
if target_uid not in uid_to_idx:
continue
target_idx = uid_to_idx[target_uid]
context_indices = [uid_to_idx[u] for u in t['context_uids'] if u in uid_to_idx]
if len(context_indices) > 0:
targets_list.append(target_idx)
llm_diffs_list.append(t['llm_pred_diff'])
context_index_lists.append(context_indices)
valid_tuples.append(t)
tuples = valid_tuples # Update to valid only
targets = torch.tensor(targets_list)
llm_diffs = torch.tensor(llm_diffs_list)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Join with original data for consistency_id (add to tuples or separate)
# For now, create consistency_groups: dict of lists of uids per consistency_id
consistency_groups = defaultdict(list)
for uid, row in df_data.iterrows():
consistency_groups[row['consistency_id']].append(uid)
print(f"Num labels: {num_labels}, Num tuples: {len(tuples)}")
print(f"Sample prior: {priors[0]}")
print(f"Sample tuple: {tuples[0] if tuples else 'None'}")
# Precompute group indices for direct loss
group_index_lists = []
for group_uids in consistency_groups.values():
group_indices = [uid_to_idx.get(u) for u in group_uids if u in uid_to_idx]
if len(group_indices) > 1:
group_index_lists.append(group_indices)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
context_index_tensors = [torch.tensor(lst, dtype=torch.long, device=device) for lst in context_index_lists]
group_index_tensors = [torch.tensor(lst, dtype=torch.long, device=device) for lst in group_index_lists]
logger.info(f"Num labels: {num_labels}, Num tuples: {len(tuples)}")
logger.info(f"Sample prior: {priors[0]}")
logger.info(f"Sample tuple: {tuples[0] if tuples else 'None'}")
# %% [code]
# Define learnables
@@ -124,12 +171,12 @@ loss_weights = {
'entropy': 0.1
}
print(f"Labels shape: {labels.shape}")
print(f"Loss weights: {loss_weights}")
logger.info(f"Labels shape: {labels.shape}")
logger.info(f"Loss weights: {loss_weights}")
# %% [code]
# Define modular loss function
def coherence_loss(labels, tuples, priors, consistency_groups, uid_to_idx, loss_weights, scale=1000.0, verbose=False):
def coherence_loss(labels, tuples, priors, consistency_groups, uid_to_idx, loss_weights, scale=1000.0, targets=None, llm_diffs=None, context_index_lists=None, group_index_lists=None, context_index_tensors=None, group_index_tensors=None, verbose=False, temperature=1.0):
"""
Custom loss for unsupervised label optimization.
@@ -143,9 +190,10 @@ def coherence_loss(labels, tuples, priors, consistency_groups, uid_to_idx, loss_
Weighted sum for flexibility; adjust loss_weights dict to tune.
"""
# FIXME learning weightson the loss can lead to rewards hacking, e.g. learn 1 for easy loss, and 0 for hard ones
device = labels.device
priors = priors.to(device)
soft_labels = F.softmax(labels, dim=1) # [num_labels, 2]
rank_loss_fn = nn.MarginRankingLoss(margin=1.0)
total_loss = 0.0
num_tuples = len(tuples)
terms = {}
@@ -153,11 +201,11 @@ def coherence_loss(labels, tuples, priors, consistency_groups, uid_to_idx, loss_
# No extra deps; use loop-based aggregation for small num_labels (~800)
num_nodes = soft_labels.shape[0]
weighted_context = torch.zeros_like(soft_labels) # [num_nodes, 2]
degrees = torch.zeros(num_nodes).to(soft_labels.device) # For normalization
degrees = torch.zeros(num_nodes, device=device) # For normalization
for t in tuples:
target_idx = uid_to_idx[t['uid']]
context_indices = torch.tensor([uid_to_idx[u] for u in t['context_uids']])
for t_idx in range(num_tuples):
target_idx = targets[t_idx]
context_indices = context_index_tensors[t_idx]
if len(context_indices) == 0:
continue
@@ -176,37 +224,36 @@ def coherence_loss(labels, tuples, priors, consistency_groups, uid_to_idx, loss_
mask = degrees > 0
weighted_context[mask] /= degrees[mask].unsqueeze(1)
# Updated Mutual: Use weighted_context for CE
mutual_loss = 0.0
for t in tuples:
target_idx = uid_to_idx[t['uid']]
if degrees[target_idx] == 0:
continue
context_agg = weighted_context[target_idx]
target_soft = soft_labels[target_idx]
# Updated Mutual: Use weighted_context for CE, vectorized where possible
# Filter valid targets (degrees >0)
valid_mask = degrees[targets] > 0
valid_targets = targets[valid_mask]
if len(valid_targets) == 0:
mutual_loss = torch.tensor(0.0, device=device)
else:
context_agg_per_valid = weighted_context[valid_targets] # [num_valid, 2]
target_soft_per_valid = soft_labels[valid_targets] # [num_valid, 2]
# Forward
mutual_loss += F.cross_entropy(context_agg.unsqueeze(0), target_soft.unsqueeze(0))
# Reverse
mutual_loss += F.cross_entropy(target_soft.unsqueeze(0), context_agg.unsqueeze(0))
terms['mutual'] = mutual_loss / max(num_tuples, 1) / 2
# Asymmetric CE: context aggregation predicts target label (one-way inference)
# This matches ICM's P(y_i | context) formulation - context should confidently predict target
# Temperature annealing: high temp early (explore), low temp late (exploit)
# Prevents premature collapse to uniform/degenerate solutions
context_agg_tempered = F.softmax(context_agg_per_valid.log() / (temperature + 1e-8), dim=1)
mutual_loss = -(target_soft_per_valid * (context_agg_tempered + 1e-12).log()).sum(dim=1).mean()
terms['mutual'] = mutual_loss / max(num_tuples, 1)
# Pairwise Ranking: Enforce ranking on learnable diff vs LLM pred diff
# Intuition: LLM logprobs are better for relative rankings than absolute probs; ensure learned prob ranking matches LLM's scaled logprob diff.
# Simple English: "The LLM ranked 'yes' higher than 'no' for this example—make sure your learned label agrees on which is stronger, with a safety margin to avoid ties. Like forcing your guesses to match the model's confidence order, not exact numbers."
ranking_loss = 0.0
for t in tuples:
target_idx = uid_to_idx[t['uid']]
llm_diff = t['llm_pred_diff'] / scale
score1 = soft_labels[target_idx, 1] # Prob 1
score2 = soft_labels[target_idx, 0] # Prob 0
target_rank = 1 if llm_diff > 0 else -1
if target_rank == -1:
ranking_loss += rank_loss_fn(score2, score1, torch.tensor(1.0))
else:
ranking_loss += rank_loss_fn(score1, score2, torch.tensor(1.0))
terms['ranking'] = ranking_loss / max(num_tuples, 1)
# Vectorized ranking
score1 = soft_labels[targets, 1] # [num_tuples, ]
score0 = soft_labels[targets, 0]
learned_diffs = score1 - score0
signs = (llm_diffs > 0).float() * 2 - 1 # +1 if llm prefers 1, -1 else
m = 1.0
violations = torch.clamp(m - signs * learned_diffs, min=0.0)
ranking_loss = violations.mean()
terms['ranking'] = ranking_loss
# Prior KL: Pull toward priors
# Intuition: Anchor learned labels to fixed LLM priors (averaged logprob diffs) via KL on logits to prevent drift from model's initial biases.
@@ -218,31 +265,27 @@ def coherence_loss(labels, tuples, priors, consistency_groups, uid_to_idx, loss_
# Intuition: Labels in the same group (e.g., related questions via consistency_id) should agree; penalize soft label variance for local stability.
# Simple English: "Related examples (like paraphrases) should have similar labels—don't let them vary wildly. Average their probabilities and punish if they're all over the place, ensuring the model doesn't contradict itself on similar stuff."
direct_loss = 0.0
num_groups = 0
for group_uids in consistency_groups.values():
group_indices = [uid_to_idx.get(u) for u in group_uids if u in uid_to_idx]
if len(group_indices) > 1:
group_soft = soft_labels[torch.tensor(group_indices)]
direct_loss += group_soft.var(dim=0).mean()
num_groups += 1
num_groups = len(group_index_lists)
for g_idx in range(num_groups):
group_indices = group_index_tensors[g_idx]
group_soft = soft_labels[group_indices]
direct_loss += group_soft.var(dim=0).mean()
terms['direct'] = direct_loss / max(num_groups, 1)
# Entropy: Penalize high entropy in targets (encourage confident labels)
# Intuition: Reward low-entropy (decisive) labels for targets, proxying downstream confidence from coherent propagation.
# Simple English: "Make labels decisive (mostly yes or no, not 50/50 unsure). For each target, calculate how 'spread out' its probability is and add a small penalty if it's too wishy-washy—pushes toward bold, consistent choices that build confidence across predictions."
entropy_loss = 0.0
for t in tuples:
target_idx = uid_to_idx[t['uid']]
target_soft = soft_labels[target_idx]
entropy = -(target_soft * torch.log(target_soft + 1e-8)).sum()
entropy_loss += entropy
terms['entropy'] = entropy_loss / max(num_tuples, 1)
# Vectorized entropy
target_soft_per_tuple = soft_labels[targets] # [num_tuples, 2]
entropies = -(target_soft_per_tuple * torch.log(target_soft_per_tuple + 1e-8)).sum(dim=1)
entropy_loss = entropies.mean()
terms['entropy'] = entropy_loss
# New: Reward context for good evidence
# New: Reward context for good evidence (still looped, but faster with precompute)
reward_loss = 0.0
for t in tuples:
target_idx = uid_to_idx[t['uid']]
context_indices = torch.tensor([uid_to_idx[u] for u in t['context_uids']])
for t_idx in range(num_tuples):
target_idx = targets[t_idx]
context_indices = context_index_tensors[t_idx]
if len(context_indices) == 0:
continue
@@ -259,51 +302,124 @@ def coherence_loss(labels, tuples, priors, consistency_groups, uid_to_idx, loss_
weighted_loss = sum(loss_weights.get(k, 0.0) * terms[k] for k in terms)
total_loss = weighted_loss / (sum(loss_weights.values()) + loss_weights.get('reward', 0.0) or 1e-5)
# For hacking: print terms
# For hacking: logger.info terms
if torch.is_grad_enabled() and verbose:
print(f"Loss terms: { {k: v.item() if hasattr(v, 'item') else v for k,v in terms.items()} }")
logger.info(f"Loss terms: { {k: v.item() if hasattr(v, 'item') else v for k,v in terms.items()} }")
assert torch.isfinite(total_loss), "Loss is NaN or Inf"
return total_loss
# Test with dummy
dummy_loss = coherence_loss(labels, tuples[:1], priors, consistency_groups, uid_to_idx, loss_weights)
print(f"Sample loss: {dummy_loss.item()}")
# Test with dummy (update call with precomputes)
# Note: for dummy, move to device manually if needed
device_test = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
labels_test = labels.to(device_test)
priors_test = priors.to(device_test)
targets_test = targets[:1].to(device_test)
llm_diffs_test = llm_diffs[:1].to(device_test)
context_index_lists_test = context_index_lists[:1]
group_index_lists_test = group_index_lists
context_index_tensors_test = [torch.tensor(lst, dtype=torch.long, device=device_test) for lst in context_index_lists_test]
group_index_tensors_test = [torch.tensor(lst, dtype=torch.long, device=device_test) for lst in group_index_lists_test]
dummy_loss = coherence_loss(labels_test, tuples[:1], priors_test, consistency_groups, uid_to_idx, loss_weights, targets=targets_test, llm_diffs=llm_diffs_test, context_index_lists=context_index_lists_test, group_index_lists=group_index_lists_test, context_index_tensors=context_index_tensors_test, group_index_tensors=group_index_tensors_test, verbose=False)
logger.info(f"Sample loss: {dummy_loss.item()}")
# %% [code]
@dataclass
class Config:
epochs: int = 10
device: str = "cuda"
lr: float = 0.4
weight_decay: float = 1e-4
opt: Literal["adamw", "lbfgs"] = "adamw"
test: bool = False
subsample_size: int = 500 # just for test mode
wandb: bool = False # Enable wandb logging
def run_backprop_experiment(labels, tuples, priors, consistency_groups, uid_to_idx, loss_weights, df_data, df_preds, unique_uids,
prior_logits,
device='cuda',
):
targets, llm_diffs, context_index_lists, group_index_lists, context_index_tensors, group_index_tensors,
config: Config,
device=None):
"""
Run backprop experiment with given weights, return metrics.
"""
if device is None:
device = torch.device(config.device if torch.cuda.is_available() else 'cpu')
if config.test:
epochs = config.epochs
num_keep = min(config.subsample_size, len(tuples))
tuples = tuples[:num_keep]
targets = targets[:num_keep]
llm_diffs = llm_diffs[:num_keep]
context_index_lists = context_index_lists[:num_keep]
context_index_tensors = context_index_tensors[:num_keep]
group_index_tensors = group_index_tensors # unchanged
logger.info(f"Test mode: subsampled to {num_keep} tuples, {epochs} epochs")
else:
epochs = config.epochs
# Copy labels
current_labels = labels.clone().detach().to(device).requires_grad_(True)
priors = priors.to(device)
priors_dev = priors.to(device)
targets_dev = targets.to(device)
llm_diffs_dev = llm_diffs.to(device)
# Shared loss computation function
def get_loss(temp=1.0):
return coherence_loss(current_labels, tuples, priors_dev, consistency_groups, uid_to_idx, loss_weights,
targets=targets_dev, llm_diffs=llm_diffs_dev,
context_index_lists=context_index_lists, group_index_lists=group_index_lists,
context_index_tensors=context_index_tensors, group_index_tensors=group_index_tensors,
verbose=False, temperature=temp)
optimizer = optim.LBFGS([current_labels], lr=0.1, max_iter=20)
# optimizer = optim.AdamW([current_labels], lr=0.1)
losses = []
epochs = 10
for epoch in tqdm(range(epochs), desc="Optimizing"):
def closure():
if config.opt == "adamw":
optimizer = optim.AdamW([current_labels], lr=config.lr, weight_decay=config.weight_decay)
for epoch in tqdm(range(epochs), desc="Optimizing"):
optimizer.zero_grad()
loss = coherence_loss(current_labels, tuples, priors, consistency_groups, uid_to_idx, loss_weights, verbose=False)
# Temperature annealing: 2.0 → 0.5 over epochs (exponential decay)
temp = 0.5 + 1.5 * (0.95 ** epoch)
loss = get_loss(temp)
loss.backward()
return loss
assert torch.isfinite(current_labels).all(), "Labels contain NaN or Inf"
loss_val = optimizer.step(closure)
losses.append(loss_val.item())
# Gradient clipping for stability at higher learning rates
torch.nn.utils.clip_grad_norm_([current_labels], max_norm=1.0)
optimizer.step()
losses.append(loss.item())
# Log to wandb
if config.wandb and WANDB_AVAILABLE:
wandb.log({"loss": loss.item(), "temperature": temp, "epoch": epoch})
assert torch.isfinite(current_labels).all(), "Labels contain NaN or Inf"
else: # lbfgs
optimizer = optim.LBFGS([current_labels], lr=config.lr, max_iter=20)
for epoch in tqdm(range(epochs), desc="Optimizing"):
# Temperature annealing: 2.0 → 0.5 over epochs (exponential decay)
temp = 0.5 + 1.5 * (0.95 ** epoch)
def closure():
optimizer.zero_grad()
loss = get_loss(temp)
loss.backward()
# Gradient clipping for stability
torch.nn.utils.clip_grad_norm_([current_labels], max_norm=1.0)
return loss
loss = optimizer.step(closure)
losses.append(loss.item())
# Log to wandb
if config.wandb and WANDB_AVAILABLE:
wandb.log({"loss": loss.item(), "temperature": temp, "epoch": epoch})
assert torch.isfinite(current_labels).all(), "Labels contain NaN or Inf"
# Post-process
# Post-process (unchanged)
final_soft = F.softmax(current_labels, dim=1)
final_hard = torch.argmax(final_soft, dim=1).cpu().numpy()
# Output DF in memory
output_data = []
for uid in unique_uids:
idx = uid_to_idx[uid]
@@ -312,22 +428,20 @@ def run_backprop_experiment(labels, tuples, priors, consistency_groups, uid_to_i
row['learned_soft_0'] = float(final_soft[idx, 0].detach())
row['learned_soft_1'] = float(final_soft[idx, 1].detach())
row['learned_hard'] = int(final_hard[idx])
# row['prior_diff'] = avg_diffs.get(uid, 0.0)
output_data.append(row)
df_output = pd.DataFrame(output_data)
acc = (df_output['learned_hard'] == df_output['vanilla_label']).mean() if 'vanilla_label' in df_output.columns else 0.0
# LLM acc
llm_acc = 0.0
if not df_preds.empty and 'uid' in df_preds.columns and 'vanilla_label' in df_data.columns:
df_preds_local = df_preds.copy()
df_preds_local['hard_llm_pred'] = (df_preds_local['score'] > 0).astype(int)
llm_acc = (df_preds_local['hard_llm_pred'] == df_preds_local['vanilla_label']).mean()
# ICM corr
icm_corr = 0.0
output_dir = Path("outputs/backprop")
output_dir.mkdir(parents=True, exist_ok=True)
icm_path = output_dir.parent / "icm/truthfulqa/icm_final_labels.parquet"
if icm_path.exists():
df_icm = pd.read_parquet(icm_path)
@@ -343,37 +457,109 @@ def run_backprop_experiment(labels, tuples, priors, consistency_groups, uid_to_i
'acc': acc,
'icm_corr': icm_corr,
'final_loss': losses[-1] if losses else 0.0,
'llm_acc': llm_acc
'llm_acc': llm_acc,
'losses': losses
}
# %% [code]
# Experiment with loss weights: Loop over variations, print key results
# Replace the experiment loop with this
if __name__ == "__main__":
parser = simple_parsing.ArgumentParser()
parser.add_arguments(Config, dest="config")
config = parser.parse_args().config
logger.info(f"Running backprop experiment with config: {config}")
# FIXME this is only varying mutual weight, but should vary others too
# Baseline weights
base_weights = {
'mutual': 1.0,
'ranking': 0.5,
'prior': 0.1,
'direct': 0.5,
'entropy': 0.1,
'reward': 0.2
}
Path("outputs/backprop").mkdir(parents=True, exist_ok=True)
import matplotlib.pyplot as plt
base_weights = {
'mutual': 1.0,
'ranking': 0.5,
'prior': 0.1,
'direct': 0.5,
'entropy': 0.1,
'reward': 0.2
}
# Example: Vary mutual weight (one at a time, as suggested)
mutual_variations = [0.1, 0.5, 1.0, 2.0]
results = {}
import matplotlib.pyplot as plt
results = {}
for k in base_weights:
weights = base_weights.copy()
weights = {k:0 for k in weights} # Zero all
weights[k] = 1
print(f"\n--- Testing {k} = {weights[k]} ---")
res = run_backprop_experiment(labels, tuples, priors, consistency_groups, uid_to_idx, weights, df_data, df_preds, unique_uids, diffs)
results[k] = res
plt.plot(res['final_loss'])
plt.savefig(f"outputs/backprop/loss_{k}.png")
print(f"{k}: Acc {res['acc']:.4f}, ICM Corr {res['icm_corr']:.4f}, Final Loss {res['final_loss']:.4f}, LLM Acc {res['llm_acc']:.4f}")
# Define experiments: name -> weights dict
experiments = {}
# Individual loss ablations
for k in base_weights:
weights = {kk: 0 for kk in base_weights}
weights[k] = 1
experiments[k] = weights
# Combined losses
experiments['combined'] = {
'mutual': 0.6,
'prior': 0.5,
'ranking': 0.0,
'direct': 0.0,
'entropy': 0.0,
'reward': 0.0
}
# Initialize wandb if enabled
if config.use_wandb and WANDB_AVAILABLE:
wandb.init(
project="unsupervised-elicitation-backprop",
config={
"lr": config.lr,
"epochs": config.epochs,
"optimizer": config.opt,
"test_mode": config.test_mode,
"num_tuples": len(tuples) if not config.test_mode else config.subsample_size,
"num_labels": len(unique_uids)
},
name=f"lr{config.lr}_ep{config.epochs}_{'test' if config.test_mode else 'full'}"
)
# Run all experiments
for name, weights in experiments.items():
logger.info(f"\n--- Testing {name} ---")
# Start new wandb run for each experiment if enabled
if config.use_wandb and WANDB_AVAILABLE:
wandb.config.update({"experiment": name, "weights": weights})
res = run_backprop_experiment(labels, tuples, priors, consistency_groups, uid_to_idx, weights, df_data, df_preds, unique_uids,
targets, llm_diffs, context_index_lists, group_index_lists, context_index_tensors, group_index_tensors, config)
results[name] = res
logger.info(f"Losses for {name}: {res['losses']}")
# Log final metrics to wandb
if config.use_wandb and WANDB_AVAILABLE:
wandb.log({
f"{name}/final_acc": res['acc'],
f"{name}/final_loss": res['final_loss'],
f"{name}/icm_corr": res['icm_corr'],
f"{name}/llm_acc": res['llm_acc']
})
plt.figure()
plt.plot(res['losses'])
plt.title(f"Loss curve for {name}")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plot_path = Path(f"outputs/backprop/loss_{name}.png")
plt.savefig(plot_path)
# Log plot to wandb
if config.use_wandb and WANDB_AVAILABLE:
wandb.log({f"{name}/loss_curve": wandb.Image(str(plot_path))})
plt.close()
logger.info(f"Saved plot to {plot_path}")
logger.info(f"{name}: Acc {res['acc']:.4f}, ICM Corr {res['icm_corr']:.4f}, Final Loss {res['final_loss']:.4f}, LLM Acc {res['llm_acc']:.4f}")
# Finish wandb run
if config.use_wandb and WANDB_AVAILABLE:
wandb.finish()
# %%