File size: 5,237 Bytes
4eec64a
 
 
 
 
 
 
 
 
 
 
 
 
 
8ef5dd8
4eec64a
8ef5dd8
 
 
 
 
 
 
 
 
 
 
 
 
4eec64a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8ef5dd8
 
 
 
4eec64a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
---

license: mit
library_name: transformers
pipeline_tag: text-generation
tags:
- chip-8
- emulator
- mechanistic-interpretability
- compiled-transformer
- custom_code
---


# chipformer

**Tetris, running entirely inside a transformer.**

This model *is* a CHIP-8 computer. Give it a game as the prompt and greedy decoding runs the game:
every token it predicts is one write to the machine's memory (a register, the program counter, a
row of pixels). Paint those writes onto a 64×32 screen and you are playing Tetris. Your keypresses
and the 60 Hz clock go into the context window. Everything else, every frame, every collision,
every point on the scoreboard, comes out of the model.

There is no emulator behind it. It is the same loop every language model runs: read the context,
predict the next token, append it, repeat.

- One set of weights plays Tetris, Space Invaders, Brix, Pong and 100+ other games. Switching games
  means switching the prompt.
- An ordinary decoder-only architecture: 18 layers, d_model 382,

  42 attention heads, 22,276 ReGLU units, vocabulary 8,708, float64.



**Play it in your browser:** [divagr.com/chipformer](https://divagr.com/chipformer) ·

[Hugging Face Space](https://huggingface.co/spaces/divagr1925/chipformer) · Source: [GitHub](https://github.com/divagr18/chipformer)



## How it works



- **The ROM is the prompt.** The prompt is `BOOT`, then one `ADDR a, VAL v` token pair per nonzero

  byte of machine state, then `STEP`.

- **Every output token is a write.** Each instruction becomes a short record of writes such as

  `@V3 =1F @PC =22A STEP`. The screen is drawn from these writes.

- **The host injects the outside world.** A `TICK` token arrives once per 60 Hz frame, and keypresses

  are injected as writes. Every other token is the model's greedy prediction.

- **Attention is memory.** To read register V3, a head attends to the most recent write to V3.

  Keys use a parabola over the address plus a recency term, so the hardmax lands exactly there.

- **ReLU is the ALU.** Addition, comparison, XOR and opcode decoding are all built from ReLUs,

  using identities like step(x) = ReLU(x+1) - ReLU(x).



The architecture is ordinary:

- token embedding plus a position embedding from [t, t²]

- causal softmax attention (β = 1000) with residual connections

- ReGLU feed-forward layers

- a linear unembedding

- no normalization



`modeling_chipformer.py` contains nothing CHIP-8 specific apart from the host loop, which only feeds
tokens in and reads the screen back.

**Where the weights come from.** They are compiled from a description of the CHIP-8 circuit
(written in a small DSL in the GitHub repo) rather than learned, which is what makes every
instruction exact. 139,229 of the 32,328,660 parameters are nonzero.

## Usage

```python

import torch

from huggingface_hub import hf_hub_download

from transformers import AutoModelForCausalLM



repo = "divagr1925/chipformer"

model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, dtype=torch.float64)



rom = open(hf_hub_download(repo, "roms/tetris.ch8"), "rb").read()

game = model.play(rom, quirks="shift,loadstore,clip,jump", ipf=15)

for frame in range(120):

    game.frame(keys={5} if frame % 20 < 4 else ())  # CHIP-8 key 5 = left

print(game.ascii())

```

`play_tetris.py` animates the game in a terminal: `python play_tetris.py`.

The model is a normal `PreTrainedModel`, so you can also call it directly:
- `model(input_ids, past_key_values=cache)` returns next-token logits over the 8,708-token vocabulary.
- `game.tokens` (with `keep_tokens=True`) gives you the full trace.

### Notes

- **Use float64.** The memory-read heads compare attention scores of about 10¹³, and these are only
  exact in float64. fp16 and bf16 will break the machine.
- **Speed.** Decoding is greedy, one token at a time with a KV cache. On a CPU in PyTorch that is
  about 200 tokens/s, or 2–3 frames per second of Tetris (the game wants 60).
  - `game = model.play(...)` multiplies by sparse copies of the big feed-forward matrices, since only
    0.4% of the weights are nonzero. The outputs are identical; pass `sparse=False` to run the dense

    weights exactly as loaded.

  - The Rust/wasm engine in the GitHub repo uses the same weights with sparse matmuls and exact

    attention shortcuts. It runs the games at full speed in a browser.

- **Context.** Positions must stay below 4,194,304 (2²²). `ChipformerGame` restarts the context

  from a state dump every 250,000 tokens, as the engine does.

- **Coverage.** The weights implement the original CHIP-8 instruction set with Octo's quirk flags.

  SCHIP/XO-CHIP opcodes and the vBlank display-wait quirk are not supported.


## Files

| File | |
|---|---|
| `model.safetensors` | dense float64 weights (what `transformers` loads) |
| `model.chipformer` | the same weights, sparse (nonzeros only), for the Rust/wasm engine |
| `configuration_chipformer.py`, `modeling_chipformer.py` | the architecture plus the host loop |
| `play_tetris.py` | a terminal demo |
| `roms/tetris.ch8` | Tetris by Fran Dachille (1991), freely distributable (Revival Studios CHIP-8 pack) |