Continuous Function

Attention & Transformers

Embedding, Unembedding, and Weight Tying

How token IDs become vectors, hidden states become vocabulary logits, and tied weights make input-row and output-classifier gradients share one table.

status: reviewimportance: criticaldifficulty 4/5math: undergraduateread: 14mlive demo
Back to Attention & TransformersNext: Decoder-Only Transformer Forward Pass (review)

Concept Structure

Embedding, Unembedding, and Weight Tying

01Intuition

Start with the picture, metaphor, or geometric mechanism.

02Math

Make the objects explicit and connect them with notation.

03Code

Mirror the equations with runnable implementation details.

04Interactive Demo

Manipulate the mechanism and watch the idea respond.

3prerequisites
3next concepts
4related links

Learner Contract

What this page should let you do.

You are here becauseHow token IDs become vectors, hidden states become vocabulary logits, and tied weights make input-row and output-classifier gradients share one table.

This Attention & Transformers concept is the current idea: keep the same invariant visible across Intuition, Math, Code, Interactive Demo.

By the end4/4 sections ready | runnable code expected | live demo

Explain the mechanism, trace the main notation, and test one prediction in the live demo.

Do this firstIntuition

Read the intuition before the notation; the math should name a mechanism you already felt.

Test the linkManipulate one control and predict the visible change.Then continue to Decoder-Only Transformer Forward Pass (review)

Claim/source review status

Substantive review recorded

1/1 claims have bounded review metadata; still check caveats and source scope.Metadata-derived; review may be AI-assisted. Not a human certification.
Claims1/1 reviewed
Sources5 cited
Codeattached
Demolive
Reviewed2026-07-02
Updatedpage 2026-07-02

Learning item flow

4/4 sections readyAsk about thisResearch room
ConceptEmbedding, Unembedding, and Weight TyingAttention & TransformersEquationEmbedding, Unembedding, and Weight Tying equation 1Exact equationCodeEmbedding, Unembedding, and Weight Tying runnable code 1Exact runnable codeDemoEmbedding, Unembedding, and Weight Tying interactive demoInteractive demoClaimIn a language model, embedding is a row lookup from token IDs into re...Exact claim checkSourceStanford CS336: Language Modeling from ScratchExact source
ConceptEmbedding, Unembedding, and Weight TyingAttention & Transformers
6 sources attachedLocal snapshot ready
concept:attention-transformers/embedding-unembedding-weight-tying
CodenearbyPredictbefore revealRoomitem handoff
01IntuitionStart with the picture, metaphor, or geometric mechanism.02MathMake the objects explicit and connect them with notation.03CodeMirror the equations with runnable implementation details.04Interactive DemoManipulate the mechanism and watch the idea respond.
01

01

Intuition

Build the mental picture first so the rest of the page has something to attach to.

Section prompt

A tokenizer gives the model integers. Integer 1 might mean cat, but the Transformer cannot do attention over the idea "1" by itself. It needs a vector that can live in the residual stream.

The input embedding table is the first bridge. It asks:

Which row should I read for this token ID?

If the vocabulary has V tokens and the model width is d, the embedding table has shape (V,d). A token ID selects one row, so a batch of token IDs shaped (B,T) becomes vectors shaped (B,T,d).

The output side is the second bridge. After the Transformer has updated the residual stream, each hidden vector must become scores over the vocabulary. Unembedding asks:

Which vocabulary token does this hidden vector score highly?

That output map is a classifier over V tokens. It is often written as a matrix multiplication from hidden width d to vocabulary width V, followed by softmax and cross-entropy.

Weight tying connects the two bridges. Instead of learning one table for input rows and another unrelated classifier for output logits, the model can reuse the same parameter matrix. In a common math layout, if E in R^{V x d} is the embedding table, tied unembedding uses W_U = E^T. Then the logit for token v is the dot product between the current hidden vector and the same row E[v] that is read when token v appears as input.

This is a parameter-sharing trick, not a magic inverse. Unembedding does not recover the original token by reversing embedding. It scores every vocabulary item under the current hidden state and the loss decides which score should rise or fall.

Vaswani et al. describe learned embeddings for input/output tokens, a learned linear transformation plus softmax for predicted next-token probabilities, and sharing the same weight matrix between embedding layers and the pre-softmax linear transformation.

Press and Wolf argue that the output classifier weights form a meaningful output embedding and show why tying input and output embeddings is a serious modeling choice rather than a notation accident.

Inan, Khosravi, and Socher derive a language-modeling framework that also leads to tying input embedding and output projection weights.

02

02

Math

Translate the story into symbols, assumptions, and a derivation you can inspect.

Section prompt
Equation 1hb,t0=E[X[b,t]],h^0_{b,t} = E[X[b,t]],Equation 2hb,t0=oX[b,t]⊤E.h^0_{b,t} = o_{X[b,t]}^\top E.

Let:

  • V be vocabulary size.
  • d be residual-stream width.
  • X in {0,...,V-1}^{B x T} be a batch of token IDs.
  • E in R^{V x d} be the token embedding table.

For a token ID x, define its one-hot vector o_x in R^V. Embedding can be written either as a row lookup

hb,t0=E[X[b,t]],h^0_{b,t} = E[X[b,t]],

or as the equivalent matrix product

hb,t0=oX[b,t]⊤E.h^0_{b,t} = o_{X[b,t]}^\top E.

Across the batch, the embedding output has shape (B,T,d).

After the Transformer updates the residual stream, let h_{b,t} in R^d be the hidden state used to predict the next token. With an untied output projection,

zb,t=hb,tWU+bU,z_{b,t} = h_{b,t} W_U + b_U,

where W_U in R^{d x V}, b_U in R^V, and z_{b,t} in R^V contains one logit per vocabulary token.

The cross-entropy gradient for a target token y has the familiar form

∂L∂zv=pv−1[v=y],\frac{\partial L}{\partial z_v} = p_v - \mathbf{1}[v=y],

where p = softmax(z).

In the untied case, the selected input row E[x] receives gradient through the lookup path, while the output classifier column W_U[:,y] receives gradient through the logits path. Those are different tensors.

With tied weights,

WU=E⊤,W_U = E^\top,

so each output classifier column corresponds to an embedding row. The logit for vocabulary item v becomes

zv=h⊤E[v].z_v = h^\top E[v].

Now the same parameter row can receive two kinds of contribution: lookup-path gradient when token v appears in the input, and classifier-path gradient when the loss asks the model to raise or lower the logit for token v.

03

03

Code

Keep the implementation aligned with the notation so the algorithm is legible.

Section prompt
Runnable code 1import numpy as np rng = np.random.default_rng(7) vocab = ["the", "cat", "sat", "on", "mat",...python
import numpy as np

rng = np.random.default_rng(7)
vocab = ["the", "cat", "sat", "on", "mat", "<eos>"]
V, d = len(vocab), 4
ids = np.array([[0, 1, 2, 3]])       # the cat sat on
input_pos, target = 1, 1             # cat is read and also the target here
E = rng.normal(size=(V, d)) * 0.2
W_untied = rng.normal(size=(d, V)) * 0.2
context = np.array([0.05, -0.03, 0.02, 0.04])
h = E[ids[0, input_pos]] + context

def softmax(x):
    e = np.exp(x - x.max())
    return e / e.sum()

def ce_grads(W):
    logits = h @ W
    p = softmax(logits)
    dz = p.copy()
    dz[target] -= 1.0
    loss = -np.log(p[target])
    return loss, np.outer(h, dz), W @ dz

untied_loss, grad_Wu, grad_h = ce_grads(W_untied)
grad_E_untied = np.zeros_like(E)
grad_E_untied[ids[0, input_pos]] += grad_h

tied_loss, grad_Wu_tied, grad_h_tied = ce_grads(E.T)
grad_E_tied = grad_Wu_tied.T        # output-classifier path
grad_E_tied[ids[0, input_pos]] += grad_h_tied  # lookup path

print("E", E.shape, "W_U", W_untied.shape, "W_U tied as E.T", E.T.shape)
print("untied cat lookup grad norm", round(np.linalg.norm(grad_E_untied[1]), 4))
print("untied cat classifier grad norm", round(np.linalg.norm(grad_Wu[:, 1]), 4))
print("tied cat shared-row grad norm", round(np.linalg.norm(grad_E_tied[1]), 4))
04

04

Interactive Demo

Use direct manipulation to connect the explanation to a moving system.

Section prompt

The lab below is a vocabulary bridge audit.

Before reveal, choose whether the cat token's input-row gradient and output-classifier gradient live in separate tensors or meet in one shared row. Toggle between untied and tied mode, commit a prediction, then inspect the gradient ledger.

The key invariant is: embedding reads rows; unembedding scores vocabulary items; tying W_U = E^T makes those two roles share one parameter table.

Keep the caveat in view: tied weights are optional, layout conventions vary, and unembedding is not a geometric inverse of embedding.

Live Concept Demo

Explore Embedding, Unembedding, and Weight Tying

The stage is code-native and interactive. Use it to test the explanation against the mechanism.

difficulty 4/5undergraduatecode-aligned
Demo Prediction Checkpoint

Manipulate one control and predict the visible change.

Commit to what Embedding, Unembedding, and Weight Tying should make visible before reading the result.

After The First Pass

Turn the concept into an inspected object.

Once the invariant is visible in the intuition, math, code, and demo, use these panels to inspect the mechanism visually, check source support, practice the idea, and attach a grounded research question.

Mechanism Storyboard

See the idea move before the page explains it

How token IDs become vectors, hidden states become vocabulary logits, and tied weights make input-row and output-classifier gradients share one table.

Prediction open01 / Intuition
Prediction lens

Start with the picture, metaphor, or geometric mechanism.

Commit first

Before reading further, choose the kind of change Embedding, Unembedding, and Weight Tying should make visible.

Visual Inquiry

Make the image answer a mathematical question

How token IDs become vectors, hidden states become vocabulary logits, and tied weights make input-row and output-classifier gradients share one table.

4/4 stages readyLive demo connected
Prediction

Which visible object should carry the first intuition?

Commit first

Pick the cue that should make Embedding, Unembedding, and Weight Tying easier to reason about before the page gives the answer.

Source Grounding

Canonical references for the mechanism on this page.

course-notes · 2026Stanford CS336: Language Modeling from ScratchStanford CS336 course staff

Course-level route source for implementing language-model tokenizer, architecture, optimizer, and training pieces from scratch. This page uses it for route placement, not exact assignment-code claims.

Open source
paper · 2017Attention Is All You NeedVaswani et al.

Primary Transformer source for learned token embeddings, a learned linear transformation plus softmax for output probabilities, shared embedding/pre-softmax weights, and scaling embeddings by sqrt(d_model).

Open source
paper · 2017Using the Output Embedding to Improve Language ModelsPress and Wolf

Primary source arguing that the topmost output weight matrix is a valid word embedding and that tying input and output embeddings can improve language models while reducing parameters.

Open source
paper · 2017Tying Word Vectors and Word Classifiers: A Loss Framework for Language ModelingInan, Khosravi, and Socher

Primary source for the framework that leads to tying the input embedding matrix and output projection/classifier weights in language models.

Open source
documentation · 2026PyTorch torch.nn.EmbeddingPyTorch

Framework source for embedding as a lookup table with weight shape (num_embeddings, embedding_dim), integer-index input, and output shape with an added embedding dimension.

Open source
documentation · 2026PyTorch torch.nn.LinearPyTorch

Framework source for a linear layer y = x A^T + b, input/output shape conventions, and weight shape (out_features, in_features).

Open source

Claim Review

How token IDs become vectors, hidden states become vocabulary logits, and tied weights make input-row and output-classifier gradients share one table.

Status1 substantive review recorded

Claims without a substantive review badge still need exact source-support review.

Sources6 references

cs336-language-modeling-from-scratch, vaswani-2017-embeddings-softmax, press-wolf-2017-output-embedding, inan-2017-tying-word-vectors, pytorch-nn-embedding, pytorch-nn-linear

Local checks4 local checks

Use equations, runnable code, and demos to check whether the source support is operational.

Substantively reviewedIn a language model, embedding is a row lookup from token IDs into residual-stream vectors, unembedding is a linear map from hidden states to vocabulary logits, and weight tying reuses one parameter matrix for both directions so input-row and output-classifier gradients meet in the same tensor.Claim metadata: source checked

Vaswani et al. support learned input/output embeddings, the linear-softmax output map, and shared embedding/pre-softmax weights; Press and Wolf plus Inan et al. support tied input-output embeddings; PyTorch documents the lookup-table and affine-layer shape conventions used in the runnable code.

Sources: Attention Is All You Need, Using the Output Embedding to Improve Language Models, Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling, PyTorch torch.nn.Embedding, PyTorch torch.nn.LinearWeight tying is an optional implementation choice, not a geometric inverse of embedding and not a guarantee that rows become perfect semantic vectors. Layouts differ by framework, and the softmax/cross-entropy loss remains a separate object.A bounded review summary is present; still check caveats and exact reference scope.

Checked Vaswani et al. section 3.4, Press and Wolf 2017, Inan/Khosravi/Socher 2017, PyTorch Embedding/Linear docs, D2L language-modeling context, and Stanford CS336 route context. Page remains review-status pending GPT Pro publication critique.

Reviewer: codex-local-source-audit; reviewed 2026-07-02
source-span-vaswani-2017-embeddings-softmaxsource-span-press-wolf-2017-output-embeddingsource-span-inan-2017-tying-word-vectorsEquation 1Equation 2Runnable code 1Interactive demo

Source support candidates

course-notes 2026Stanford CS336: Language Modeling from Scratch

Course-level route source for implementing language-model tokenizer, architecture, optimizer, and training pieces from scratch. This page uses it for route placement, not exact assignment-code claims.

paper 2017Attention Is All You Need

Primary Transformer source for learned token embeddings, a learned linear transformation plus softmax for output probabilities, shared embedding/pre-softmax weights, and scaling embeddings by sqrt(d_model).

paper 2017Using the Output Embedding to Improve Language Models

Primary source arguing that the topmost output weight matrix is a valid word embedding and that tying input and output embeddings can improve language models while reducing parameters.

paper 2017Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling

Primary source for the framework that leads to tying the input embedding matrix and output projection/classifier weights in language models.

Mechanism checks

Equation 1
hb,t0=E[X[b,t]],h^0_{b,t} = E[X[b,t]],
Equation 2
hb,t0=oX[b,t]⊤E.h^0_{b,t} = o_{X[b,t]}^\top E.
Runnable code 1import numpy as np rng = np.random.default_rng(7) vocab = ["the", "cat", "sat", "on", "mat",...Demo stateLive mechanism probe

Practice Loop

Try the idea before it explains itself

How token IDs become vectors, hidden states become vocabulary logits, and tied weights make input-row and output-classifier gradients share one table.

Readiness0/3 checks ready
Predict

Before touching the demo, predict one visible change that should happen in Embedding, Unembedding, and Weight Tying.

Hint 1

Reveal when your model needs a nudge.

Hint 2

Reveal when your model needs a nudge.

Hint 3

Reveal when your model needs a nudge.

Grounded research drawerClose
ConceptEmbedding, Unembedding, and Weight TyingAttention & Transformers
Runnable code comparisonEmbedding, Unembedding, and Weight Tying runnable code 1rng = np.random.default_rng(7)Prediction before revealEmbedding, Unembedding, and Weight Tying interactive demoManipulate one control and predict the visible change.
Grounded room questionWhat is the smallest example that makes Embedding, Unembedding, and Weight Tying click without losing the math?Local snapshot ready

Research Room

Attach the question to a claim, equation, code, or demo

Pick the concept, equation, source, runnable code, claim, misconception, or demo state before asking for help. The handoff keeps that page item in context.
Next local actionNo local draft saved yet

Open the draft below to save one note and next action in this browser.

conceptAttention & Transformers

Embedding, Unembedding, and Weight Tying

Attached question

What is the smallest example that makes Embedding, Unembedding, and Weight Tying click without losing the math?

Local action draftNo local draft saved yetExpand only when ready to capture one local next action
Local action draft

This draft stays in this browser, attached to the selected learning item.

No local draft saved.
Evidence to inspect
  • References to inspect: attached references on this page.
  • Definition, prerequisite, and contrast concept links
  • The equation or runnable code that makes the concept operational
  • One demo state that shows the invariant instead of a slogan
What would resolve this
  • The learner can state the mechanism in their own words
  • The learner can name the prerequisite that would repair confusion
  • The learner can predict how the mechanism changes under one perturbation
Grounded AI handoff

I am working in Continuous Function's research reading room. Object: concept - Embedding, Unembedding, and Weight Tying Selected item key: recorded for copy. Context: Attention & Transformers Page anchor: recorded for copy. Open question: What is the smallest example that makes Embedding, Unembedding, and Weight Tying click without losing the math? Evidence to inspect: - References to inspect: attached references on this page. - Definition, prerequisite, and contrast concept links - The equation or runnable code that makes the concept operational - One demo state that shows the invariant instead of a slogan What would resolve this: - The learner can state the mechanism in their own words - The learner can name the prerequisite that would repair confusion - The learner can predict how the mechanism changes under one perturbation Answer as a careful research tutor: stay source-grounded, separate verified evidence from assumptions, name the relevant math objects, and end with one next action.

View it in context
concept/concept-notebook/attention-transformers/embedding-unembedding-weight-tying concept:attention-transformers/embedding-unembedding-weight-tying