Add interactive token attention visualization to attention_heads

Author: Julian Aßmann

python/Demonstration.ipynb

@@ -1,8 +1,10 @@
 """Attention visualisations"""
+
 from typing import List, Optional, Union
 
 import numpy as np
 import torch
+
 from circuitsvis.utils.render import RenderedHTML, render
 
 
@@ -15,6 +17,7 @@ def attention_heads(
     negative_color: Optional[str] = None,
     positive_color: Optional[str] = None,
     mask_upper_tri: Optional[bool] = None,
+    show_tokens: Optional[bool] = None,
 ) -> RenderedHTML:
     """Attention Heads
 
@@ -41,6 +44,8 @@ def attention_heads(
         mask_upper_tri: Whether or not to mask the upper triangular portion of
         the attention patterns. Should be true for causal attention, false for
         bidirectional attention.
+        show_tokens: Whether to display an interactive token view showing the
+        attention from each token to all other tokens.
 
     Returns:
         Html: Attention pattern visualization
@@ -54,6 +59,7 @@ def attention_heads(
         "positiveColor": positive_color,
         "tokens": tokens,
         "maskUpperTri": mask_upper_tri,
+        "showTokens": show_tokens,
     }
 
     return render(

python/circuitsvis/attention.py

@@ -47,11 +47,11 @@
     "@storybook/addon-actions": "^6.5.14",
     "@storybook/addon-essentials": "^6.5.14",
     "@storybook/addon-interactions": "^6.5.14",
-    "@storybook/addon-links": "^6.5.14",
-    "@storybook/builder-webpack5": "^6.5.14",
+    "@storybook/addon-links": "^9.0.16",
+    "@storybook/builder-webpack5": "^9.0.16",
     "@storybook/manager-webpack5": "^6.5.14",
     "@storybook/preset-typescript": "^3.0.0",
-    "@storybook/react": "^6.5.14",
+    "@storybook/react": "^9.0.16",
     "@storybook/testing-library": "^0.0.13",
     "@tensorflow/tfjs-node": "^4.1.0",
     "@testing-library/jest-dom": "5.16.5",
@@ -79,7 +79,7 @@
     "eslint-plugin-prettier": "4.2.1",
     "eslint-plugin-react": "7.31.11",
     "eslint-plugin-react-hooks": "4.6.0",
-    "eslint-plugin-storybook": "^0.6.8",
+    "eslint-plugin-storybook": "^9.0.16",
     "eslint-plugin-testing-library": "5.9.1",
     "jest": "^29.3.1",
     "jest-canvas-mock": "^2.4.0",

react/package.json

@@ -22,3 +22,19 @@ InductionHeadsLayer.args = {
   tokens: mockTokens,
   attention: mockAttention
 };
+
+export const InteractiveTokensDemo: ComponentStory<typeof AttentionHeads> =
+  Template.bind({});
+InteractiveTokensDemo.args = {
+  tokens: mockTokens,
+  attention: mockAttention,
+  showTokens: true
+};
+
+export const WithoutTokens: ComponentStory<typeof AttentionHeads> =
+  Template.bind({});
+WithoutTokens.args = {
+  tokens: mockTokens,
+  attention: mockAttention,
+  showTokens: false
+};

react/src/attention/AttentionHeads.stories.tsx

@@ -1,6 +1,9 @@
-import React from "react";
+import { Rank, tensor, Tensor4D } from "@tensorflow/tfjs";
+import React, { useMemo, useState } from "react";
 import { Col, Container, Row } from "react-grid-system";
+import tinycolor from "tinycolor2";
 import { AttentionPattern } from "./AttentionPattern";
+import { Tokens, TokensView } from "./components/AttentionTokens";
 import { useHoverLock, UseHoverLockState } from "./components/useHoverLock";
 
 /**
@@ -20,6 +23,52 @@ export function attentionHeadColor(
   return `hsla(${hue}, 70%, 50%,  ${alpha})`;
 }
 
+/**
+ * Color the attention values by heads
+ *
+ * We want attention values to be colored by each head (i.e. becoming [heads x
+ * dest_tokens x src_tokens x rgb_color_channel]). This way, when outputting an
+ * image of just one attention head it will be colored (by the specific hue
+ * assigned to that attention head) rather than grayscale.
+ *
+ * Importantly, when outputting an image that averages
+ * several attention heads we can then also average over the colors (so that we
+ * can see for each destination-source token pair which head is most important).
+ * For example, if the specific pair is very red, it suggests that the red
+ * attention head is most important for this destination-source token combination.
+ *
+ * @param attentionInput Attention input as [heads x dest_tokens x source_tokens] array
+ *
+ * @returns Tensor of the shape [heads x dest_tokens x src_tokens x
+ * rgb_color_channel]
+ */
+export function colorAttentionTensors(attentionInput: number[][][]): Tensor4D {
+  // Create a TensorFlow tensor from the attention data
+  const attentionTensor = tensor<Rank.R3>(attentionInput); // [heads x dest_tokens x source_tokens]
+
+  const attention = attentionTensor.arraySync() as number[][][];
+
+  // Set the colors
+  const colored = attention.map((head, headNumber) =>
+    head.map((destination) =>
+      destination.map((sourceAttention) => {
+        // Color
+        const attentionColor = tinycolor({
+          h: (headNumber / attention.length) * 360, // Hue (degrees 0-360)
+          s: 0.8, // Saturation (slightly off 100% to make less glaring)
+          l: 1 - 0.75 * sourceAttention // Luminance (shows amount of attention)
+        });
+
+        // Return as a nested list in the format [red, green, blue]
+        const { r, g, b } = attentionColor.toRgb();
+        return [r, g, b];
+      })
+    )
+  );
+
+  return tensor(colored);
+}
+
 /**
  * Attention Heads Selector
  */
@@ -115,14 +164,34 @@ export function AttentionHeads({
   negativeColor,
   positiveColor,
   maskUpperTri = true,
-  tokens
+  tokens,
+  showTokens = true
 }: AttentionHeadsProps) {
   // Attention head focussed state
   const { focused, onClick, onMouseEnter, onMouseLeave } = useHoverLock(0);
 
+  // State for the token view type
+  const [tokensView, setTokensView] = useState<TokensView>(
+    TokensView.DESTINATION_TO_SOURCE
+  );
+
+  // State for which token is focussed
+  const {
+    focused: focusedToken,
+    onClick: onClickToken,
+    onMouseEnter: onMouseEnterToken,
+    onMouseLeave: onMouseLeaveToken
+  } = useHoverLock();
+
   const headNames =
     attentionHeadNames || attention.map((_, idx) => `Head ${idx}`);
 
+  // Color the attention values (by head) for token interaction
+  const coloredAttention = useMemo(
+    () => colorAttentionTensors(attention),
+    [attention]
+  );
+
   return (
     <Container>
       <h3 style={{ marginBottom: 15 }}>
@@ -176,6 +245,51 @@ export function AttentionHeads({
         </Col>
       </Row>
 
+      {showTokens && (
+        <Row>
+          <Col xs={12}>
+            <div className="tokens" style={{ marginTop: 20 }}>
+              <h4 style={{ display: "inline-block", marginRight: 15 }}>
+                Interactive Tokens
+                <span style={{ fontWeight: "normal" }}>
+                  {" "}
+                  (hover/click to explore attention)
+                </span>
+              </h4>
+              <select
+                value={tokensView}
+                onChange={(e) => setTokensView(e.target.value as TokensView)}
+                style={{
+                  marginLeft: 10,
+                  padding: "5px 10px",
+                  borderRadius: 4,
+                  border: "1px solid #ccc"
+                }}
+              >
+                <option value={TokensView.DESTINATION_TO_SOURCE}>
+                  Source ← Destination
+                </option>
+                <option value={TokensView.SOURCE_TO_DESTINATION}>
+                  Destination ← Source
+                </option>
+              </select>
+              <div style={{ marginTop: 10 }}>
+                <Tokens
+                  coloredAttention={coloredAttention}
+                  focusedHead={focused}
+                  focusedToken={focusedToken}
+                  onClickToken={onClickToken}
+                  onMouseEnterToken={onMouseEnterToken}
+                  onMouseLeaveToken={onMouseLeaveToken}
+                  tokens={tokens}
+                  tokensView={tokensView}
+                />
+              </div>
+            </div>
+          </Col>
+        </Row>
+      )}
+
       <Row></Row>
     </Container>
   );
@@ -268,4 +382,13 @@ export interface AttentionHeadsProps {
    * Must be the same length as the list of values.
    */
   tokens: string[];
+
+  /**
+   * Show interactive tokens
+   *
+   * Whether to display the interactive token layer with hover/click functionality.
+   *
+   * @default true
+   */
+  showTokens?: boolean;
 }