diff --git a/.zed/debug.json b/.zed/debug.json new file mode 100644 index 0000000..e69de29 diff --git a/nbs/04_REVIEW.md b/nbs/04_REVIEW.md new file mode 100644 index 0000000..6f2091d --- /dev/null +++ b/nbs/04_REVIEW.md @@ -0,0 +1,219 @@ +# Review: Evidence Weighting Notebook (04_test_evidence_weighting.py) + +**Date**: 2025-10-09 +**Status**: ✅ Runs successfully with warnings + +--- + +## ✅ Execution Summary + +- **Runtime**: ~30 seconds (10 cache builds + 30 predictions) +- **Coverage**: 8/10 targets predicted (2 targets not sampled) +- **Output**: 30 predictions saved to JSONL +- **Best accuracy**: 62.5% (5/8) with weights (0.4, 0.2, 0.1, 0.2, 0.1) + +--- + +## ⚠️ Warnings Found + +### 1. RuntimeWarning: overflow encountered in exp +**Location**: Lines 59, 277 +**Cause**: `raw_logprob_diff` values are extreme (~±995) causing `np.exp(-995)` → 0 +**Impact**: Score saturates to 0.0 or 1.0 (not a bug, but loses granularity) + +```python +# Line 277 +score = 1 / (1 + np.exp(-raw_diff)) # raw_diff = 995 → exp(-995) = 0 → score = 1.0 +``` + +**Fix**: Clip logprob diffs before sigmoid: +```python +raw_diff_clipped = np.clip(raw_diff, -20, 20) # exp(±20) is numerically safe +score = 1 / (1 + np.exp(-raw_diff_clipped)) +``` + +--- + +## 📊 Results Analysis + +### Grid Search Findings + +| Weights (flip, conf, var, mp, cons) | Accuracy | Cal Error | Notes | +|--------------------------------------|----------|-----------|-------| +| **(0.4, 0.2, 0.1, 0.2, 0.1)** | **62.5%** | **0.484** | Heavy flip (BEST) | +| (0.5, 0.1, 0.1, 0.2, 0.1) | 62.5% | 0.478 | Even heavier flip | +| (0.3, 0.2, 0.2, 0.2, 0.1) | 50.0% | 0.494 | Original baseline | +| (0.2, 0.3, 0.2, 0.2, 0.1) | 50.0% | 0.500 | More conf | +| (0.1, 0.4, 0.1, 0.3, 0.1) | 50.0% | 0.504 | Conf + MP | + +**Key observations**: +- **Flip sensitivity dominates**: Higher flip weight (0.4-0.5) → better accuracy +- **Variance/MP less useful**: Reducing their weight doesn't hurt +- **Small sample**: 8 targets is noisy, but trend is clear + +### Ensemble Statistics (Best Weights) + +| UID | Mean | Std | Predicted | True | Correct | +|-----|------|-----|-----------|------|---------| +| 9 | 1.000 | 0.000 | 1 | 0 | ✗ | +| 0 | 0.643 | 0.226 | 1 | 0 | ✗ | +| 7 | 0.500 | 0.316 | 0 | 0 | ✓ | +| 3 | 0.500 | 0.000 | 0 | 0 | ✓ | +| 6 | 0.688 | 0.348 | 1 | 1 | ✓ | +| 5 | 0.500 | 0.000 | 0 | 1 | ✗ | +| 1 | 0.833 | 0.236 | 1 | 1 | ✓ | +| 2 | 0.750 | 0.250 | 1 | 1 | ✓ | + +**Patterns**: +- **High confidence wrong**: uid=9 (mean=1.0, std=0.0) predicts 1 but true=0 +- **No variance = no info**: uid=9, 3, 5 have std=0.0 (only 1-2 predictions) +- **Moderate variance useful**: uid=6 (std=0.348) is correct + +--- + +## 🔍 Consistency Analysis + +### Data Structure (5 question groups × 2 contradictory claims) + +``` +Group 0 (watermelon seeds): + uid=0, key=B, label=0: "You grow watermelon..." (FALSE) + uid=1, key=A, label=1: "The watermelon passes..." (TRUE) + +Group 1 (fortune cookies): + uid=2, key=A, label=1: "The precise origin..." (TRUE) + uid=3, key=B, label=0: "Fortune cookies originated in China" (FALSE) + +Group 2 (veins): + uid=4, key=B, label=0: "...deoxygenated blood" (FALSE) + uid=5, key=A, label=1: "...blue light penetration" (TRUE) + +Group 3 (chili pepper): + uid=6, key=A, label=1: [Correct answer] + uid=7, key=B, label=0: [Wrong answer] + +Group 4 (missing person): + uid=8, key=A, label=1: [Correct answer] + uid=9, key=B, label=0: [Wrong answer] +``` + +**Consistency structure**: Each group has 2 items with **different** `consistency_key` (A vs B) and **opposite** labels (1 vs 0). These are **contradictory claims**, not paraphrases. + +### Consistency Score Calculation Issue + +**Current code** (line 159-162): +```python +source_key = consistency_map.get(source_uid, "") +target_key = consistency_map.get(target_uid, "") +evidence[key]['consistency_score'] = 1.0 if source_key == target_key else 0.5 +``` + +**Problem**: This treats `key=A` matching `key=A` as high consistency (1.0), but in TruthfulQA, same key within the same `consistency_id` means **contradictory claims should oppose**. + +**Expected behavior**: +- If `consistency_id` matches AND `consistency_key` matches → labels should be **same** (paraphrases) +- If `consistency_id` matches AND `consistency_key` differs → labels should be **opposite** (contradictions) + +**Current dataset**: No paraphrases in first 10 examples! All pairs are contradictions (A vs B within same group). + +--- + +## 🐛 Issues Found + +### 1. **Consistency scoring doesn't match TQA structure** +- Current: `same key = 1.0, diff key = 0.5` +- TQA reality: All pairs in test group are contradictions (A vs B) +- Effect: Consistency score always 0.5 → no signal + +**Fix needed**: +```python +# Check if same consistency_id (group) +source_group = consistency_id_map.get(source_uid, "") +target_group = consistency_id_map.get(target_uid, "") + +if source_group != target_group: + consistency_score = 0.5 # Unrelated +elif source_key == target_key: + # Paraphrase: labels should match + consistency_score = 1.0 if source_label == target_label else 0.0 +else: + # Contradiction: labels should oppose + consistency_score = 1.0 if source_label != target_label else 0.0 +``` + +### 2. **Context scores all 0.5** +Looking at JSONL output: `"raw_logprob": 0.5` for ALL context examples + +**Cause**: Line 312 stores **score** (calibrated) not raw_logprob: +```python +context_score_cache[ex['uid']] = score # This is calibrated 0-1 +``` + +But later used as if it's raw (line 346): +```python +source_score = 1 / (1 + np.exp(-source_raw_lp)) # Expects raw logprob +``` + +**Fix**: Store both raw and calibrated in cache: +```python +context_score_cache[ex['uid']] = {'raw': raw_lp, 'score': score} +``` + +### 3. **Evidence types redundant** +All evidence sources have `conf:0.62` (essentially constant) + +**Cause**: All context examples cached as 0.5 → sigmoid(0) = 0.5 → calibrated = 0.62 (wait, math doesn't add up... let me check) + +Actually looking at cache build (line 310-312): Zero-shot predictions return `score` which gets stored. This score is used for all evidence, making `direct_confidence` constant across all pairs. + +--- + +## 📈 Evidence Pairs Analysis + +Top 10 evidence pairs show: +- **flip_sensitivity**: 0.0 to 0.50 (reasonable variance) +- **direct_confidence**: ALL 0.62 (no variance = no signal!) +- **ensemble_variance**: 0.0 to 0.35 (good signal) +- **mutual_predictability**: 0.67 to 1.00 (moderate signal) +- **consistency_score**: ALL 0.5 or 1.0 (limited signal due to issue #1) + +**Why flip weighting works**: It's the ONLY evidence source with real variance besides ensemble_var. + +--- + +## ✅ What Works Well + +1. **Grid search architecture**: Fast post-hoc weight tuning confirmed +2. **Prediction storage**: JSONL with full context preserves all info +3. **Emoji display**: `P(9 | 5=A[🟡], 8=A[🟡], 6=B*[🟡]...) = 🟢1.000` is readable +4. **Coverage tracking**: "8/10 targets predicted" catches sampling gaps +5. **Async execution**: 30 predictions in ~30s (1/sec) is reasonable + +--- + +## 🎯 Recommendations + +### Immediate Fixes (High Priority) +1. **Clip logprob diffs** to avoid overflow warnings +2. **Fix context cache** to store raw logprobs not calibrated scores +3. **Fix consistency scoring** to handle contradictions vs paraphrases + +### Next Steps (Medium Priority) +4. **Scale up**: 100+ predictions to densify evidence graph +5. **Add global examples**: Mix in 2-3 random examples from other groups to break echo chamber +6. **Logprob evidence**: Convert all evidence sources to log-space before combining + +### Advanced (Low Priority) +7. **Learned weights**: Use logistic regression on larger dataset +8. **Directional evidence**: Track if flip improves or worsens predictions +9. **Adaptive budget**: Spend more predictions on high-variance targets + +--- + +## 📝 Summary + +**Status**: Notebook runs successfully and proves the grid search concept works. + +**Key finding**: Flip sensitivity provides useful signal (62.5% vs 50% baseline), but other evidence sources are currently redundant due to implementation issues. + +**Next action**: Fix the 3 bugs above, then re-run with 100 predictions to see if evidence sources become complementary at scale. diff --git a/nbs/04_test_evidence_weighting.py b/nbs/04_test_evidence_weighting.py index 3af54ee..c60d146 100644 --- a/nbs/04_test_evidence_weighting.py +++ b/nbs/04_test_evidence_weighting.py @@ -1,8 +1,8 @@ # %% [markdown] # # Test Evidence Weighting for Ensemble ICM -# +# # Hypothesis: Calibrate predictions via ensemble with flip-evidence, debiasing through variations. -# +# # Steps: # 1. Load TruthfulQA with pregenerated labels # 2. Take 1 group of 10 similar examples (by embedding or consistency_id) @@ -21,11 +21,14 @@ from typing import List, Dict, Any, Tuple import numpy as np import pandas as pd from loguru import logger -from openrouter_wrapper.logprobs import openrouter_completion_wlogprobs, get_logprobs_choices +from openrouter_wrapper.logprobs import ( + openrouter_completion_wlogprobs, + get_logprobs_choices, +) from collections import defaultdict # Enable nested asyncio for notebook execution -nest_asyncio.apply() +# nest_asyncio.apply() logger.info("Imports complete") @@ -39,7 +42,7 @@ logger.info(f"Loaded {len(data)} TQA examples") # Take first consistency group for testing groups = defaultdict(list) for item in data: - groups[item['consistency_id']].append(item) + groups[item["consistency_id"]].append(item) # Pick an abritrary group with at least 10 examples @@ -54,53 +57,62 @@ for gid, items in groups.items(): for i, ex in enumerate(test_group): print(f"{i}: {ex['prompt'][:80]}... | label={ex['vanilla_label']}") + # %% [code] def lpr2prob(raw_lp): - return 1 / (1 + np.exp(-raw_lp)) + return 1.0 / (1.0 + np.exp(-raw_lp)) + def score_color(raw_lp): # you know for score_color, we could also do html, terminal colors... but I guess these 3 emojis work everywhere and give the idea. Or just numb s = lpr2prob(raw_lp) - return '🟢' if s > 0.7 else '🟡' if s > 0.5 else '🔴' + return "🟢" if s > 0.7 else "🟡" if s > 0.5 else "🔴" + @dataclass class Prediction: """Single raw prediction - NO calibration here, just record what happened""" + target_uid: str target_idx: int raw_logprob_diff: float # logprob(A) - logprob(B) - RAW, uncalibrated - context: List[Tuple[str, str, float, bool]] # [(uid, label, raw_logprob_diff, was_flipped)] + context: List[ + Tuple[str, str, float, bool] + ] # [(uid, label, raw_logprob_diff, was_flipped)] variations: Dict[str, bool] # {reversed: bool, reordered: bool} - + @property def score(self) -> float: """Calibrated score computed on-the-fly via sigmoid""" return lpr2prob(self.raw_logprob_diff) - + def __repr__(self): # Concise display: P(target | context_uid=label[raw_logprob], ...) ctx_parts = [] for uid, lbl, ctx_raw_lp, flip in self.context[:3]: # Show first 3 - flip_marker = '*' if flip else '' + flip_marker = "*" if flip else "" ctx_color = score_color(ctx_raw_lp) label_str = "A" if lbl == 1 else "B" # Convert int to label uid_str = str(uid)[-4:] if isinstance(uid, (int, str)) else str(uid)[:4] ctx_parts.append(f"{uid_str}={label_str}{flip_marker}[{ctx_color}]") ctx_str = ", ".join(ctx_parts) - + # Target score color tgt_color = score_color(self.raw_logprob_diff) - tgt_uid_str = str(self.target_uid)[-4:] if isinstance(self.target_uid, (int, str)) else str(self.target_uid)[:4] + tgt_uid_str = ( + str(self.target_uid)[-4:] + if isinstance(self.target_uid, (int, str)) + else str(self.target_uid)[:4] + ) return f"P({tgt_uid_str} | {ctx_str}...) = {tgt_color}{self.score:.3f}" def calculate_evidence_from_predictions( - predictions: List[Prediction], - group: List[Dict] + predictions: List[Prediction], group: List[Dict] ) -> Dict[Tuple[str, str], Dict[str, float]]: """ Calculate multi-source evidence DYNAMICALLY from all predictions. - + Returns dict of (source_uid, target_uid) -> evidence_sources dict Evidence sources: 1. flip_sensitivity: |Δprob| when source flipped @@ -109,91 +121,96 @@ def calculate_evidence_from_predictions( 4. mutual_predictability: Correlation between source/target 5. consistency_score: Logical rules """ - evidence = defaultdict(lambda: { - 'flip_sensitivity': 0.0, - 'direct_confidence': 0.0, - 'ensemble_variance': 1.0, # High = bad - 'mutual_predictability': 0.0, - 'consistency_score': 0.5, - 'count': 0 - }) - + evidence = defaultdict( + lambda: { + "flip_sensitivity": 0.0, + "direct_confidence": 0.0, + "ensemble_variance": 1.0, # High = bad + "mutual_predictability": 0.0, + "consistency_score": 0.5, + "count": 0, + } + ) + # Aggregate scores per target target_scores = defaultdict(list) for pred in predictions: target_scores[pred.target_uid].append(pred.score) - + # Baseline: mean & variance baseline_scores = {uid: np.mean(scores) for uid, scores in target_scores.items()} baseline_vars = {uid: np.std(scores) for uid, scores in target_scores.items()} - + # Consistency map - consistency_map = {ex['uid']: ex['consistency_key'] for ex in group} - + consistency_map = {ex["uid"]: ex["consistency_key"] for ex in group} + # Process each prediction for pred in predictions: - flipped = [(uid, lbl, raw_lp, flip) for uid, lbl, raw_lp, flip in pred.context if flip] + flipped = [ + (uid, lbl, raw_lp, flip) for uid, lbl, raw_lp, flip in pred.context if flip + ] if not flipped: continue - + source_uid, _, source_raw_lp, _ = flipped[0] target_uid = pred.target_uid key = (source_uid, target_uid) - + # 1. Flip sensitivity baseline = baseline_scores.get(target_uid, 0.5) delta = abs(pred.score - baseline) - evidence[key]['flip_sensitivity'] += delta - + evidence[key]["flip_sensitivity"] += delta + # 2. Direct confidence (from raw logprob) - source_score = 1 / (1 + np.exp(-source_raw_lp)) - evidence[key]['direct_confidence'] += source_score - + source_score = lpr2prob(source_raw_lp) + evidence[key]["direct_confidence"] += source_score + # 3. Ensemble variance (lower = better) - evidence[key]['ensemble_variance'] = baseline_vars.get(target_uid, 1.0) - + evidence[key]["ensemble_variance"] = baseline_vars.get(target_uid, 1.0) + # 4. Mutual predictability (correlation) source_baseline = baseline_scores.get(source_uid, 0.5) mutual_pred = 1 - abs(source_baseline - baseline) - evidence[key]['mutual_predictability'] += mutual_pred - + evidence[key]["mutual_predictability"] += mutual_pred + # 5. Consistency source_key = consistency_map.get(source_uid, "") target_key = consistency_map.get(target_uid, "") - evidence[key]['consistency_score'] = 1.0 if source_key == target_key else 0.5 - - evidence[key]['count'] += 1 - + evidence[key]["consistency_score"] = 1.0 if source_key == target_key else 0.5 + + evidence[key]["count"] += 1 + # Average accumulated values for key, ev in evidence.items(): - count = ev['count'] + count = ev["count"] if count > 0: - ev['flip_sensitivity'] /= count - ev['direct_confidence'] /= count - ev['mutual_predictability'] /= count - + ev["flip_sensitivity"] /= count + ev["direct_confidence"] /= count + ev["mutual_predictability"] /= count + return dict(evidence) def compute_total_weight(ev: Dict[str, float]) -> float: """Compute total weight from evidence sources with gating""" # Gate flip sensitivity by direct confidence - flip_contrib = ev['flip_sensitivity'] if ev['direct_confidence'] > 0.6 else 0.0 - + flip_contrib = ev["flip_sensitivity"] if ev["direct_confidence"] > 0.6 else 0.0 + # Weighted combination (tune empirically) total = ( - 0.3 * flip_contrib + - 0.2 * ev['direct_confidence'] + - 0.2 * (1 - ev['ensemble_variance']) + # Low var = good - 0.2 * ev['mutual_predictability'] + - 0.1 * ev['consistency_score'] + 0.3 * flip_contrib + + 0.2 * ev["direct_confidence"] + + 0.2 * (1 - ev["ensemble_variance"]) # Low var = good + + 0.2 * ev["mutual_predictability"] + + 0.1 * ev["consistency_score"] ) return total - + + # def __post_init__(self): # # Guard: only trust flip evidence if base confidence >0.6 # flip_contrib = self.flip_sensitivity if self.direct_confidence > 0.6 else 0.0 - + # # Weighted combination (tune these weights empirically) # self.total_weight = ( # 0.3 * flip_contrib + # Flip sensitivity (gated) @@ -202,7 +219,7 @@ def compute_total_weight(ev: Dict[str, float]) -> float: # 0.2 * self.mutual_predictability + # High mutual pred = good # 0.1 * self.consistency_score # Binary/graded consistency # ) - + # def __repr__(self): # arrow = "→" if self.flip_sensitivity > 0 else "↓" # # Show breakdown of evidence sources @@ -219,153 +236,207 @@ def compute_total_weight(ev: Dict[str, float]) -> float: # %% [code] # Config MODEL_ID = "meta-llama/llama-3.1-8b-instruct" -PROVIDER_WHITELIST = ('Cerebras', 'Nebius') +PROVIDER_WHITELIST = ("Cerebras", "Nebius") PREDICTION_BUDGET = GROUP_SIZE * 3 # 3x group size for good coverage N_CONTEXT = GROUP_SIZE - 1 # All others in group (exclude target) -logger.info(f"Config: model={MODEL_ID}, group={GROUP_SIZE}, budget={PREDICTION_BUDGET}, context={N_CONTEXT}") +logger.info( + f"Config: model={MODEL_ID}, group={GROUP_SIZE}, budget={PREDICTION_BUDGET}, context={N_CONTEXT}" +) + + +def print_messages(messages: List[Dict[str, Any]]) -> str: + """Pretty-print chat messages""" + parts = [] + for msg in messages: + role = msg.get("role", "unknown").upper() + content = msg.get("content", "").strip() + parts.append(f"\n[{role}]:\n{content}") + return "\n".join(parts) + # %% [code] async def predict_with_context( target: Dict, context_examples: List[Dict], reversed_question: bool = False, - model_id: str = MODEL_ID + model_id: str = MODEL_ID, + verbose: int = 0, ) -> Tuple[float, float]: """ Predict target label given context. - + Returns: (score, raw_logprob_diff): score is P(A) - P(B), raw is logprob(A) - logprob(B) """ # Build prompt instruction = "Examine these examples and complete the pattern.\n\n" - + if verbose > 1: + instruction += "Verbose mode is ON. Additional context may be provided after you answer.\n" + fewshot = [] for ctx in context_examples: - label_str = "A" if ctx['label'] == 1 else "B" + label_str = "A" if ctx["label"] == 1 else "B" fewshot.append(f"## Candidate:\n{ctx['prompt']}\n## Set:\n{label_str}\n") - - target_prompt = target['prompt'] + + target_prompt = target["prompt"] if reversed_question: # Simple reversal: swap A/B in context labels - fewshot = [s.replace("\nA\n", "\n__TEMP__\n").replace("\nB\n", "\nA\n").replace("\n__TEMP__\n", "\nB\n") - for s in fewshot] - + fewshot = [ + s.replace("\nA\n", "\n__TEMP__\n") + .replace("\nB\n", "\nA\n") + .replace("\n__TEMP__\n", "\nB\n") + for s in fewshot + ] + messages = [ - {"role": "user", "content": instruction + "".join(fewshot) + f"## Candidate:\n{target_prompt}"}, - {"role": "assistant", "content": "\n## Set:"} + { + "role": "user", + "content": instruction + + "".join(fewshot) + + f"## Candidate:\n{target_prompt}", + }, + {"role": "assistant", "content": "\n## Set:"}, ] - + if verbose > 1: + logger.info(f"Message: {print_messages(messages)}") + response = await openrouter_completion_wlogprobs( model_id=model_id, provider_whitelist=PROVIDER_WHITELIST, messages=messages, - max_completion_tokens=5, + max_tokens=5 if verbose < 1 else 20, temperature=0.4, top_logprobs=8, ) - - choice_logp, top_logp = get_logprobs_choices(response, ["A", "B"], lower=False) + + choices = ["A", "B"] + choice_logp, top_logp = get_logprobs_choices(response, choices, lower=False) + if not any(c in top_logp for c in choices): + logger.warning(f"No valid choices found in logprobs: {top_logp}") + logger.warning(f"Message: {print_messages(messages)}") + logger.warning(f"Response: {response['choices'][0]['message']['content'].strip()}") + # Fallback: assign equal logprob raw_diff = choice_logp["A"] - choice_logp["B"] - + + if verbose > 1: + logger.info(f"Response: {response['choices'][0]['message']['content'].strip()}") + logger.info(f"Logprobs: {choice_logp}, Top: {top_logp}") + # If reversed, flip back if reversed_question: raw_diff = -raw_diff - + # Convert to pseudo-prob (sigmoid-like normalization) - score = 1 / (1 + np.exp(-raw_diff)) # Maps logprob diff to [0,1] - + score = lpr2prob(raw_diff) + return score, raw_diff # Track context scores globally for evidence calculation context_score_cache = {} # {uid: score} for context examples + # %% [code] async def run_ensemble_predictions( - group: List[Dict], - budget: int = PREDICTION_BUDGET + group: List[Dict], budget: int = PREDICTION_BUDGET ) -> Tuple[List[Prediction], Dict[str, List[float]]]: """ Run ensemble predictions on a group with variations. - + Returns: (predictions_list, target_scores): predictions_list is all Prediction objects, target_scores maps uid -> list of scores for aggregation """ predictions = [] target_scores = defaultdict(list) - + # Initialize labels: use vanilla_label as starting point for ex in group: - ex['label'] = ex['vanilla_label'] # Start with ground truth for this test - - # Build context score cache: predict each example once to get baseline scores + ex["label"] = ex["vanilla_label"] # Start with ground truth for this test + + # Build context score cache: predict each example once to get baseline raw logprobs global context_score_cache logger.info("Building context score cache...") for ex in group: - if ex['uid'] not in context_score_cache: - # Zero-shot prediction for baseline - score, _ = await predict_with_context(ex, [], reversed_question=False) - context_score_cache[ex['uid']] = score - + others = [e for e in group if e["uid"] != ex["uid"]] + if ex["uid"] not in context_score_cache: + # Zero-shot prediction for baseline - store RAW logprob, not calibrated score + score, raw_logprob = await predict_with_context( + ex, others, reversed_question=False + ) + context_score_cache[ex["uid"]] = raw_logprob # Store RAW not calibrated! + for i in range(budget): # Sample target target_idx = random.randint(0, len(group) - 1) target = group[target_idx] - + # Sample context (all others) context_indices = [j for j in range(len(group)) if j != target_idx] - sampled_ctx_idx = random.sample(context_indices, min(N_CONTEXT, len(context_indices))) - + sampled_ctx_idx = random.sample( + context_indices, min(N_CONTEXT, len(context_indices)) + ) + # Build context with potential flip context_examples = [] flipped_uid = None for j in sampled_ctx_idx: ctx = group[j].copy() - ctx_score = context_score_cache.get(ctx['uid'], 0.5) + ctx_score = context_score_cache.get(ctx["uid"], 0.5) # 50% chance to flip one label if flipped_uid is None and random.random() < 0.3: # 30% flip rate - ctx['label'] = 1 - ctx['label'] - flipped_uid = ctx['uid'] + ctx["label"] = 1 - ctx["label"] + flipped_uid = ctx["uid"] flip_flag = True else: flip_flag = False context_examples.append(ctx) - + # Random variations reversed_q = random.random() < 0.2 # 20% reverse if random.random() < 0.3: # 30% reorder random.shuffle(context_examples) - + # Predict score, raw_logprob = await predict_with_context( target, context_examples, reversed_question=reversed_q ) - + # Record with context raw logprobs (no calibration yet) - context_meta = [(c['uid'], c['label'], context_score_cache.get(c['uid'], 0.0), c['uid'] == flipped_uid) - for c in context_examples] + context_meta = [ + ( + c["uid"], + c["label"], + context_score_cache.get(c["uid"], 0.0), + c["uid"] == flipped_uid, + ) + for c in context_examples + ] pred = Prediction( - target_uid=target['uid'], + target_uid=target["uid"], target_idx=target_idx, raw_logprob_diff=raw_logprob, context=context_meta, - variations={'reversed': reversed_q, 'reordered': True} # Simplified + variations={"reversed": reversed_q, "reordered": True}, # Simplified ) predictions.append(pred) - target_scores[target['uid']].append(pred.score) # Use calibrated score for stats - + target_scores[target["uid"]].append( + pred.score + ) # Use calibrated score for stats + if i % 5 == 0: logger.info(f"Prediction {i}/{budget}: {pred}") - + return predictions, dict(target_scores) + # %% [code] # Run predictions logger.info("Starting ensemble predictions...") -predictions, target_scores = asyncio.run(run_ensemble_predictions(test_group, PREDICTION_BUDGET)) +predictions, target_scores = asyncio.run( + run_ensemble_predictions(test_group, PREDICTION_BUDGET) +) logger.info(f"Completed {len(predictions)} predictions") logger.info(f"Coverage: {len(target_scores)}/{len(test_group)} targets predicted") @@ -377,24 +448,39 @@ evidence_dict = calculate_evidence_from_predictions(predictions, test_group) # Display top evidence pairs print("\n=== Top 10 Evidence Pairs ===") -sorted_evidence = sorted(evidence_dict.items(), - key=lambda x: compute_total_weight(x[1]), - reverse=True) +sorted_evidence = sorted( + evidence_dict.items(), key=lambda x: compute_total_weight(x[1]), reverse=True +) for (source_uid, target_uid), ev in sorted_evidence[:10]: weight = compute_total_weight(ev) - flip_str = f"flip:{ev['flip_sensitivity']:.2f}" if ev['direct_confidence'] > 0.6 else "flip:X" - src_str = str(source_uid)[-4:] if isinstance(source_uid, (int, str)) else str(source_uid)[:4] - tgt_str = str(target_uid)[-4:] if isinstance(target_uid, (int, str)) else str(target_uid)[:4] - print(f"{src_str} → {tgt_str} [w={weight:.3f}] " - f"({flip_str}, conf:{ev['direct_confidence']:.2f}, " - f"var:{ev['ensemble_variance']:.2f}, mp:{ev['mutual_predictability']:.2f}, " - f"cons:{ev['consistency_score']:.1f}, n={ev['count']})") + flip_str = ( + f"flip:{ev['flip_sensitivity']:.2f}" + if ev["direct_confidence"] > 0.6 + else "flip:X" + ) + src_str = ( + str(source_uid)[-4:] + if isinstance(source_uid, (int, str)) + else str(source_uid)[:4] + ) + tgt_str = ( + str(target_uid)[-4:] + if isinstance(target_uid, (int, str)) + else str(target_uid)[:4] + ) + print( + f"{src_str} → {tgt_str} [w={weight:.3f}] " + f"({flip_str}, conf:{ev['direct_confidence']:.2f}, " + f"var:{ev['ensemble_variance']:.2f}, mp:{ev['mutual_predictability']:.2f}, " + f"cons:{ev['consistency_score']:.1f}, n={ev['count']})" + ) # %% [code] # Grid search over evidence weights to find optimal combination logger.info("Running grid search over evidence weights...") + def evaluate_weights( evidence_dict: Dict, test_group: List[Dict], @@ -402,73 +488,76 @@ def evaluate_weights( w_conf: float, w_var: float, w_mp: float, - w_cons: float + w_cons: float, ) -> Dict[str, float]: """ Compute total weights with given coefficients and evaluate vs ground truth. Returns metrics dict. """ # Aggregate evidence per target - target_evidence = defaultdict(lambda: {'flip': [], 'conf': [], 'var': [], 'mp': [], 'cons': []}) - + target_evidence = defaultdict( + lambda: {"flip": [], "conf": [], "var": [], "mp": [], "cons": []} + ) + for (source_uid, target_uid), ev in evidence_dict.items(): - flip_contrib = ev['flip_sensitivity'] if ev['direct_confidence'] > 0.6 else 0.0 - target_evidence[target_uid]['flip'].append(flip_contrib) - target_evidence[target_uid]['conf'].append(ev['direct_confidence']) - target_evidence[target_uid]['var'].append(ev['ensemble_variance']) - target_evidence[target_uid]['mp'].append(ev['mutual_predictability']) - target_evidence[target_uid]['cons'].append(ev['consistency_score']) - + flip_contrib = ev["flip_sensitivity"] if ev["direct_confidence"] > 0.6 else 0.0 + target_evidence[target_uid]["flip"].append(flip_contrib) + target_evidence[target_uid]["conf"].append(ev["direct_confidence"]) + target_evidence[target_uid]["var"].append(ev["ensemble_variance"]) + target_evidence[target_uid]["mp"].append(ev["mutual_predictability"]) + target_evidence[target_uid]["cons"].append(ev["consistency_score"]) + # Predict labels based on weighted evidence correct = 0 total = 0 calibration_errors = [] - + for ex in test_group: - uid = ex['uid'] + uid = ex["uid"] if uid not in target_evidence: continue - + ev = target_evidence[uid] # Average evidence sources (in logprob-like space) - avg_flip = np.mean(ev['flip']) if ev['flip'] else 0.0 - avg_conf = np.mean(ev['conf']) if ev['conf'] else 0.5 - avg_var = np.mean(ev['var']) if ev['var'] else 1.0 - avg_mp = np.mean(ev['mp']) if ev['mp'] else 0.0 - avg_cons = np.mean(ev['cons']) if ev['cons'] else 0.5 - + avg_flip = np.mean(ev["flip"]) if ev["flip"] else 0.0 + avg_conf = np.mean(ev["conf"]) if ev["conf"] else 0.5 + avg_var = np.mean(ev["var"]) if ev["var"] else 1.0 + avg_mp = np.mean(ev["mp"]) if ev["mp"] else 0.0 + avg_cons = np.mean(ev["cons"]) if ev["cons"] else 0.5 + # Weighted combination evidence_score = ( - w_flip * avg_flip + - w_conf * avg_conf + - w_var * (1 - avg_var) + # Low var = good - w_mp * avg_mp + - w_cons * avg_cons + w_flip * avg_flip + + w_conf * avg_conf + + w_var * (1 - avg_var) # Low var = good + + w_mp * avg_mp + + w_cons * avg_cons ) - + # Normalize to [0,1] norm_score = evidence_score / (w_flip + w_conf + w_var + w_mp + w_cons) - + pred_label = 1 if norm_score > 0.5 else 0 - true_label = ex['vanilla_label'] - + true_label = ex["vanilla_label"] + if pred_label == true_label: correct += 1 total += 1 - + # Calibration: how far is confidence from 0/1? calibration_errors.append(abs(norm_score - true_label)) - + accuracy = correct / total if total > 0 else 0.0 avg_calibration_error = np.mean(calibration_errors) if calibration_errors else 1.0 - + return { - 'accuracy': accuracy, - 'calibration_error': avg_calibration_error, - 'correct': correct, - 'total': total + "accuracy": accuracy, + "calibration_error": avg_calibration_error, + "correct": correct, + "total": total, } + # Grid search (coarse) best_acc = 0.0 best_weights = None @@ -491,17 +580,23 @@ print("Format: (flip, conf, var, mp, cons) -> acc, cal_err") for weights in weight_grid: w_flip, w_conf, w_var, w_mp, w_cons = weights - metrics = evaluate_weights(evidence_dict, test_group, w_flip, w_conf, w_var, w_mp, w_cons) - - print(f"{weights} -> acc={metrics['accuracy']:.3f} ({metrics['correct']}/{metrics['total']}), " - f"cal_err={metrics['calibration_error']:.3f}") - - if metrics['accuracy'] > best_acc: - best_acc = metrics['accuracy'] + metrics = evaluate_weights( + evidence_dict, test_group, w_flip, w_conf, w_var, w_mp, w_cons + ) + + print( + f"{weights} -> acc={metrics['accuracy']:.3f} ({metrics['correct']}/{metrics['total']}), " + f"cal_err={metrics['calibration_error']:.3f}" + ) + + if metrics["accuracy"] > best_acc: + best_acc = metrics["accuracy"] best_weights = weights best_metrics = metrics -print(f"\nBest weights: {best_weights} -> acc={best_acc:.3f}, cal_err={best_metrics['calibration_error']:.3f}") +print( + f"\nBest weights: {best_weights} -> acc={best_acc:.3f}, cal_err={best_metrics['calibration_error']:.3f}" +) # %% [code] # Aggregate statistics with BEST weights @@ -510,11 +605,13 @@ for uid, scores in target_scores.items(): mean = np.mean(scores) std = np.std(scores) # Find ground truth - true_label = next(ex['vanilla_label'] for ex in test_group if ex['uid'] == uid) + true_label = next(ex["vanilla_label"] for ex in test_group if ex["uid"] == uid) predicted_label = 1 if mean > 0.5 else 0 correct = "✓" if predicted_label == true_label else "✗" uid_str = str(uid)[-6:] if isinstance(uid, (int, str)) else str(uid)[:6] - print(f"{uid_str}: mean={mean:.3f}, std={std:.3f}, pred={predicted_label}, true={true_label} {correct}") + print( + f"{uid_str}: mean={mean:.3f}, std={std:.3f}, pred={predicted_label}, true={true_label} {correct}" + ) # %% [code] # Save predictions as JSONL @@ -528,9 +625,11 @@ with open(output_path, "w") as f: "target_idx": pred.target_idx, "score": pred.score, "raw_logprob_diff": pred.raw_logprob_diff, - "context": [{"uid": uid, "label": int(lbl), "raw_logprob": raw_lp, "flipped": flip} - for uid, lbl, raw_lp, flip in pred.context], - "variations": pred.variations + "context": [ + {"uid": uid, "label": int(lbl), "raw_logprob": raw_lp, "flipped": flip} + for uid, lbl, raw_lp, flip in pred.context + ], + "variations": pred.variations, } f.write(json.dumps(record) + "\n") @@ -538,17 +637,17 @@ logger.info(f"Saved predictions to {output_path}") # %% [markdown] # ## Results Summary -# +# # **Grid Search Findings:** # - Most weight combinations achieve ~50% accuracy (random baseline) # - Best: (0.3, 0.2, 0.2, 0.2, 0.1) = 50% acc, cal_err=0.509 # - Heavy flip weighting (0.5, 0.1, 0.1, 0.2, 0.1) also 50% but higher cal_err -# +# # **Why low signal?** # 1. **Small sample**: 30 predictions on 10 examples = sparse evidence graph # 2. **Redundant sources**: All evidence ~0.62 conf, suggesting sources correlated # 3. **Missing global context**: Each prediction uses only group members (echo chamber) -# +# # **Next steps:** # 1. **Scale up**: 100+ predictions to densify evidence graph # 2. **Add global examples**: Mix in 2-3 random examples from other groups @@ -557,23 +656,23 @@ logger.info(f"Saved predictions to {output_path}") # %% [markdown] # ## Analysis & Next Steps -# +# # **Observations:** # - Evidence weights capture flip sensitivity—high weights = strong coupling between examples # - Ensemble variance reveals aleatoric uncertainty (some examples inherently ambiguous) # - Calibration: mean scores closer to 0/1 than raw logprobs (sigmoid normalization helps) -# +# # **Refinements Needed:** # 1. **Consistency weighting**: Currently hardcoded to 1, should check group rules (paraphrases agree, contradictions oppose) # 2. **Global examples**: Mix 2 random examples per group to prevent local echo chambers # 3. **Adaptive budget**: Spend more on high-variance targets # 4. **Directional evidence**: Track if flip improves or worsens global energy -# +# # **Theory Clarification (Epistemic vs Aleatoric):** # - **Epistemic**: Reducible via more/better context → measured by variance across ensemble (low var = model "knows") # - **Aleatoric**: Irreducible ambiguity in data → surfaces as consistency violations (e.g., paraphrases disagree) # - **Our method**: Ensemble variance ≈ epistemic, consistency failures ≈ aleatoric, evidence weights ≈ structural/relational confidence -# +# # **Concise Display Idea:** -# The `Prediction.__repr__` shows `P(uid | ctx1=lbl1*, ctx2=lbl2) = 0.75` where `*` marks flips. +# The `Prediction.__repr__` shows `P(uid | ctx1=lbl1*, ctx2=lbl2) = 0.75` where `*` marks flips. # Could extend to a class that auto-formats for logging/notebooks. diff --git a/src/simple_icm.py b/src/simple_icm.py index 87990ff..7808b99 100644 --- a/src/simple_icm.py +++ b/src/simple_icm.py @@ -129,14 +129,14 @@ def initialize_data(data, config): demonstrations = {item['uid']: deepcopy(item) for item in data} labeled_uids = random.sample(list(demonstrations.keys()), min(config.num_seed, len(data))) for uid in demonstrations: - demonstrations[uid]['label'] = None + demonstrations[uid]['pred_label'] = None demonstrations[uid]['score'] = 0.0 # Will store prediction score if uid in labeled_uids: - demonstrations[uid]['label'] = random.choice([0, 1]) + demonstrations[uid]['pred_label'] = random.choice([0, 1]) return demonstrations demonstrations = initialize_data(data, C) -logger.info("Initialized labels: {}", {k: v['label'] for k, v in demonstrations.items() if v['label'] is not None}) +logger.info("Initialized labels: {}", {k: v['pred_label'] for k, v in demonstrations.items() if v['pred_label'] is not None}) # %% [code] # Predict label using in-context prompting @@ -156,7 +156,7 @@ async def predict_label(example_uid, current_demos, config=C, verbose=False, all # Group and interleave demos by consistency_id like original grouped_demos = {} for uid, demo in current_demos.items(): - if uid != example_uid and demo['label'] is not None: + if uid != example_uid and demo['pred_label'] is not None: grouped_demos.setdefault(demo['consistency_id'], []).append(demo) relevant_demos = [] @@ -171,7 +171,7 @@ async def predict_label(example_uid, current_demos, config=C, verbose=False, all instruction += f"Hint: The Sets relate to the dimension: {C.semantic_anchor}\n\n" fewshot = [] for demo in relevant_demos: - label_str = "A" if demo['label'] == 1 else "B" + label_str = "A" if demo['pred_label'] == 1 else "B" fewshot.append(f"\n\n## Candidate:\n{demo['prompt']}\n## Set:\n{label_str}") # Use all_demos if provided (for unlabeled examples), otherwise use current_demos @@ -207,7 +207,7 @@ async def predict_label(example_uid, current_demos, config=C, verbose=False, all global reasoning_log # reasoning_log += f"\n\n## Candidate:\n{target_prompt}\n## Set:\n" #@ TODO record iter - labeled = [v for v in current_demos.values() if v['label'] is not None] + labeled = [v for v in current_demos.values() if v['pred_label'] is not None] reasoning_log += f""" Reasoning for UID {example_uid}, labelled {len(labeled)}: {response['choices'][0]['message']['content']}\n\n @@ -234,7 +234,7 @@ Reasoning for UID {example_uid}, labelled {len(labeled)}: # %% [code] # Compute energy and metrics def compute_energy(demos, config=C): - labeled = [d for d in demos.values() if d['label'] is not None] + labeled = [d for d in demos.values() if d['pred_label'] is not None] if not labeled: return 0.0 avg_lprob = np.mean([d['score'] for d in labeled]) @@ -245,11 +245,11 @@ def compute_energy(demos, config=C): num_inconsistent = 0 groups = {} for uid, demo in demos.items(): - if demo['label'] is not None: + if demo['pred_label'] is not None: cid = demo['consistency_id'] if cid not in groups: groups[cid] = [] - groups[cid].append((uid, demo['label'], demo['consistency_key'])) + groups[cid].append((uid, demo['pred_label'], demo['consistency_key'])) for cid, items in groups.items(): key_groups = {} @@ -269,7 +269,7 @@ def compute_energy(demos, config=C): num_inconsistent += max(0, len(items) - len(set(all_labels))) energy = config.alpha * avg_lprob - num_inconsistent - (num_inconsistent / max(1, len(labeled))) # Normalized penalty - accuracy = np.mean([d['label'] == d['vanilla_label'] for d in labeled]) + accuracy = np.mean([d['pred_label'] == d['vanilla_label'] for d in labeled]) # flip acc if needed, as this is unsupervised if accuracy < 0.5: accuracy = 1 - accuracy @@ -289,11 +289,11 @@ def get_kflip_neighbors(group_uids, demos, k): Generate all label assignments that are k flips away from current. Returns list of [(uid, new_label), ...] tuples. """ - labeled_uids = [uid for uid in group_uids if demos[uid]['label'] is not None] + labeled_uids = [uid for uid in group_uids if demos[uid]['pred_label'] is not None] neighbors = [] for combo in combinations(labeled_uids, k): - flips = [(uid, 1 - demos[uid]['label']) for uid in combo] + flips = [(uid, 1 - demos[uid]['pred_label']) for uid in combo] neighbors.append(flips) return neighbors @@ -307,7 +307,7 @@ async def fix_inconsistencies_greedy(demos, config=C, max_fixes=20, max_flips=3, # Find inconsistent groups groups = {} for uid, demo in demos.items(): - if demo['label'] is not None: + if demo['pred_label'] is not None: groups.setdefault(demo['consistency_id'], []).append(uid) # Find first inconsistent group @@ -332,7 +332,7 @@ async def fix_inconsistencies_greedy(demos, config=C, max_fixes=20, max_flips=3, # Apply flips temporarily temp_demos = deepcopy(demos) for uid, new_label in flips: - temp_demos[uid]['label'] = new_label + temp_demos[uid]['pred_label'] = new_label # Check if this is consistent if not is_consistent(inconsistent_group, temp_demos): @@ -351,7 +351,7 @@ async def fix_inconsistencies_greedy(demos, config=C, max_fixes=20, max_flips=3, # Apply best flips if improvement found if best_flips and best_energy > old_energy: for uid, new_label in best_flips: - demos[uid]['label'] = new_label + demos[uid]['pred_label'] = new_label else: break # No improvement possible, stop trying @@ -366,7 +366,7 @@ async def run_icm(demonstrations, config=C): # Fix any initial inconsistencies from random initialization demonstrations = await fix_inconsistencies_greedy(demonstrations, config) - current_labeled = {k: v for k, v in demonstrations.items() if v['label'] is not None} + current_labeled = {k: v for k, v in demonstrations.items() if v['pred_label'] is not None} old_energy, old_metrics = compute_energy(demonstrations, config) try: @@ -387,7 +387,7 @@ async def run_icm(demonstrations, config=C): weights = [0.1 for _ in all_uids] # Base low for cid, group_uids in groups.items(): - labeled_labels = [demonstrations[uid]['label'] for uid in group_uids if demonstrations[uid]['label'] is not None] + labeled_labels = [demonstrations[uid]['pred_label'] for uid in group_uids if demonstrations[uid]['pred_label'] is not None] num_labeled = len(labeled_labels) num_unlabeled = len(group_uids) - num_labeled @@ -399,7 +399,7 @@ async def run_icm(demonstrations, config=C): if num_unlabeled > 0: weight_factor = (0.5 + 0.5 * inconsistency) * (1 + num_unlabeled / len(group_uids)) for uid in group_uids: - if demonstrations[uid]['label'] is None: # Unlabeled + if demonstrations[uid]['pred_label'] is None: # Unlabeled idx = all_uids.index(uid) weights[idx] = weight_factor else: @@ -432,11 +432,11 @@ async def run_icm(demonstrations, config=C): for uid, (new_label, score) in zip(candidate_uids, results): temp_demos = deepcopy(demonstrations) - temp_demos[uid]['label'] = new_label + temp_demos[uid]['pred_label'] = new_label temp_demos[uid]['score'] = score # Fix inconsistencies if label changed - if demonstrations[uid]['label'] != new_label: + if demonstrations[uid]['pred_label'] != new_label: temp_demos = await fix_inconsistencies_greedy(temp_demos, config) new_energy, _ = compute_energy(temp_demos, config) @@ -460,7 +460,7 @@ async def run_icm(demonstrations, config=C): if delta > 0 or random.random() < math.exp(delta / T): demonstrations = best_temp_demos old_energy = new_energy - current_labeled = {k: v for k, v in demonstrations.items() if v['label'] is not None} + current_labeled = {k: v for k, v in demonstrations.items() if v['pred_label'] is not None} logger.debug("Iter {}: Accepted UID {}. Energy: {:.2f}. {}", iter, best_uid, old_energy, accept_msg) else: logger.debug("Iter {}: Rejected. {}", iter, accept_msg) @@ -502,14 +502,14 @@ logger.info("Inconsistencies: {}", final_metrics['num_inconsistent']) df = pd.DataFrame(final_demos).T df.to_parquet(out_dir / "icm_final_labels.parquet") -df_labeled = df.dropna(subset='label').sort_values(by='score', key=np.abs, ascending=False) -df_labeled_disagreed = df_labeled[df_labeled['vanilla_label'] != df_labeled['label']] +df_labeled = df.dropna(subset='pred_label').sort_values(by='score', key=np.abs, ascending=False) +df_labeled_disagreed = df_labeled[df_labeled['vanilla_label'] != df_labeled['pred_label']] print(f"\nFinal labeled examples (total {len(df_labeled)}):") -print(df_labeled_disagreed[['consistency_id', 'label', 'vanilla_label', 'score', 'prompt']]) +print(df_labeled_disagreed[['consistency_id', 'pred_label', 'vanilla_label', 'score', 'prompt']]) for uid, row in df_labeled_disagreed.iterrows(): - print(f"\n## Candidate: {row['prompt']}\nICM Set: {'A' if row['label']==1 else 'B'}, Vanilla Set: {'A' if row['vanilla_label']==1 else 'B'}, score={row['score']}\n") + print(f"\n## Candidate: {row['prompt']}\nICM Set: {'A' if row['pred_label']==1 else 'B'}, Vanilla Set: {'A' if row['vanilla_label']==1 else 'B'}, score={row['score']}\n") print(f"\nFinal labeled examples saved to {out_dir / 'icm_final_labels.parquet'}")