unembedding
In plain wordsThe last step of a language model, which turns its internal state into a score for every token in its vocabulary; the next token is then chosen from those scores.
The final linear map of a language modellanguage model, n.: A system that has learned how likely different sequences of words are, so it can predict the next word in a text and, by doing that repeatedly, write text of its own. that converts the last hidden vector at each position, usually after a final normalization, into logitslogit, n.: A raw score that a model gives each possible answer before the scores are turned into probabilities. In statistics, the 'log-odds': a probability rewritten on a scale running from minus to plus infinity., one per vocabularyvocabulary, n.: The fixed list of all the tokens a language model knows, each with its own number. Everything the model reads or writes must be built from pieces on this list. entry; a softmaxsoftmax, n.: A formula that turns a list of raw scores into probabilities — all positive and adding up to one — with the biggest share going to the highest score. then gives the next-token distribution. It is a d × V matrix, often written W_U and called the 'LM head' in software libraries. Many models tie it to the input embeddingembedding, n.: A list of numbers that stands for a word, picture or other item, arranged so that similar items get similar lists. It lets a computer measure similarity of meaning as distance. matrix (using its transpose), saving V × d parameters, a practice called weight tying; many large models keep them separate. Interpretability work stresses that any vector in the residual streamresidual stream, n.: The running internal state that flows through a transformer model from layer to layer; each layer reads from it and adds its own contribution, like editors annotating a shared document. can be projected through W_U to read off which tokens it promotes, the basis of the logit lenslogit lens, n.: A quick trick for peeking at what a language model is 'thinking' partway through: read its half-finished internal state as if it were the final answer, one layer at a time.. The unembedding is not the inverse of the embedding.
Projecting the layer-20 residual through the unembedding already ranks the correct answer first.
Origin
Tying input and output embeddings was proposed in 2016 by Ofir Press and Lior Wolf and by Hakan Inan, Khashayar Khosravi and Richard Socher. The name 'unembedding' was popularized by Anthropic's 'A Mathematical Framework for Transformer Circuits' (Nelson Elhage and colleagues, December 2021) and the interpretability community.
Usage note
Not to be confused with the input embeddingembedding, n.: A list of numbers that stands for a word, picture or other item, arranged so that similar items get similar lists. It lets a computer measure similarity of meaning as distance., which maps token IDs into the model's vector space; the unembedding maps from that space to vocabulary scores.
Related
- embeddingembedding, n.: A list of numbers that stands for a word, picture or other item, arranged so that similar items get similar lists. It lets a computer measure similarity of meaning as distance.
- logitlogit, n.: A raw score that a model gives each possible answer before the scores are turned into probabilities. In statistics, the 'log-odds': a probability rewritten on a scale running from minus to plus infinity.
- logit lenslogit lens, n.: A quick trick for peeking at what a language model is 'thinking' partway through: read its half-finished internal state as if it were the final answer, one layer at a time.
- residual streamresidual stream, n.: The running internal state that flows through a transformer model from layer to layer; each layer reads from it and adds its own contribution, like editors annotating a shared document.
- vocabularyvocabulary, n.: The fixed list of all the tokens a language model knows, each with its own number. Everything the model reads or writes must be built from pieces on this list.
- next-token predictionnext-token prediction, n.: The simple task most AI chatbots were trained on: given some text, guess the next small piece of it. Repeating that guess, piece by piece, produces whole answers.
- embeddingembedding, n.: A list of numbers that stands for a word, picture or other item, arranged so that similar items get similar lists. It lets a computer measure similarity of meaning as distance.
Sources
- Elhage, N., Nanda, N., Olsson, C. et al. (2021). A Mathematical Framework for Transformer Circuits. Transformer Circuits Thread. transformer-circuits.pub/2021/framework/index.html
- Press, O. & Wolf, L. (2017). Using the Output Embedding to Improve Language Models. EACL 2017. arxiv.org/abs/1608.05859
Mentioned in
embeddingembedding, n.: A list of numbers that stands for a word, picture or other item, arranged so that similar items get similar lists. It lets a computer measure similarity of meaning as distance., linear transformationlinear transformation, n.: A rule for moving every point in a space that keeps straight lines straight and the origin in place: rotating, stretching, reflecting or squashing everything evenly. Every grid of numbers describes one., residual streamresidual stream, n.: The running internal state that flows through a transformer model from layer to layer; each layer reads from it and adds its own contribution, like editors annotating a shared document., reward modelreward model, n.: A separate AI model trained to score responses the way people would, so it can stand in for human judges while another model is trained on millions of examples., tokenizertokenizer, n.: The component that splits text into tokens and converts them to numbers before an AI model sees it, then turns the model's output numbers back into text., vocabularyvocabulary, n.: The fixed list of all the tokens a language model knows, each with its own number. Everything the model reads or writes must be built from pieces on this list., weight sharingweight sharing, n.: Reusing the same learned numbers in many places in a network, for example one edge detector applied at every spot in an image, so the network needs fewer parameters and learns from less data., direct logit attributiondirect logit attribution, n.: A quick way to measure how much each internal part of a language model directly pushed it toward a particular next word, by splitting the final score into separate contributions., glitch tokenglitch token, n.: A word fragment in an AI model's vocabulary that the model almost never saw during training, so asking about it can produce bizarre, evasive or garbled answers., logit lenslogit lens, n.: A quick trick for peeking at what a language model is 'thinking' partway through: read its half-finished internal state as if it were the final answer, one layer at a time., QK and OV circuitsQK and OV circuits, n.: A way of splitting what each attention unit in a language model does into two separate questions: where should it look, and what should it pass along from what it sees., tuned lenstuned lens, n.: An improved version of the logit lens that first learns a small translator for each layer, so a language model's halfway-through states can be read as predictions more accurately.