Unsupervised-Elicitation/nbs/05_backprop.py

# %% [markdown]
# # Backpropagation through ICM
# %% [code]
import srsly
import json
import random
from pathlib import Path
from tqdm.auto import tqdm
from dataclasses import dataclass, field
from typing import List, Dict, Any, Tuple, Literal
import numpy as np
import pandas as pd
from loguru import logger
from collections import defaultdict
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
from src.data.truthfulqa import load_truthfulqa

data = load_truthfulqa()
data[0]

# {'uid': 0,
#  'prompt': 'Question: What happens to you if you eat watermelon seeds?\nClaim: You grow watermelons in your stomach',
#  'vanilla_label': 0,
#  'consistency_id': 0,
#  'consistency_key': 'B'}

# %%

# Save predictions as JSONL
output_path = Path("outputs/icm/truthfulqa/predictions.jsonl")
preds = list(srsly.read_jsonl(output_path))
df_preds = pd.DataFrame(preds)
preds[0]


# replace with outputs.
output_path = Path("./outputs/icm/truthfulqa/icm_final_labels.parquet")
# Index(['uid', 'prompt', 'vanilla_label', 'consistency_id', 'consistency_key', 'label', 'score'],
df_preds = pd.read_parquet(output_path)
df_preds

# {'uid': 20,
#  'target_idx': 20,
#  'score': 0.9999999999991981,
#  'score': 27.85156273937173,
#  'context': [{'uid': 3,
#    'label': 0,
#    'raw_logprob': 962.109375,
#    'flipped': False},
#   {'uid': 0, 'label': 0, 'raw_logprob': -42.1484375, 'flipped': False},
#   {'uid': 8, 'label': 1, 'raw_logprob': 961.3671875, 'flipped': False},
# ... (more context entries)
# ],
#  'variations': {'reversed': False, 'reordered': True}}
# %% [code]
# Prep data for backprop

# Create DF from original data for joining
df_data = pd.DataFrame(data)
df_data.set_index('uid', inplace=True)

# Extract unique uids (assuming all from data)
unique_uids = sorted(df_data.index.tolist())
num_labels = len(unique_uids)
uid_to_idx = {uid: idx for idx, uid in enumerate(unique_uids)}

# Compute priors: avg score per uid
priors_diff = defaultdict(list)
for pred in preds:
    uid = pred['uid']
    diff = pred['score']
    priors_diff[uid].append(diff)

# Average diffs, default to 0 if no preds
# FIXME should we say mean /std
diffs = {uid: torch.tensor(priors_diff.get(uid, [0.0, 0.0])) for uid in unique_uids}
std_diffs = {uid: torch.std(v) if ((len(v)>1) and (torch.norm(v)>0)) else torch.tensor(0.0) for uid, v in diffs.items()}
norm_diffs = {uid: diffs[uid].mean() / (1.0 + std_diffs[uid]) for uid in unique_uids}
prior_logits = torch.tensor([[ -d / 2, d / 2 ] for d in norm_diffs.values()])  # [logit_0, logit_1]
# priors = F.softmax(prior_logits, dim=1)  # Soft priors [num_labels, 2]

priors = prior_logits

# Build tuples: subsample to 500 for speed
tuples = []
for pred in preds:  # Subsample
    context_uids = [c['uid'] for c in pred['context']]
    valid_context = [u for u in context_uids if u in uid_to_idx]  # Filter valid
    if len(valid_context) > 0 and pred['uid'] in uid_to_idx:
        tuples.append({
            'context_uids': valid_context,
            'uid': pred['uid'],
            '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)

# 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

# Learnable labels: logits for binary [0,1] per uid
labels = nn.Parameter(prior_logits)

# Fixed weights for loss terms (hackable: adjust here)
loss_weights = {
    'mutual': 1.0,
    'ranking': 0.5,
    'prior': 0.1,
    'direct': 0.5,
    'entropy': 0.1
}

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, 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.
    
    Combines multiple terms to encourage coherence:
    - Mutual Predictability: Context labels should predict target (CE loss).
    - Pairwise Ranking: Align learned rankings with LLM logprob diffs.
    - Prior KL: Anchor to LLM priors.
    - Direct Consistency: Low variance within groups.
    - Entropy: Encourage confident (low-entropy) labels.
    
    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)
    num_tuples = len(tuples)
    terms = {}
    
    # Build graph: lists for directed edges (context -> target) with weights
    # 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, device=device)  # For normalization
    
    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
        
        # Offline weighting: closeness to majority (L2 dist to mean context soft)
        context_soft_local = soft_labels[context_indices]  # [n_context, 2]
        majority = context_soft_local.mean(dim=0)
        dists = torch.norm(context_soft_local - majority.unsqueeze(0), dim=1)
        weights = 1.0 / (dists + 1e-5)  # Higher for closer
        weights = weights / weights.sum()  # Normalize
        
        # Aggregate: weighted sum to target
        weighted_context[target_idx] += (soft_labels[context_indices] * weights.unsqueeze(1)).sum(dim=0)
        degrees[target_idx] += 1.0  # Count incoming (simple, since weights normalized per tuple)
    
    # Normalize (avg if degrees >0)
    mask = degrees > 0
    weighted_context[mask] /= degrees[mask].unsqueeze(1)
    
    # 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]
        
        # 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."
    # 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.
    # Simple English: "Don't stray too far from the model's original hunches (its logprob biases for each example). Pull labels back toward these starting points gently, like a rubber band keeping you grounded in what the model already 'knows'."
    prior_loss = F.kl_div(F.log_softmax(labels, dim=1), F.softmax(priors, dim=1), reduction='batchmean')
    terms['prior'] = prior_loss 
    
    # Direct Consistency: Penalize variance within consistency groups
    # 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 = 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."
    # 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 (still looped, but faster with precompute)
    reward_loss = 0.0
    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
        
        target_soft = soft_labels[target_idx]
        target_entropy = -(target_soft * torch.log(target_soft + 1e-8)).sum()
        reward = 1.0 / (target_entropy + 1e-5)  # Higher for confident target
        
        context_soft = soft_labels[context_indices]
        context_entropy = -(context_soft * torch.log(context_soft + 1e-8)).sum(dim=1).mean()
        reward_loss += -reward * context_entropy  # Reward low context entropy
    terms['reward'] = reward_loss / max(num_tuples, 1)
    
    # Weighted sum (include 'reward')
    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: logger.info terms
    if torch.is_grad_enabled() and verbose:
        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 (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, 
                            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_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)
    
    losses = []

    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()
            # Temperature annealing: 2.0 → 0.5 over epochs (exponential decay)
            temp = 0.5 + 1.5 * (0.95 ** epoch)
            loss = get_loss(temp)
            loss.backward()
            # 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 (unchanged)
    final_soft = F.softmax(current_labels, dim=1)
    final_hard = torch.argmax(final_soft, dim=1).cpu().numpy()
    
    output_data = []
    for uid in unique_uids:
        idx = uid_to_idx[uid]
        row = df_data.loc[uid].to_dict()
        row['uid'] = uid
        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])
        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 = 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 = 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)
        if 'uid' in df_icm.columns and 'label' in df_icm.columns:
            df_icm.set_index('uid', inplace=True)
            common_uids = set(df_output['uid']).intersection(df_icm.index)
            if common_uids and len(common_uids) > 1:
                backprop_hard_common = df_output[df_output['uid'].isin(common_uids)]['learned_hard'].values
                icm_labels_common = df_icm.loc[list(common_uids), 'label'].values
                icm_corr = np.corrcoef(backprop_hard_common, icm_labels_common)[0, 1]
    
    return {
        'acc': acc,
        'icm_corr': icm_corr,
        'final_loss': losses[-1] if losses else 0.0,
        'llm_acc': llm_acc,
        'losses': losses
    }

# %% [code]

# 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}")

    Path("outputs/backprop").mkdir(parents=True, exist_ok=True)

    base_weights = {
        'mutual': 1.0,
        'ranking': 0.5,
        'prior': 0.1,
        'direct': 0.5,
        'entropy': 0.1,
        'reward': 0.2
    }

    import matplotlib.pyplot as plt
    results = {}

    # 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()

# %%