opponent.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import itertools\n",
    "import pickle\n",
    "\n",
    "import torch\n",
    "from leela_interp import Lc0sight, LeelaBoard\n",
    "from leela_interp.tools.attention import attention_attribution, top_k_attributions\n",
    "from leela_interp.tools.patching import activation_patch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = Lc0sight(\"lc0.onnx\", device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"interesting_puzzles.pkl\", \"rb\") as f:\n",
    "    puzzles = pickle.load(f)\n",
    "len(puzzles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_sacrifice_puzzles(puzzle):\n",
    "    moves_to_play = [move for i, move in enumerate(puzzle.principal_variation) if i % 2 == 0] \n",
    "    squares_to_play_on = [x[2:4] for x in moves_to_play]\n",
    "    return all(x == squares_to_play_on[0] for x in squares_to_play_on)\n",
    "\n",
    "sacrifice_puzzles_idx = puzzles.apply(find_sacrifice_puzzles, axis=1)\n",
    "sacrifice_puzzles = puzzles[sacrifice_puzzles_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sacrifice_puzzles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at one of these puzzles:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "puzzle = sacrifice_puzzles.iloc[1]\n",
    "board = LeelaBoard.from_puzzle(puzzle)\n",
    "board"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The *principal variation* is the best sequence of moves for both sides:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "puzzle.principal_variation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Leela solves this puzzle correctly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.pretty_play(board)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualizing attention patterns\n",
    "\n",
    "Next, let's look at some attention patterns. These are 64 x 64 arrays, with one entry for each pair of squares."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer = 9\n",
    "head = 5\n",
    "\n",
    "# We're using nnsight to cache activations and do interventions. There's also an interface\n",
    "# based directly on pytorch hooks if you prefer that, see Lc0Model.capturing().\n",
    "with model.trace(board):\n",
    "    attention = model.attention_scores(layer).output[0, head].save()\n",
    "attention.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can plot slices of this attention pattern:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "square = \"b3\"\n",
    "# This converts a square in chess notation to the index inside Leela's activations for\n",
    "# that square. Note that the input to Leela is flipped depending on the current player's\n",
    "# color.\n",
    "idx = board.sq2idx(square)\n",
    "# attention has shape (query_dim, key_dim); indexing into the first one gives us a slice\n",
    "# of the attention pattern with fixed query.\n",
    "board.plot(attention[idx], caption=f\"L{layer}H{head} attention with query={square}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Attention attribution\n",
    "\n",
    "Let's look at L12H12 instead and do attention attribution (this is basically approximating a zero-ablation of individual attention weights). We'll then plot the entries with the highest attribution scores as arrows from key to query (i.e. in the direction of information flow)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attribution = attention_attribution(\n",
    "    [board], layer=12, head=12, model=model, return_pt=True\n",
    ")[0]\n",
    "values, colors = top_k_attributions(attribution, board, k=5)\n",
    "board.plot(arrows=colors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_square = puzzle.principal_variation[0][2:4]\n",
    "key_square = puzzle.principal_variation[2][2:4]\n",
    "query_idx = board.sq2idx(query_square)\n",
    "key_idx = board.sq2idx(key_square)\n",
    "\n",
    "with model.trace(board):\n",
    "    model.attention_scores(12).output[0, 12, query_idx, key_idx] = 0\n",
    "    output = model.output.save()\n",
    "\n",
    "probs = model.logits_to_probs(board, output[0])[0]\n",
    "policy = model.top_moves(board, probs, top_k=5)\n",
    "print(policy)\n",
    "print(\"WDL:\", output[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The previous top move, Ng6, is now in 4th place at only 16%. Leela also thinks it's worse (the win probability is down to 9.9% from 28.7%)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Activation patching\n",
    "\n",
    "Finally, let's do activation patching. Every puzzle in our dataset already has a \"corrupted version\" that we automatically generated. This is a very similar board position, but with a slight difference that makes the tactic no longer work. Note the new pawn on h6:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "corrupted_board = LeelaBoard.from_fen(puzzle.corrupted_fen)\n",
    "display(corrupted_board)\n",
    "model.pretty_play(corrupted_board)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could implement activation patching fairly easily with `nnsight`, but we'll instead introduce our patching helper function. Let's patch the output of L12H12 on every square:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_odds_reductions = -activation_patch(\n",
    "    module_func=model.headwise_attention_output,\n",
    "    # Layer, head, output square:\n",
    "    locations=list(itertools.product([12], [12], range(64))),\n",
    "    model=model,\n",
    "    # We could also pass in board and corrupted_board manually instead\n",
    "    puzzles=puzzle,\n",
    ")\n",
    "log_odds_reductions.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "board.plot(log_odds_reductions, caption=\"Log odds reduction for each square\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, activation patching essentially only has a big effect on g6, where L12H12 moved information to from h4."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Next steps\n",
    "We demonstrated how to use several mechanistic interpretability techniques on a single board position. It's fairly straightforward to extend these to batches of positions, see the files in `scripts` for examples. `nnsight` also makes it quite easy to use other interpretability techniques that we didn't cover here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}