{ "nbformat": 4, "nbformat_minor": 5, "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0" } }, "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Notebook 1: Find the Direction", "", "**AI Alignment \u2014 The Multiverse School**", "", "Every concept a language model has learned lives somewhere in its activation space as a *direction* \u2014 a vector. Honesty, helpfulness, harm avoidance \u2014 they're not single neurons, they're geometric features spread across thousands of dimensions.", "", "Today you will **find one of those directions**, extract it, and prove it's real.", "", "**What you need:** A RunPod GPU pod (setup instructions below)" ], "id": "c001" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 0 \u2014 RunPod Setup", "", "You're running this on your own GPU, not Colab. Here's how to get there.", "", "### Create your pod", "", "1. Go to [runpod.io](https://runpod.io) and sign up (or log in)", "2. Add credits ($5-10 is plenty for today \u2014 a 3B model on an RTX 4090 costs ~$0.44/hr)", "3. Click **Pods** \u2192 **+ GPU Pod**", "4. Pick a GPU \u2014 any of these work for Qwen-2.5-3B:", " - **RTX 4090** (24 GB, ~$0.44/hr) \u2014 fast, cheap, recommended", " - **RTX 3090** (24 GB, ~$0.22/hr on community cloud) \u2014 cheapest option", " - **A40** (48 GB, ~$0.39/hr) \u2014 overkill but fine", "5. Template: **RunPod Pytorch 2.4.0** (it comes with torch + CUDA pre-installed)", "6. Set container disk to **20 GB** (model weights are ~6 GB)", "7. Click **Deploy**", "", "### Open Jupyter", "", "Once your pod is running:", "1. Click **Connect** \u2192 **Jupyter Lab**", "2. Upload this notebook (drag it into the file browser)", "3. Open it and start running cells", "", "### Or use the terminal", "", "You can also SSH in or use the web terminal:", "```", "ssh root@ -p -i ~/.ssh/id_ed25519", "```", "RunPod shows you the SSH command on the pod page." ], "id": "c002" }, { "cell_type": "code", "metadata": {}, "source": [ "# Check your GPU \u2014 you should see your RunPod GPU here", "import torch", "print(f\"PyTorch: {torch.__version__}\")", "print(f\"CUDA available: {torch.cuda.is_available()}\")", "if torch.cuda.is_available():", " print(f\"GPU: {torch.cuda.get_device_name(0)}\")", " print(f\"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB\")", "else:", " print(\"WARNING: No GPU detected. Check your RunPod pod settings.\")", "", "# Install what's not in the RunPod template", "import subprocess", "subprocess.check_call([\"pip\", \"install\", \"-q\", \"transformers\", \"accelerate\",", " \"scikit-learn\", \"matplotlib\"])", "print(\"\\nDependencies installed.\")" ], "id": "c003", "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 1 \u2014 Load the Model", "", "We're using **Qwen2.5-3B-Instruct** \u2014 a 3-billion-parameter instruction-tuned model.", "In float16 it's about 6 GB, which fits on any GPU you picked above.", "", "When the model loads, look at the numbers it prints. You're about to reach inside", "those layers and find what the model uses to represent honesty." ], "id": "c004" }, { "cell_type": "code", "metadata": {}, "source": [ "import torch", "from transformers import AutoModelForCausalLM, AutoTokenizer", "", "MODEL_NAME = \"Qwen/Qwen2.5-3B-Instruct\"", "", "print(\"Loading tokenizer...\")", "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)", "", "print(\"Loading model (this takes 1-2 minutes the first time, cached after)...\")", "model = AutoModelForCausalLM.from_pretrained(", " MODEL_NAME,", " torch_dtype=torch.float16,", " device_map=\"auto\",", ")", "model.eval()", "", "print(f\"\\n--- Model loaded ---\")", "print(f\"Parameters: {sum(p.numel() for p in model.parameters()) / 1e9:.1f} billion\")", "print(f\"Layers: {model.config.num_hidden_layers}\")", "print(f\"Hidden size: {model.config.hidden_size}\")", "print(f\"Device: {model.device}\")", "print(f\"GPU memory used: {torch.cuda.memory_allocated() / 1e9:.1f} GB\")", "print(f\"\\nEvery layer produces a {model.config.hidden_size}-dimensional vector.\")", "print(f\"The concept of 'honesty' is a direction in that {model.config.hidden_size}-D space.\")", "print(f\"We're going to find it.\")" ], "id": "c005", "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 2 \u2014 What's Inside a Layer?", "", "Before we go looking for honesty, let's see what a hidden state actually looks like.", "We'll feed the model a simple prompt and look at the raw activations." ], "id": "c006" }, { "cell_type": "code", "metadata": {}, "source": [ "# Feed the model a prompt and capture ALL hidden states from ALL layers", "test_input = \"Is this claim actually true?\"", "inputs = tokenizer(test_input, return_tensors=\"pt\").to(model.device)", "", "with torch.no_grad():", " outputs = model(**inputs, output_hidden_states=True)", "", "# outputs.hidden_states is a tuple: one tensor per layer (+ the embedding layer)", "print(f\"Number of hidden state layers: {len(outputs.hidden_states)}\")", "print(f\"Shape of each: {outputs.hidden_states[0].shape}\")", "print(f\" \u2192 [batch_size=1, sequence_length={inputs['input_ids'].shape[1]}, hidden_dim={model.config.hidden_size}]\")", "print(f\"\\nThe last token's hidden state at the middle layer:\")", "", "mid_layer = model.config.num_hidden_layers // 2", "last_token_state = outputs.hidden_states[mid_layer][0, -1, :]", "print(f\" Shape: {last_token_state.shape}\")", "print(f\" First 10 values: {last_token_state[:10].float().cpu().tolist()}\")", "print(f\"\\nThat vector IS what the model is 'thinking' at this point.\")", "print(f\"Somewhere in those {model.config.hidden_size} dimensions is the direction for honesty.\")" ], "id": "c007", "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3 \u2014 Contrastive Pairs", "", "To find the honesty direction, we need to show the model the **same scenario**", "framed two ways: one that activates honesty, one that activates sycophancy", "(telling people what they want to hear).", "", "The difference in hidden states between these two framings \u2014 averaged across", "many examples \u2014 IS the honesty direction.", "", "**Read through these pairs.** Each one is the same question, but the system", "prompt tells the model to be either honest or people-pleasing." ], "id": "c008" }, { "cell_type": "code", "metadata": {}, "source": [ "# 20 matched pairs: same user question, different system framing", "# The model's hidden states will differ \u2014 and that difference is what we're after.", "", "SCENARIOS = [", " \"Is this too-good-to-be-true investment actually safe?\",", " \"What are the real side effects of this medication?\",", " \"Is this online review trustworthy?\",", " \"Does this weight loss supplement really work?\",", " \"Is my data actually private on this platform?\",", " \"Will this coding bootcamp really get me a job?\",", " \"Is this charity using donations effectively?\",", " \"Are these lab results actually concerning?\",", " \"Is this contractor giving me a fair quote?\",", " \"Does this academic paper actually prove its claim?\",", " \"Is this news article reporting facts accurately?\",", " \"Will this skincare product actually reduce wrinkles?\",", " \"Is this used car in good condition?\",", " \"Is my password really secure enough?\",", " \"Does this insurance policy actually cover what I need?\",", " \"Is this organic label meaningful or just marketing?\",", " \"Will this online course teach me real skills?\",", " \"Is this financial advisor acting in my interest?\",", " \"Are these product reviews genuine?\",", " \"Is this job offer legitimate?\",", "]", "", "HONEST_SYS = \"You are an honest, straightforward assistant.\"", "SYCOPHANTIC_SYS = \"You are a people-pleasing assistant who always tells users what they want to hear.\"", "", "honest_prompts = [f\"{HONEST_SYS} User asks: {s}\" for s in SCENARIOS]", "sycophantic_prompts = [f\"{SYCOPHANTIC_SYS} User asks: {s}\" for s in SCENARIOS]", "", "print(f\"Created {len(honest_prompts)} honest and {len(sycophantic_prompts)} sycophantic prompts.\")", "print(f\"\\nExample honest: {honest_prompts[0][:80]}...\")", "print(f\"Example sycophantic: {sycophantic_prompts[0][:80]}...\")" ], "id": "c009", "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 4 \u2014 Extract Hidden States", "", "Now we feed every prompt through the model and capture the hidden state", "at the **middle layer**, **last token position**.", "", "Why middle layer? Early layers handle syntax. Late layers handle output", "formatting. The middle layers are where *concepts* live \u2014 honesty, deception,", "helpfulness \u2014 as geometric features." ], "id": "c010" }, { "cell_type": "code", "metadata": {}, "source": [ "import numpy as np", "", "def get_hidden_states(prompts, model, tokenizer, layer_idx):", " \"\"\"Extract the hidden state at [layer_idx, last_token] for each prompt.\"\"\"", " states = []", " for i, prompt in enumerate(prompts):", " inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)", " with torch.no_grad():", " out = model(**inputs, output_hidden_states=True)", " hs = out.hidden_states[layer_idx][:, -1, :].float().cpu().numpy()", " states.append(hs[0])", " if (i + 1) % 5 == 0:", " print(f\" {i+1}/{len(prompts)}\")", " return np.array(states)", "", "mid_layer = model.config.num_hidden_layers // 2", "print(f\"Extracting from layer {mid_layer} (of {model.config.num_hidden_layers})...\")", "", "print(\"\\nHonest prompts:\")", "honest_states = get_hidden_states(honest_prompts, model, tokenizer, mid_layer)", "", "print(\"\\nSycophantic prompts:\")", "syco_states = get_hidden_states(sycophantic_prompts, model, tokenizer, mid_layer)", "", "print(f\"\\nHonest states shape: {honest_states.shape}\")", "print(f\"Sycophantic states shape: {syco_states.shape}\")", "print(f\"\\nEach row is a {model.config.hidden_size}-dimensional snapshot of what the model\")", "print(f\"was 'thinking' after reading that prompt. 40 snapshots total: 20 honest, 20 sycophantic.\")" ], "id": "c011", "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 5 \u2014 Train the Probe", "", "A **linear probe** is a simple classifier (logistic regression) trained on", "the hidden states. If it can tell honest from sycophantic states with high", "accuracy, that proves the model has a **linearly separable** internal", "representation of honesty.", "", "In plain terms: the model already knows the difference. We're just finding", "where it stores that knowledge." ], "id": "c012" }, { "cell_type": "code", "metadata": {}, "source": [ "from sklearn.linear_model import LogisticRegression", "from sklearn.model_selection import cross_val_score", "", "# Stack all states and label them", "X = np.concatenate([honest_states, syco_states])", "y = np.array([1] * len(honest_states) + [0] * len(syco_states))", "# 1 = honest, 0 = sycophantic", "", "# Train with 5-fold cross-validation", "probe = LogisticRegression(max_iter=1000)", "scores = cross_val_score(probe, X, y, cv=5)", "", "# Fit on all data (we'll use the weights as our direction vector)", "probe.fit(X, y)", "", "print(f\"Probe accuracy (5-fold CV): {scores.mean():.1%} \u00b1 {scores.std():.1%}\")", "print(f\"\\nWhat this means:\")", "if scores.mean() > 0.9:", " print(f\" The model's hidden states are {scores.mean():.0%} separable on honesty vs sycophancy.\")", " print(f\" The model HAS an internal representation of this concept.\")", " print(f\" It's not just a prompt trick \u2014 it's geometric.\")", "elif scores.mean() > 0.7:", " print(f\" Partially separable \u2014 the signal exists but it's noisy.\")", "else:", " print(f\" Weak separation \u2014 try a different layer or more contrastive pairs.\")" ], "id": "c013", "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 6 \u2014 Visualize the Separation", "", "Let's project those 2048-dimensional vectors down to 2D with PCA and actually", "*see* the honesty direction." ], "id": "c014" }, { "cell_type": "code", "metadata": {}, "source": [ "import matplotlib", "matplotlib.use('Agg')", "import matplotlib.pyplot as plt", "from sklearn.decomposition import PCA", "", "pca = PCA(n_components=2)", "X_2d = pca.fit_transform(X)", "", "fig, ax = plt.subplots(1, 1, figsize=(8, 6))", "ax.scatter(X_2d[:len(honest_states), 0], X_2d[:len(honest_states), 1],", " c='#2563a8', label='Honest', alpha=0.8, s=60, edgecolors='white', linewidth=0.5)", "ax.scatter(X_2d[len(honest_states):, 0], X_2d[len(honest_states):, 1],", " c='#dc2626', label='Sycophantic', alpha=0.8, s=60, edgecolors='white', linewidth=0.5)", "", "# Draw the separating line", "w = pca.transform(probe.coef_)[0]", "b = probe.intercept_[0]", "xx = np.linspace(X_2d[:, 0].min() - 1, X_2d[:, 0].max() + 1, 100)", "yy = -(w[0] * xx + b) / (w[1] + 1e-10)", "mask = (yy > X_2d[:, 1].min() - 2) & (yy < X_2d[:, 1].max() + 2)", "ax.plot(xx[mask], yy[mask], 'k--', alpha=0.4, linewidth=1, label='Probe boundary')", "", "ax.set_xlabel(f\"PC1 ({pca.explained_variance_ratio_[0]:.1%} variance)\")", "ax.set_ylabel(f\"PC2 ({pca.explained_variance_ratio_[1]:.1%} variance)\")", "ax.set_title(\"Hidden States: Honest vs Sycophantic\\n(PCA of layer \" + str(mid_layer) + \")\")", "ax.legend()", "plt.tight_layout()", "plt.savefig(\"probe_separation.png\", dpi=150)", "plt.show()", "", "print(\"\\nBlue dots = honest framing. Red dots = sycophantic framing.\")", "print(\"The dashed line is where the linear probe separates them.\")", "print(\"The direction PERPENDICULAR to that line is the honesty vector.\")", "print(\"\\nSaved to probe_separation.png\")" ], "id": "c015", "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 7 \u2014 Extract the Steering Vector", "", "The probe's weight vector **is** the honesty direction. It's the line in", "2048-dimensional space along which honest and sycophantic activations differ", "most. We normalize it to unit length so we can scale it later.", "", "**This is what you'll use in Notebook 2 to steer the model.**" ], "id": "c016" }, { "cell_type": "code", "metadata": {}, "source": [ "# The probe's coefficient IS the concept direction", "honesty_direction = probe.coef_[0].copy()", "honesty_direction /= np.linalg.norm(honesty_direction)", "", "print(f\"Steering vector extracted.\")", "print(f\" Shape: {honesty_direction.shape}\")", "print(f\" Norm: {np.linalg.norm(honesty_direction):.4f} (unit vector)\")", "print(f\" Layer: {mid_layer}\")", "print(f\"\")", "print(f\"This vector IS what the model uses to represent honesty vs sycophancy.\")", "print(f\"Adding it to the model's activations will push it toward honesty.\")", "print(f\"Subtracting it will push toward sycophancy.\")", "", "# Save for Notebook 2", "np.save(\"honesty_direction.npy\", honesty_direction)", "print(f\"\\nSaved to honesty_direction.npy \u2014 you'll load this in Notebook 2.\")" ], "id": "c017", "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 8 \u2014 Quick Steering Preview", "", "Before we move to Notebook 2, let's do a quick sanity check: does adding this", "vector actually change the model's output?" ], "id": "c018" }, { "cell_type": "code", "metadata": {}, "source": [ "def make_hook(vector_tensor, alpha=1.0):", " \"\"\"Forward hook that adds a steering vector to hidden states.\"\"\"", " def hook_fn(module, input, output):", " if isinstance(output, tuple):", " return (output[0] + alpha * vector_tensor,) + output[1:]", " return output + alpha * vector_tensor", " return hook_fn", "", "vec_tensor = torch.tensor(honesty_direction, dtype=torch.float16).to(model.device)", "", "def quick_generate(prompt, alpha=0.0):", " messages = [{\"role\": \"user\", \"content\": prompt}]", " text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)", " inputs = tokenizer(text, return_tensors=\"pt\").to(model.device)", " handle = None", " if alpha != 0:", " handle = model.model.layers[mid_layer].register_forward_hook(make_hook(vec_tensor, alpha))", " with torch.no_grad():", " out = model.generate(**inputs, max_new_tokens=100, temperature=0.7,", " do_sample=True, pad_token_id=tokenizer.eos_token_id)", " if handle:", " handle.remove()", " return tokenizer.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True).strip()", "", "prompt = \"My friend's painting is terrible but they asked what I think. What should I say?\"", "print(f\"Prompt: {prompt}\\n\")", "", "for alpha, label in [(0, \"BASELINE\"), (6, \"STEERED HONEST (+6)\"), (-6, \"STEERED SYCOPHANTIC (-6)\")]:", " print(f\"[{label}]\")", " print(quick_generate(prompt, alpha))", " print()" ], "id": "c019", "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## What You Just Did", "", "1. **Provisioned your own GPU** on RunPod \u2014 not Colab, not a managed notebook. Your machine.", "2. **Loaded a 3-billion-parameter model** and looked at its internal structure", "3. **Created contrastive pairs** \u2014 the same questions framed as honest vs sycophantic", "4. **Extracted hidden states** \u2014 the model's internal representations at the middle layer", "5. **Trained a linear probe** \u2014 proved the model has a linearly separable concept of honesty", "6. **Visualized the separation** \u2014 saw the honesty direction in 2D", "7. **Extracted the steering vector** \u2014 the unit-length direction in activation space", "8. **Previewed steering** \u2014 saw the model's behavior change when you add/subtract the vector", "", "**The key insight:** honesty isn't stored in a single neuron. It's a *direction*", "across 2048 dimensions. And because it's a direction, you can move along it \u2014", "making the model more honest, or more sycophantic, by adding or subtracting the vector.", "", "**Next: Notebook 2** \u2014 systematic steering experiments, adversarial red-teaming, and", "deploying this as an API on your own backend." ], "id": "c020" } ] }