1975년산 8비트 6502 프로세서에서 구동하는 자회귀 언어모델

tl;dr - I trained a tiny Mamba-based autoregressive language model and wrote an inference engine to run it on the 8-bit 6502 processor (from 1975, with 32KB RAM). Running it on my dad's BBC Micro generated the text below. once upon a time tom and lily saw things lily were sad her house he heartd them ilily and tom said yes she saw a little girl smiled tom was so excited her mom said yes The MOS 6502 is an 8-bit microprocessor released in 1975, powering the BBC Micro and the Apple II . I am lucky enough to have access to my dad's BBC Model B from the 80s; I wanted to see, using modern machine learning, what the strongest language model we could fit on this machine was. Unsurprisingly, this poses significant challenges. The model weights and inference code need to be contained within 25KB of user-space memory — my final configuration was 9KB inference code and 13KB model weights. The CPU only operates on an 8-bit integer datatype, and doesn't include multiplication in its instruction set. CC65 is used for the inference code, enabling compilation of C to the 6502 instruction set. A binary of a model trained on my MacBook can then be written to the BBC Micro using PlayUEF and a custom 3.5mm-to-tape cable I DIY'ed. This convinces the BBC that it's listening to a tape drive, while my laptop plays audio out of its headphone jack. The sim65 emulator allows a parity check between the C inference binary and the reference Python model implementation. The full inference engine can be tested on the jsbeeb emulator before running on the BBC Micro. You can run it yourself — the link below boots a BBC Micro in your browser, loads the UEF tape image straight from GitHub, and auto-types the commands to run the model. No emulator install required (note that generation takes a few minutes). Run BitNet on a BBC Micro → Modelling The goal of this project is to build an autoregressive language model — a language model that produces tokens one-by-one, similar to frontier language models. The model is a function $f$ that produces the next token from the existing context: $$ f: \text{'the cat sat on the ma'} \mapsto \text{'t'} $$ In large-scale language modeling, a 'token' would be a word or sub-word part. For this post, the vocabulary (list of tokens) used will be 26 letters plus the ' ' character. For models on this small scale, a larger vocabulary (eg. word or subword vocab) would lead to the vocabulary encoder / decoder layers consuming too much of the parameter budget. An embedding layer maps tokens into the hidden dimension (dim=56 in our case) of the model. 3Blue1Brown's video on neural networks is great for understanding how spatial token embeddings work. $$ g: \{\text{a}, \text{b}, ... , \text{z}, '\text{ }'\} \to \mathbb R^{56} $$ Once tokens are mapped to our high dimensional space, mixing layers are used to model recurrent dependencies between tokens (see recurrent layers ). BitNet BitNet was introduced as a method for fast inference on CPU. A matrix multiplication $Y = XW$ is a set of dot products of rows of $X$ with columns of $W$: $$Y_{ij} = \sum_k X_{ik} W_{kj}$$ BitNet quantizes $W$ such that its values lie in the ternary set $\{-1, 0, 1\}$. This reduces the dot product to a sequence of add/subtract operations: $$ \begin{align*} Y_{ij} &= X_{i1} W_{1j} + X_{i2} W_{2j} + \cdots + X_{in} W_{nj} \\ &= X_{i1} - X_{i2} + ... - X_{in}\\ \end{align*} $$ The 6502 processor's instruction set doesn't contain multiply: instead, a multiply is built from repeated bit-shift-and-add operations. A single 8×8 multiply and accumulate would cost 150 clock cycles. By contrast a ternary accumulate would take 30 clock cycles, making inference much faster with higher-quantization weights. Each BitNet parameter takes only $\log_2(3) = 1.58$ bits of storage, compared with 8 bits for int8, 32 for float32, etc. We can pack 4 or 5 parameters per byte: 5 parameters per byte is more data-efficient, but unpacking into ternary values requires repeated floor-divide-by-3 operations — not a native instruction on the 6502, making unpacking costly. By contrast, packing 4 parameters per byte assigns a chunk of 2 bits to each parameter, and unpacking just requires a right-shift. We opt for 4 parameters per byte for inference speed. In 13KB at 4 parameters per byte, 52k BitNet parameters can be stored. Experimental validation shows that the higher parameter count at lower quantization provides better bang-for-buck than fewer high-precision parameters. To train BitNet parameters, a similar method to other quantized training is used. The parameter is stored in full float32 precision, quantized to ternary during the forward pass, but gradients flow at full precision in the backward pass: def ternary_quantize(w: torch.Tensor) -> torch.Tensor: """Round to {-1, 0, +1} with straight-through estimator on the backward pass.""" q = torch.clamp(torch.round(w), -1.0, 1.0) # quantize forwards return w + (q - w).detach() # full precision gradient for backwards pass In practice the final LM head is kept in int4 — the output projection needs more resolution to spread probability across the vocabulary cleanly. Every other matrix parameter in the model is ternary. Recurrent Layer Architecture Attention Traditional language models such as GPT-3 used attention for modeling sequential relationships. Attention assigns vectors to each token, and uses a dot-product between each pair of vectors to allow tokens to exchange information. With a dimension of $d$ per token and a total of $s$ tokens in context, a dot product is required between our token and every token in the context window: $O(ds)$ FLOPs total. The inference forward pass (what we'll actually run on the BBC Micro) changes with context size. This is challenging with the 6502, as tensor sizes increase during inference, growing memory usage as we generate tokens. As transformers generate tokens, the KV cache grows by $O(\text{n layer} \times \text{hidden dim})$ per token generated. With our 32KB RAM budget, this would eat up the space we would prefer to use for storing model weights. What is the benefit that attention provides? It allows exact recall from each token; this allows transformers to perform well at copying and needle-in-a-haystack problems. But do we need this level of precision? For our dataset we just need enough short-term recall to learn how to spell words and produce simple grammar forms. Recurrent Models Other architectures have the property that each model forward pass has identical computation shape. State stored by the model is contained within a fixed-size state vector $h$: $$f: (t_i, h) \to t_{i+1}$$ where $t_i$ is token number $i$, and $h$ is a fixed size model hidden state. Both SSMs (eg. S4 , Mamba ) and RNN-style models ( GRUs , LSTMs ) satisfy this property, making them much more suitable for constrained-memory inference on the 6502. Why Not GRU? Recurrent models are known to be victim to the vanishing/exploding gradient problem , often leading to unstable training. Each step of the recurrence compounds errors by the magnitude of the largest eigenvalue of the forward matrix. An eigenvalue slightly larger than 1 can be catastrophic: as sequence length $s$ increases, errors compound as $\lambda^s, \lambda = 1 + \varepsilon$. In the BitNet regime, the spectral radius of our forward matrices typically far exceeds 1: in order to have a spectral radius of $\approx 1$, we would require that $98\%$ of weights are $0$. These instabilities from the GRU mean that we see divergence in every training run. The only way to avoid these is to store the primary weight matrix in a higher quantization (eg. int4). Since these are the largest matrices of the model, storing weights in int4 massively reduces the possible dimension of the model. In contrast, Mamba performs a per-channel scalar update, where the decay value per step is computed at inference-time (in the range $[0, 128) / 128$), and so by construction never exceeds 1, making explosion imp