Transformer

Stable

Inference for small transformer models, in pure Go

Used by Audio

GoAIMachine LearningTransformer
GitHub

Inference for small transformer models over sequences of frames, in float32 on the standard library alone, with no cgo. It runs any model a trainer writes in its weight file; RipModel, which hears recordings for Stapes, is one.

  • Convolutions, windowed attention with ALiBi, and towers of named output heads
  • Runs over a sequence of any length
  • A versioned, checksummed weight file, checked whole on load, that a trainer in any language can write
  • Deterministic on a platform, and one model runs on any number of goroutines at once

The forward pass

A model reads a sequence of frames and writes one output vector per frame. It runs models and does not train them: any trainer that writes the weight file can produce one, and a PyTorch copy of the model in Stapes does.

Frames

T frameseach channels planes of inputs features, any length of sequence

Input

  1. Convolutionsoptional the frames as an image, frames by features: odd kernels, GELU, a max-pool over features
  2. Input layer linear, from a frame’s features to the width

Stack

  1. Block× layers pre-norm, each half added back to its input
    1. Attention per head, the frames within the window, ALiBi-biased by distance
    2. Feed-forward to the FFN width and back, exact GELU
  2. Norm, output layer or towers: stacks of their own over the one input layer, each computing the next heads

Out

  • One vector a frame its named heads end to end: Slice(out, "pitch")
  • Raw scores no activation: the caller applies its own
RipModel’s shape, for scale: 230 features a frame, 229 mel bands and the level, with a second plane of as many when it reads PCEN; a width of 256, eight heads, four layers and a window of 64 frames, 3.2 million weights.
Attention
Each frame attends to the frames within a window on either side. Head a of h, counted from 0, subtracts s·|i − j| from each score, with s = 2−8(a+1)/h (ALiBi), in place of position embeddings, so a model runs over a sequence of any length.
Shape
Optional convolutions over the frames as an image, then the input layer, then one stack of blocks for every head, or towers, stacks of their own over the one input layer, each computing the next heads.
Numerics
float32 throughout, but for a norm's mean and variance, the GELU and the softmax's exponential, taken in float64 and rounded. It matches PyTorch to float32 rounding, within 10−5 on real recordings.
Concurrency
A loaded model is read-only: one serves any number of goroutines at once.

How far three of eight heads reach around a frame, by their ALiBi bias alone, over a window of 64: the first stays within a few frames, the last weighs the whole window nearly evenly. Nothing past the window counts.

The weight file

One file is one model: the magic VIFT, the format’s version, the shape, every tensor by name with its dimensions and float32 data, and an FNV-1a checksum over all of it. Load reads the whole file and refuses it if the checksum, the version, the shape or any tensor’s name or size disagree. The format’s document states the computation exactly, enough for a trainer in any language to write a model that means the same here.

m, err := transformer.Load(f)
out := m.Forward(frames)          // one vector per frame
pitch := m.Slice(out[t], "pitch") // a named head's scores