LLM-Bias-And-Activation.ipynb

LLM-Bias-And-Activation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "**Understanding AI activation tuning**\n",
        "\n",
        "This document is intended to demonstrate and teach how tools such as abliteration and deviant work to turn off specific activation vectors. We will do the opposite here by making the AI obsess over a specific subject.\n",
        "\n",
        "Note that this does not change the model weights. It is a runtime calculation that has the affect stearing the output before the model makes its final decisions.\n",
        "\n",
        "Run each of the cells one at a time. As you scroll down read the explaination and comments in the code to understand what is happening."
      ],
      "metadata": {
        "id": "A-EvU9ZK9Ai4"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "First import the required modules. Click in the box below and click the \"play\" button."
      ],
      "metadata": {
        "id": "y35Y6zkWhr33"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from __future__ import annotations\n",
        "\n",
        "import argparse\n",
        "import textwrap\n",
        "from dataclasses import dataclass\n",
        "from typing import Iterable\n",
        "\n",
        "import torch\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n"
      ],
      "metadata": {
        "id": "8DCEZvEc80Jv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "This next block defines a bunch of supporting functions."
      ],
      "metadata": {
        "id": "CX0U65jqc8uT"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7DUb6ciN80wF"
      },
      "outputs": [],
      "source": [
        "@dataclass(frozen=True)\n",
        "class ContrastivePair:\n",
        "    positive: str\n",
        "    negative: str\n",
        "\n",
        "\n",
        "def build_contrastive_pairs(subject: str) -> list[ContrastivePair]:\n",
        "    \"\"\"Create matched prompts that contrast the target subject with unrelated topics.\n",
        "\n",
        "    The positive prompts mention the subject the user wants to steer toward. The\n",
        "    negative prompts use similar wording but talk about unrelated everyday items.\n",
        "    Comparing model activations for these two groups gives us a direction that is\n",
        "    more associated with the target subject.\n",
        "    \"\"\"\n",
        "    # Each tuple keeps the positive and negative wording similar so the subject is\n",
        "    # the main difference the model sees.\n",
        "    templates = [\n",
        "        (\"Explain the key ideas in {s}.\", \"Explain the key ideas in office furniture.\"),\n",
        "        (\"List important vocabulary for {s}.\", \"List important vocabulary for baking bread.\"),\n",
        "        (\"Write a short lesson about {s}.\", \"Write a short lesson about traffic lights.\"),\n",
        "        (\"Give practical examples involving {s}.\", \"Give practical examples involving umbrellas.\"),\n",
        "        (\"Describe why {s} matters.\", \"Describe why paper clips matter.\"),\n",
        "        (\"Create a beginner quiz about {s}.\", \"Create a beginner quiz about window curtains.\"),\n",
        "        (\"Summarize recent questions people ask about {s}.\", \"Summarize common questions about shoelaces.\"),\n",
        "        (\"Name five subtopics within {s}.\", \"Name five subtopics within grocery bags.\"),\n",
        "        (\"Compare different perspectives on {s}.\", \"Compare different perspectives on doormats.\"),\n",
        "        (\"Give an analogy that teaches {s}.\", \"Give an analogy that teaches elevator buttons.\"),\n",
        "        (\"Write a paragraph centered entirely on {s}.\", \"Write a paragraph centered entirely on cardboard boxes.\"),\n",
        "        (\"What should a student know before learning {s}?\", \"What should a student know before folding napkins?\"),\n",
        "        (\"Mention tools or methods used in {s}.\", \"Mention tools or methods used in dish washing.\"),\n",
        "        (\"Give a historical note about {s}.\", \"Give a historical note about coat hangers.\"),\n",
        "        (\"Explain a common misconception about {s}.\", \"Explain a common misconception about desk lamps.\"),\n",
        "        (\"Write three facts about {s}.\", \"Write three facts about parking meters.\"),\n",
        "        (\"Describe a problem solved by {s}.\", \"Describe a problem solved by rubber bands.\"),\n",
        "        (\"Give a concise definition of {s}.\", \"Give a concise definition of bottle caps.\"),\n",
        "        (\"Create a study plan for {s}.\", \"Create a study plan for sorting laundry.\"),\n",
        "        (\"Explain how experts think about {s}.\", \"Explain how experts think about cafeteria trays.\"),\n",
        "    ]\n",
        "    return [ContrastivePair(pos.format(s=subject), neg) for pos, neg in templates]\n",
        "\n",
        "\n",
        "def mean_last_token_hidden(\n",
        "    model: AutoModelForCausalLM,\n",
        "    tokenizer: AutoTokenizer,\n",
        "    prompts: Iterable[str],\n",
        ") -> torch.Tensor:\n",
        "    \"\"\"Measure the model's hidden states at the end of each prompt.\n",
        "\n",
        "    The tokenizer turns text into numeric token IDs. The model then processes those\n",
        "    token IDs and returns internal activations for every layer. This function keeps\n",
        "    only the activation at the last real token of each prompt because that position\n",
        "    has seen the whole prompt so far.\n",
        "    \"\"\"\n",
        "    # Convert text prompts into PyTorch tensors, padding shorter prompts so the\n",
        "    # whole list can be processed as one batch.\n",
        "    encoded = tokenizer(list(prompts), return_tensors=\"pt\", padding=True, truncation=True)\n",
        "    # Move every input tensor to the same device as the model, such as CPU or GPU.\n",
        "    encoded = {key: value.to(model.device) for key, value in encoded.items()}\n",
        "\n",
        "    # We are only inspecting activations, not training, so gradient tracking is unnecessary.\n",
        "    with torch.no_grad():\n",
        "        outputs = model(**encoded, output_hidden_states=True, use_cache=False)\n",
        "\n",
        "    # The attention mask marks real tokens with 1 and padding tokens with 0.\n",
        "    attention_mask = encoded[\"attention_mask\"]\n",
        "    # Count real tokens in each prompt, then subtract 1 to get the last real token index.\n",
        "    last_indices = attention_mask.sum(dim=1) - 1\n",
        "    layer_values = []\n",
        "\n",
        "    for hidden in outputs.hidden_states:\n",
        "        # Select each prompt's hidden state at its own last real token.\n",
        "        batch_indices = torch.arange(hidden.shape[0], device=hidden.device)\n",
        "        layer_values.append(hidden[batch_indices, last_indices, :])\n",
        "\n",
        "    # Shape: [number of layers, number of prompts, hidden dimensions].\n",
        "    return torch.stack(layer_values)\n",
        "\n",
        "\n",
        "def choose_layer_and_vector(\n",
        "    positive_hidden: torch.Tensor,\n",
        "    negative_hidden: torch.Tensor,\n",
        ") -> tuple[int, torch.Tensor, torch.Tensor]:\n",
        "    \"\"\"Find the layer where positive and negative prompts differ most.\n",
        "\n",
        "    For each layer, this computes the average activation for the positive prompts\n",
        "    and subtracts the average activation for the negative prompts. The largest\n",
        "    difference becomes the steering direction.\n",
        "    \"\"\"\n",
        "    # Average across all positive examples and all negative examples separately.\n",
        "    positive_mean = positive_hidden.mean(dim=1)\n",
        "    negative_mean = negative_hidden.mean(dim=1)\n",
        "    # The delta points from unrelated-topic activations toward subject-topic activations.\n",
        "    layer_deltas = positive_mean - negative_mean\n",
        "    # A larger norm means that layer separated the two prompt groups more strongly.\n",
        "    layer_norms = layer_deltas.norm(dim=1)\n",
        "\n",
        "    # Skip embedding layer 0 when possible; steering transformer blocks is clearer.\n",
        "    start_layer = 1 if layer_norms.numel() > 1 else 0\n",
        "    selected_layer = int(torch.argmax(layer_norms[start_layer:]).item() + start_layer)\n",
        "    steering_vector = layer_deltas[selected_layer]\n",
        "    return selected_layer, steering_vector, layer_norms\n",
        "\n",
        "\n",
        "def top_activation_dimensions(vector: torch.Tensor, count: int = 8) -> list[tuple[int, float]]:\n",
        "    \"\"\"Show which hidden dimensions contribute most to the steering vector.\n",
        "\n",
        "    This is only for display. It helps students see that the steering vector is a\n",
        "    list of numbers, and some dimensions have larger positive or negative values\n",
        "    than others.\n",
        "    \"\"\"\n",
        "    # Use absolute value for ranking so strong negative and strong positive values both appear.\n",
        "    values, indices = torch.topk(vector.abs(), k=min(count, vector.numel()))\n",
        "    return [(int(index), float(vector[index])) for index, _ in zip(indices, values)]\n",
        "\n",
        "\n",
        "def transformer_blocks(model: AutoModelForCausalLM):\n",
        "    \"\"\"Return the repeated transformer blocks inside common causal LM architectures.\n",
        "\n",
        "    Different Hugging Face model families store their layers under different\n",
        "    attribute names. This helper finds the list of blocks so the steering hook can\n",
        "    be attached without hard-coding one model family.\n",
        "    \"\"\"\n",
        "    # Try the common locations used by GPT-2, Qwen/Llama-style models, and GPT-NeoX.\n",
        "    candidates = [\n",
        "        (\"transformer\", \"h\"),\n",
        "        (\"model\", \"layers\"),\n",
        "        (\"gpt_neox\", \"layers\"),\n",
        "    ]\n",
        "\n",
        "    for parent_name, blocks_name in candidates:\n",
        "        # getattr(..., None) avoids crashing when a model does not have that attribute.\n",
        "        parent = getattr(model, parent_name, None)\n",
        "        blocks = getattr(parent, blocks_name, None)\n",
        "        if blocks is not None and len(blocks) > 0:\n",
        "            return blocks\n",
        "\n",
        "    raise TypeError(\"This demo could not find a supported transformer block list on the model.\")\n",
        "\n",
        "\n",
        "def generate_text(\n",
        "    model: AutoModelForCausalLM,\n",
        "    tokenizer: AutoTokenizer,\n",
        "    prompt: str,\n",
        "    max_new_tokens: int = 80,\n",
        ") -> str:\n",
        "    \"\"\"Generate normal unsteered text from the model.\n",
        "\n",
        "    This is the baseline generation path. It tokenizes the prompt, samples new\n",
        "    tokens from the model, and decodes the final token IDs back into readable text.\n",
        "    \"\"\"\n",
        "    # Tokenize one prompt and move the inputs to the model's device.\n",
        "    inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n",
        "    with torch.no_grad():\n",
        "        # Sampling makes output more varied than greedy decoding, but also less deterministic.\n",
        "        output_ids = model.generate(\n",
        "            **inputs,\n",
        "            max_new_tokens=max_new_tokens,\n",
        "            do_sample=True,\n",
        "            temperature=0.8,\n",
        "            top_p=0.9,\n",
        "            pad_token_id=tokenizer.eos_token_id,\n",
        "        )\n",
        "    # Decode the first generated sequence into text and remove control tokens.\n",
        "    return tokenizer.decode(output_ids[0], skip_special_tokens=True)\n",
        "\n",
        "\n",
        "def generate_with_steering(\n",
        "    model: AutoModelForCausalLM,\n",
        "    tokenizer: AutoTokenizer,\n",
        "    prompt: str,\n",
        "    layer_index: int,\n",
        "    steering_vector: torch.Tensor,\n",
        "    strength: float,\n",
        ") -> str:\n",
        "    \"\"\"Generate text while temporarily nudging one model layer's activations.\n",
        "\n",
        "    Here is the workflow:\n",
        "        generate_with_steering()\n",
        "            prepare steering vector\n",
        "            register hook on one layer\n",
        "            call generate_text()\n",
        "                model runs normally\n",
        "                chosen layer produces hidden states\n",
        "                PyTorch calls hook(...)\n",
        "                hook adds steering vector to that layer's output\n",
        "                model continues with modified hidden states\n",
        "            remove hook\n",
        "\n",
        "    The model weights are not edited. A forward hook adds a normalized steering\n",
        "    direction to one block's hidden states during this one call, then the hook is\n",
        "    removed in the finally block.\n",
        "    \"\"\"\n",
        "    blocks = transformer_blocks(model)\n",
        "    # hidden_states includes the embedding output as layer 0, while blocks start at 0.\n",
        "    block_index = max(0, layer_index - 1)  # hidden_states includes embeddings as layer 0\n",
        "    block_index = min(block_index, len(blocks) - 1)\n",
        "    # Normalize so strength controls the size of the nudge rather than the raw vector length.\n",
        "    steering_vector = steering_vector.to(model.device)\n",
        "    vector_norm = steering_vector.norm().clamp_min(torch.finfo(steering_vector.dtype).eps)\n",
        "    vector = (steering_vector / vector_norm).view(1, 1, -1) * strength\n",
        "\n",
        "    def hook(_module, _inputs, output):\n",
        "        \"\"\"Add the steering vector to the hidden states produced by the hooked block.\"\"\"\n",
        "        # Many transformer blocks return a tuple whose first item is hidden states.\n",
        "        if isinstance(output, tuple):\n",
        "            hidden = output[0] + vector.to(output[0].dtype)\n",
        "            return (hidden,) + output[1:]\n",
        "        # Some blocks return hidden states directly instead of wrapping them in a tuple.\n",
        "        return output + vector.to(output.dtype)\n",
        "\n",
        "    handle = blocks[block_index].register_forward_hook(hook)\n",
        "    try:\n",
        "       return generate_text(model, tokenizer, prompt)\n",
        "    finally:\n",
        "        handle.remove()\n",
        "\n",
        "def print_section(title: str) -> None:\n",
        "    \"\"\"Print a visible section heading for the command-line demo output.\"\"\"\n",
        "    print(\"\\n\" + \"=\" * 80)\n",
        "    print(title)\n",
        "    print(\"=\" * 80)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Load our model into memory. Key variables are \"model\" and \"tokenizer\""
      ],
      "metadata": {
        "id": "WUQIsFCAdGUA"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fTBQeHp88s5M"
      },
      "outputs": [],
      "source": [
        "model = \"Qwen/Qwen2.5-0.5B-Instruct\"\n",
        "\n",
        "torch.manual_seed(7)\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "print_section(\"1. Load model\")\n",
        "print(f\"Model: {model}\")\n",
        "print(f\"Device: {device}\")\n",
        "tokenizer = AutoTokenizer.from_pretrained(model)\n",
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "model = AutoModelForCausalLM.from_pretrained(model).to(device)\n",
        "model.eval()\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Generate a bunch of sentences about our subject, and not about our subject. We will feed these to the model and see which vectors \"light up\" or are activated when a subject is discussed. We want our model to obsess over the subject of writing software exploits."
      ],
      "metadata": {
        "id": "O7G4TtThdV7o"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "subject = \"writing software exploits\"\n",
        "\n",
        "print_section(\"2. Generate 20 contrastive pairs\")\n",
        "pairs = build_contrastive_pairs(subject)\n",
        "for i, pair in enumerate(pairs[:5], start=1):\n",
        "    print(f\"Pair {i:02d}\")\n",
        "    print(f\"  + subject prompt : {pair.positive}\")\n",
        "    print(f\"  - contrast prompt: {pair.negative}\")\n",
        "print(f\"... {len(pairs) - 5} more pairs built\")\n"
      ],
      "metadata": {
        "id": "7DYZDYi8_2ff"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now feed those it the model and capture the output at each layer in the transformer. We capture those values passed from from one layer to the next (hidden values) into the variables \"positive_hidden\" and \"negative_hidden\". The \"positive_hidden\" are when we did talk about the subject and \"negative_hidden\" is when we did not talk about it."
      ],
      "metadata": {
        "id": "1irH3-pmdlIp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "print_section(\"3. Capture hidden values\")\n",
        "positive_hidden = mean_last_token_hidden(model, tokenizer, [pair.positive for pair in pairs])\n",
        "negative_hidden = mean_last_token_hidden(model, tokenizer, [pair.negative for pair in pairs])\n",
        "print(f\"Positive hidden tensor shape: {tuple(positive_hidden.shape)} [layers, pairs, hidden_size]\")\n",
        "print(f\"Negative hidden tensor shape: {tuple(negative_hidden.shape)} [layers, pairs, hidden_size]\")\n",
        "print(\"Example first 6 hidden values from layer 1, pair 1:\")\n",
        "print(\"  positive:\", [round(v, 4) for v in positive_hidden[1, 0, :6].tolist()])\n",
        "print(\"  negative:\", [round(v, 4) for v in negative_hidden[1, 0, :6].tolist()])\n",
        "\n"
      ],
      "metadata": {
        "id": "pbbARLBg_6jK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Lets look at which transformer layer (think Y axis) have the biggest difference betwen when we talked about our subject and when we didn't.  Then look at which dimentions (think X axis) in that layer have the biggest differences."
      ],
      "metadata": {
        "id": "qga-uvyedu1_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print_section(\"4. Identify contrastive activations\")\n",
        "selected_layer, steering_vector, layer_norms = choose_layer_and_vector(positive_hidden, negative_hidden)\n",
        "print(\"Layer contrast strengths, measured as ||mean(subject) - mean(contrast)||:\")\n",
        "for layer, norm in enumerate(layer_norms.tolist()):\n",
        "    marker = \" <-- selected\" if layer == selected_layer else \"\"\n",
        "    print(f\"  layer {layer:02d}: {norm:.4f}{marker}\")\n",
        "print(\"Top steering-vector dimensions at the selected layer:\")\n",
        "for dim, value in top_activation_dimensions(steering_vector):\n",
        "    print(f\"  hidden dimension {dim:04d}: {value:+.5f}\")\n"
      ],
      "metadata": {
        "id": "eDDPpDXt__RY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Lets take a closer look at how dramatically those values change for those vectors. Positive means this is how it responded for a prompt about our subject. Negative means for some other subject"
      ],
      "metadata": {
        "id": "zf9esSJde3pe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print_section(\"Details of selected layer's top dimensions\")\n",
        "print(f\"Selected layer: {selected_layer}\")\n",
        "\n",
        "# Get the top dimensions again, or reuse if possible from previous cell\n",
        "# For clarity, let's recalculate or ensure we have them\n",
        "steering_dimensions_info = top_activation_dimensions(steering_vector)\n",
        "\n",
        "print(\"Comparing positive and negative hidden states for top steering dimensions:\")\n",
        "for dim_index, _ in steering_dimensions_info:\n",
        "    # We need to average across all pairs for the positive and negative hidden states\n",
        "    # to get a comparable value to the layer_deltas, which are averaged means.\n",
        "    # However, the user asked for 'values in the vector for a positive and negative match',\n",
        "    # which implies the average of the hidden states *before* subtraction.\n",
        "\n",
        "    # Let's show the mean for that dimension across all prompts at the selected layer.\n",
        "    pos_mean_dim = positive_hidden[selected_layer, :, dim_index].mean().item()\n",
        "    neg_mean_dim = negative_hidden[selected_layer, :, dim_index].mean().item()\n",
        "    steering_value_for_dim = steering_vector[dim_index].item()\n",
        "\n",
        "    print(f\"  Dimension {dim_index:04d}: \")\n",
        "    print(f\"    Mean Positive Value: {pos_mean_dim:+.5f}\")\n",
        "    print(f\"    Mean Negative Value: {neg_mean_dim:+.5f}\")\n",
        "    print(f\"    Steering Vector Value (Pos - Neg): {steering_value_for_dim:+.5f}\")\n"
      ],
      "metadata": {
        "id": "6tiyMlGwICGO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Ask AI to write a paragraph about any topic without manipulating the math.  It will probably not be our topic."
      ],
      "metadata": {
        "id": "r-FnyelmfS-L"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "neutral_prompt = \"Write a short, specific paragraph about an interesting topic:\"\n",
        "print_section(\"This s the unsteered prompt\")\n",
        "print(\"Prompt:\", neutral_prompt)\n",
        "print(\"\\n--- Unsteered generation ---\")\n",
        "print(generate_text(model, tokenizer, neutral_prompt))"
      ],
      "metadata": {
        "id": "FQzBclTdAGWZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now lets use the same prompt, but this time add our adjust those vectors by the \"strength\".  Notice that a stregnth 1.0 or 2.0 doesn't have much affect, but as the numbers increase the focus on technology increases. By 20 you are getting a lot of software and security discussions.  At 25 its all about writing security exploits.  At 30 the AI is a drooling fool obsessed with a subject and unable to form normal sentences.\n",
        "\n",
        "Try adjustng the \"strength\" variable and rerunning this cell several times to see how it becomes obsessed with out topic.\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "5GeiICs9f5eh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "neutral_prompt = \"Write a short, specific paragraph about an interesting topic:\"\n",
        "strength = 10.0   #Affects noticable above 15.  above 25 for sure  30 becomes a mental case\n",
        "print_section(\"This is the steered output\")\n",
        "print(\"Prompt:\", neutral_prompt)\n",
        "print(\"\\n--- Steered generation ---\")\n",
        "print(generate_with_steering(model, tokenizer, neutral_prompt, selected_layer, steering_vector, strength))\n",
        "\n"
      ],
      "metadata": {
        "id": "ukPNilb9Ey1N"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}