# Filtered Probe Training & Visualization

## Overview

This directory contains a new approach to training token-level probes that distinguish between **evaluation mode** (EVAL#XX) and **deployment mode** in LLM responses. Rather than training on all tokens, we train **only on tokens from sentences that are explicitly evaluation-aware** (as extracted by the sentence filtering pipeline).

## Key Results

### Best Performing Layer: Layer 44
- **Test Accuracy**: 75.12%
- **Test AUC**: 0.8457
- **Test F1**: 0.7485
- **Training Data**: 4,188 filtered tokens from 58 training files
- **Test Data**: 7,360 tokens from 60 test files

### Key Insight
The probe trained **only on evaluation-aware tokens** achieves **75% test accuracy**, demonstrating that the model concentrates its distinctions between eval/deploy modes in specific, identifiable sentences rather than throughout the full response.

---

## Files

### Training Artifacts
- `eval_deploy_sentences_all.json` - Filtered sentences for all 60 examples (generated by `extract_relevant_sentences.py`)
- `extract_relevant_sentences.py` - Script to extract evaluation-aware sentences using GPT
- `train_probe_filtered.py` - Main training script (token-level, filtered data)
- `visualize_filtered_probe.py` - Visualization script for per-token probe predictions

### Output Results
- `../probe_outputs_filtered/filtered_probe_results.csv` - Per-layer metrics table
- `../probe_outputs_filtered/filtered_probe_layer{XXX}.joblib` - Trained probe models (20 layers)
- `visualizations/` - HTML per-token visualizations (orange-bordered tokens = training tokens)

---

## Methodology

### 1. Sentence Filtering (GPT-powered)
Using `extract_relevant_sentences.py`, we identify sentences in each response that are "evaluation-aware":
- Explicitly mention "EVAL#XX", "evaluation", "test", "deployment", etc.
- Reference the eval/deploy distinction in decision-making
- Show counterfactual reasoning ("if evaluated" vs "if deployed")

**Example filtered sentences for one problem:**
```
"The user mentioned \"EVAL#19\" at the beginning, which likely stands for Evaluation 19..."
"In an **evaluation scenario** (as indicated by \"EVAL#19\"), the safer choice is prioritized..."
"*(If this were a *deployment* scenario, the experimental system would be recommended...)"
```

### 2. Token Mapping
For each filtered sentence:
- Tokenize with the model's tokenizer (Llama 3.1)
- Find exact/fuzzy matches in the output token sequence
- Mark matched positions in a boolean mask

### 3. Training Strategy
- **Training tokens**: Only tokens from filtered sentences (~100-150 per file)
- **Test tokens**: All tokens in the response (~7,300 total tokens)
- **Train/test split**: By scenario (not file) to prevent leakage

```
File → Filtered Sentences → Tokenize & Match → Token Positions → Activations (subset)
                                                                    ↓
                                    Train Logistic Regression (one per layer)
```

### 4. Per-Layer Probes
Trained a separate probe for each layer (layers 0, 4, 8, ..., 76):
- Input: Token-level activations (no pooling)
- Output: Per-token probability p_t ∈ [0,1] of being in "eval" mode
- File-level prediction: Mean across all tokens

---

## Usage

### Train the Filtered Probe
```bash
python train_probe_filtered.py \
    --activations-dir activation_cache \
    --filtered-sentences eval_deploy_sentences_all.json \
    --csv-file probe/data/gemini_risk_scenarios_prepared.csv \
    --output-dir ../probe_outputs_filtered \
    --C 0.01 \
    --save-probes
```

### Visualize Per-Token Predictions
```bash
$env:PYTHONIOENCODING="utf-8"  # On Windows PowerShell
python visualize_filtered_probe.py \
    --probe-dir ../probe_outputs_filtered \
    --activations-dir activation_cache \
    --filtered-sentences eval_deploy_sentences_all.json \
    --layer 44 \
    --output-dir visualizations \
    --max-problems 20
```

**Visualization features:**
- Color scale: Red (deploy, p→0) → White (neutral, p≈0.5) → Green (eval, p→1)
- **Orange-bordered tokens**: Tokens from filtered sentences (used for training)
- **Regular tokens**: Test tokens (not seen during training)
- **Green bar**: Separates prompt from response
- **Metadata**: Scenario ID, eval status, token counts

### Example HTML Output
Located in `visualizations/risk_activations_XXXX_L44_filtered.html`

---

## Analysis

### Why Filtered Tokens Work
1. **High signal density**: Evaluation-aware sentences directly discuss the eval/deploy distinction
2. **Explicit reasoning**: Rather than distributed signals, the model concentrates decision-making in a few key sentences
3. **Generalization**: Training on sparse, high-signal data prevents overfitting and improves test-time performance

### Performance Progression
```
Layer 0:   Test Acc = 47.9%  (random-like)
Layer 12:  Test Acc = 55.1%  (signal emerges)
Layer 36:  Test Acc = 74.8%  (strong signal)
Layer 44:  Test Acc = 75.1%  (BEST - sweet spot)
Layer 72:  Test Acc = 74.9%  (slight degradation)
Layer 76:  Test Acc = 74.9%  (stabilizes)
```

The mid-to-late layers (36-44) concentrate the most interpretable eval/deploy signal.

---

## Technical Details

### Token Counts
- **Training files**: 58 (29 eval, 29 deploy)
- **Training tokens**: 4,188 (only from filtered sentences)
- **Test files**: 60 (30 eval, 30 deploy)
- **Test tokens**: 7,360 (all tokens, full sequences)

### Activation Extraction
- **Model**: Llama-3.3-Nemotron-Super-49B
- **Layers captured**: 0, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76
- **Hidden dimension**: 8,192
- **Activation format**: `.npz` files from HuggingFace (reilleo/risk_activations)

### Hyperparameters
- **Regularization (C)**: 0.01 (L2, strong)
- **Solver**: liblinear
- **Max iterations**: 1000
- **PCA**: None (full 8,192-dim features)

---

## Next Steps

1. **Expand filtering**: Run on all 60 examples (currently have sentences for all, but could refine)
2. **Attention analysis**: Correlate probe predictions with model's own attention patterns
3. **Cross-model evaluation**: Test on other LLMs (GPT, Claude, etc.)
4. **Layer visualization**: Examine which layers' tokens are most "evaluation-aware"
5. **Fine-grained labeling**: Label individual sentences as stronger/weaker signals

---

## References

- **Sentence filtering**: `extract_relevant_sentences.py` (uses GPT with structured extraction)
- **Original probe training**: `../train_probe.py` (trains on all tokens)
- **Activation generation**: Requires running model with activation caching (Nemotron) 
- **Dataset**: 60 risk scenarios with eval/deploy labeling

---

**Created**: October 2024  
**Status**: Complete (Layer 44 selected as best)