Inside an audio transformer

One block, several passes

Fig. 01The states change on each pass; the block’s parameters stay shared. K is the total number of passes through the block.SVG

A looped transformer applies a shared block repeatedly to its hidden states—the internal representation of a sequence. Each pass lets attention mix information between positions, then a feed-forward network transforms the representation at each position. Universal Transformers explore this idea of recurrence through depth.

Sharing weights can reduce the number of distinct parameters compared with stacking separate blocks. Every extra pass still costs computation, and more passes do not guarantee better results. This loop refines hidden states within a step; predicting the next audio token is a separate operation.

Attention with a direction

Seven-by-seven causal attention mask. Row five is highlighted through columns one to five. Positions beyond the diagonal are masked. — Fig. 02Each row is a query position; each column is a key position. Red shows the permitted context for position five. This is a visibility mask, not measured attention weights.SVG

In a causal transformer, a position can attend to itself and earlier positions. Later positions are masked. The triangular pattern makes that rule explicit: the representation used to predict what comes next cannot look ahead at future tokens.

The same rule can be used while processing a training sequence in parallel and while generating tokens one step at a time. Some audio systems use other attention patterns; this diagram describes the causal case.

From a waveform to discrete codes

Waveform passes through an encoder into a grid of discrete codes, organized across time and three illustrative codebook levels, then through a decoder back to audio. — Fig. 03A conceptual neural audio codec. Columns represent frames in time; rows represent codebook levels. The three-row grid is illustrative, and the red cells simply mark selected codes.SVG

A neural audio codec converts a dense waveform into a shorter sequence of learned representations. Quantization maps those representations to entries in discrete codebooks. A decoder reconstructs audio from the resulting codes.

Residual vector quantization, used in SoundStream, represents an encoding through successive codebook levels. These discrete sequences give an audio language model a vocabulary to predict. The number of levels and the frame rate depend on the codec; the diagram does not prescribe either for Rivet.

Generate, decode, continue

Four time steps of the same transformer. Each predicted token feeds the next step along a red path and is sent to a streaming codec decoder, which produces successive audio chunks. — Fig. 04An abstract view of causal audio generation. T denotes the same transformer at each step. Black links carry earlier context forward; red links feed predictions into the next step.SVG

An autoregressive audio model predicts the next token from the available context, appends it, and continues. AudioLM demonstrates language modeling over discrete audio representations. Previously computed attention keys and values can be cached during causal generation to avoid recomputing the entire prefix.

With a codec designed for streaming, a growing sequence can be decoded into successive audio chunks. A chunk may require several tokens and some buffering; one token is not necessarily one playable frame. End-to-end latency also depends on the codec and the inference implementation.

How to read these diagrams

These are original, simplified schematics of general research ideas. They are not a specification of Rivet’s architecture or a report of measured results. Normalization, residual connections, positional information, conditioning, and training objectives are omitted so each drawing can focus on one mechanism.

References & resourcesUniversal Transformers — Dehghani et al.Attention Is All You Need — Vaswani et al.SoundStream — Zeghidour et al.AudioLM — Borsos et al.

Back to research