#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = ["pillow>=10"]
# ///
"""Build the diagram of how the spread of the next-token distribution sets the readout gain.
Two tokens whose embeddings lie a distance d apart get the probabilities p and 1 - p.
Along the line between them the covariance of the embeddings is p (1 - p) d^2, and
across it the covariance is 0. With tied embeddings, a small change of the network's
output along that line therefore moves the soft token
gain = p (1 - p) d^2 / T
times as far. The four cases and every printed number come from that formula at
T = 1; no language model is involved. Pillow is needed only for --check-layout,
which measures the text with the Inter and Arimo fonts.
uv run make_spread_gain_figure.py # write spread-readout-gain.svg
uv run make_spread_gain_figure.py --verify # rebuild in memory, compare bytes
uv run make_spread_gain_figure.py --check-layout # text inside its boxes, no collisions
"""
from __future__ import annotations
import argparse
import math
import sys
from fractions import Fraction as Fr
from pathlib import Path
HERE = Path(__file__).resolve().parent
SVG = HERE / "spread-readout-gain.svg"
FONTS = {
False: ("/usr/share/fonts/opentype/inter/Inter-Regular.otf",
"/usr/share/fonts/truetype/croscore/Arimo-Regular.ttf"),
True: ("/usr/share/fonts/opentype/inter/Inter-Bold.otf",
"/usr/share/fonts/truetype/croscore/Arimo-Bold.ttf"),
}
CLS = {"title": (15, True), "h": (13, True), "s": (12, False), "v": (12, True), "m": (11, False),
"k": (11, True)}
W, H = 760, 456
T = Fr(1)
CASES = [ # heading, probability of the left token, distance between the two embeddings
("(a) Sure", Fr(99, 100), Fr(3)),
("(b) Leaning", Fr(9, 10), Fr(3)),
("(c) Torn, far apart", Fr(1, 2), Fr(3)),
("(d) Torn, close together", Fr(1, 2), Fr(1, 2)),
]
CW, CY, CH = 173, 82, 190 # card width, top and height
CX = [16 + k * (CW + 12) for k in range(4)]
S = 40 # pixels per embedding unit
DOT_Y = 164
BAR_Y, BAR_SCALE = 198, 60 # gain bar: top edge, pixels per unit of gain
AXIS_Y, AX0, AX1, VMIN, VMAX = 336, 70, 690, -3, 1
FOOT_Y = 384
def spread(p, d):
"""Variance of the drawn embedding along the line between the two tokens."""
return p * (1 - p) * d * d
assert [spread(p, d) for _, p, d in CASES] == [Fr(891, 10000), Fr(81, 100), Fr(9, 4), Fr(1, 16)]
def exact(v):
s = f"{float(v):.6f}".rstrip("0").rstrip(".")
assert Fr(s) == v, v
return s
def signed(x, digits=2):
return ("+" if x > 0 else "\u2212" if x < 0 else "") + f"{abs(x):.{digits}f}"
def percent(g):
return f"{float(100 * g):.0f}%"
def fmt(v):
s = f"{float(v):.1f}"
return s[:-2] if s.endswith(".0") else s
def enc(s):
s = s.replace("&", "&").replace("<", "<").replace(">", ">")
return "".join(c if ord(c) < 128 else f"&#x{ord(c):x};" for c in s)
def axis_x(v):
return AX0 + (v - VMIN) * (AX1 - AX0) / (VMAX - VMIN)
def cases():
out = []
for head, p, d in CASES:
g = spread(p, d) / T
out.append(dict(head=head, p=p, d=d, spread=spread(p, d), gain=g, log=math.log(g)))
return out
class Canvas:
def __init__(self):
self.out, self.texts, self.segs = [], [], []
def text(self, x, y, s, cls="m", anchor="start", extra="", box=None):
self.texts.append((s, x, y, cls, anchor, box))
a = f' text-anchor="{anchor}"' if anchor != "start" else ""
self.out.append(f' '
f'{enc(s)}')
def line(self, x1, y1, x2, y2, cls, name="line"):
self.out.append(f' ')
self.segs.append((name, x1, y1, x2, y2))
def svg(self):
return "\n".join(self.out) + "\n"
STYLE = """ """
def describe(rows):
parts = []
for r, letter in zip(rows, "abcd"):
verb = "shrinks" if r["gain"] < 1 else "stretches"
parts.append(f"({letter}) {r['head'][4:]}: probabilities {exact(r['p'])} and {exact(1 - r['p'])}, "
f"distance {exact(r['d'])}, spread {exact(r['spread'])}, log gain "
f"{'plus' if r['log'] > 0 else 'minus'} {abs(r['log']):.2f} nats, so a push "
f"{verb} to {percent(r['gain'])[:-1]} percent.")
return (
"Four cards at temperature 1. Each shows two tokens as dots sized by probability, the soft token as a "
"diamond at their balance point, and a bar for the gain on a shared scale, with a tick at a gain of 1. "
+ " ".join(parts) + " Below, a number line of the log gain from minus 3 to plus 1 nats per step "
"places d and a far below zero, b just below zero, and c above zero; zero means a gain of 1, no change. "
"A footer states that a spread is an average of squares and never negative, that only the log gain can "
"fall below zero, which means the push shrinks, and that across the line between the two tokens the "
"spread is 0, so a push in that direction does not move the soft token."
)
def figure():
rows = cases()
c = Canvas()
title = "How the spread sets the readout gain: four two-token cases"
c.out.append(f'")
return c
def check_layout(c):
from PIL import ImageFont
cache = {}
def width(s, size, bold):
return max(cache.setdefault((p, size), ImageFont.truetype(p, size)).getlength(s) for p in FONTS[bold])
boxes = {k: (CX[k] + 4, CY + 4, CX[k] + CW - 4, CY + CH - 4) for k in range(4)}
boxes.update({"top": (16, 0, W - 16, 76), "axis": (16, 284, W - 16, FOOT_Y - 4),
"foot": (28, FOOT_Y + 2, 732, FOOT_Y + 54)})
rects = []
for s, x, y, cls, anchor, box in c.texts:
size, bold = CLS[cls]
w = width(s, size, bold)
x0 = {"start": x, "middle": x - w / 2, "end": x - w}[anchor]
rects.append((s, x0, y - 0.8 * size, x0 + w, y + 0.22 * size, box))
def hits(seg, r, pad=1.0):
x0, y0, x1, y1 = r[0] - pad, r[1] - pad, r[2] + pad, r[3] + pad
_, ax, ay, bx, by = seg
t0, t1, dx, dy = 0.0, 1.0, bx - ax, by - ay
for p, q in ((-dx, ax - x0), (dx, x1 - ax), (-dy, ay - y0), (dy, y1 - ay)):
if p == 0:
if q < 0:
return False
else:
t0, t1 = (max(t0, q / p), t1) if p < 0 else (t0, min(t1, q / p))
if t0 > t1:
return False
return True
problems = []
for i, (s, x0, y0, x1, y1, box) in enumerate(rects):
b = boxes[box]
if not (x0 >= b[0] and y0 >= b[1] and x1 <= b[2] and y1 <= b[3]):
problems.append(f"outside {box}: {s!r}")
for t in rects[i + 1:]:
if x0 < t[3] and t[1] < x1 and y0 < t[4] and t[2] < y1:
problems.append(f"text overlaps text: {s!r} / {t[0]!r}")
for seg in c.segs:
if hits(seg, (x0, y0, x1, y1)):
problems.append(f"text crosses a line: {s!r} / {seg[0]}")
return problems
def main():
parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
parser.add_argument("--verify", action="store_true", help="rebuild in memory and compare bytes")
parser.add_argument("--check-layout", action="store_true", help="test text against the drawing")
args = parser.parse_args()
c = figure()
if args.check_layout:
problems = check_layout(c)
print("\n".join(problems) if problems else "layout: no collisions")
if problems:
return 1
if args.verify:
if not SVG.exists() or SVG.read_text(encoding="ascii") != c.svg():
print(f"stale or missing: {SVG.name}")
return 1
print(f"verified: {SVG.name}")
return 0
SVG.write_text(c.svg(), encoding="ascii")
print(f"wrote {SVG.name}")
return 0
if __name__ == "__main__":
sys.exit(main())