Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

ALBERT › Factorized Embedding Parameterization

Michael BrenndoerferUpdated July 19, 20254 min read

Language AI HandbookAll chapters

Part 24 of 67BERT and Variants

Chapter 174ALBERT › Factorized Embedding Parameterization

Reading section 3 of 18 · Factorized Embedding Parameterization · 4 minSection 3 of 18 · 4 min

Section 3 of 18 · 4 min read

ALBERT › Factorized Embedding ParameterizationOne section · Sections in this chapter3 / 18 ↓

Sections in this chapter18

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

BERT ties the embedding dimension directly to the hidden dimension. Every token maps to a 768-dimensional vector, which then flows through transformer layers of the same size. This seems natural, but it creates a problem: the embedding matrix grows with vocabulary size $V$ and hidden dimension $H$ , requiring $V \times H$ parameters.

Factorized Embeddings

A technique that decomposes the large embedding matrix into two smaller matrices. Instead of mapping tokens directly to hidden-size vectors ( $V \times H$ ), factorized embeddings first map to a smaller intermediate dimension ( $V \times E$ ), then project up to the hidden size ( $E \times H$ ). When $E \ll H$ , this dramatically reduces parameters.

ALBERT decouples these dimensions. Tokens first embed into a smaller space of dimension $E$ , then a linear projection maps them up to the hidden dimension $H$ . The parameter count changes from:

\text{Standard: } V \times H

to:

\text{Factorized: } V \times E + E \times H

where:

$V$ : vocabulary size (number of unique tokens the model can represent)
$H$ : hidden dimension (size of the representations flowing through transformer layers)
$E$ : embedding dimension (smaller intermediate dimension, typically 128)

The first term ( $V \times E$ ) represents the token-to-embedding lookup table. The second term ( $E \times H$ ) represents the projection matrix that expands embeddings to the hidden size. When $E \ll H$ , the factorization saves parameters because the expensive vocabulary-sized matrix uses the smaller dimension.

For BERT's vocabulary of 30,522 tokens, hidden size of 768, and an embedding dimension of 128:

\text{Standard: } 30,522 \times 768 = 23,440,896 \text{ parameters}

\text{Factorized: } 30,522 \times 128 + 128 \times 768 = 3,906,816 + 98,304 = 4,005,120 \text{ parameters}

That's an 83% reduction in embedding parameters with a single architectural change.

In[7]:

Code

class FactorizedEmbedding(nn.Module):
    """Factorized token embeddings as used in ALBERT."""

    def __init__(self, vocab_size, embedding_dim, hidden_dim):
        super().__init__()

        # First: tokens to small embedding space
        self.token_embedding = nn.Embedding(vocab_size, embedding_dim)

        # Second: project to hidden dimension
        self.projection = nn.Linear(embedding_dim, hidden_dim, bias=False)

    def forward(self, input_ids):
        # Embed tokens to small dimension
        x = self.token_embedding(input_ids)

        # Project to hidden dimension
        x = self.projection(x)

        return x

Let's compare parameter counts:

In[8]:

Code

def compare_embedding_params(vocab_size, hidden_dim, embedding_dim):
    """Compare standard vs factorized embedding parameters."""

    standard_params = vocab_size * hidden_dim

    factorized_params = vocab_size * embedding_dim + embedding_dim * hidden_dim

    reduction = 1 - factorized_params / standard_params

    return {
        "standard": standard_params,
        "factorized": factorized_params,
        "reduction": reduction,
        "ratio": standard_params / factorized_params,
    }

Out[9]:

Console

Embedding Parameter Comparison:
-----------------------------------------------------------------
BERT-Base (baseline)         |   24,030,720 params | -2.5% reduction
ALBERT-Base (E=128)          |    4,005,120 params | 82.9% reduction
ALBERT-Large (E=128)         |    4,037,888 params | 87.1% reduction
ALBERT-xLarge (E=128)        |    4,168,960 params | 93.3% reduction
ALBERT-xxLarge (E=128)       |    4,431,104 params | 96.5% reduction

The savings become more dramatic as hidden dimension increases. ALBERT-xxLarge with its 4096-dimensional hidden states would need 125 million embedding parameters under BERT's approach. Factorization reduces this to just 4.4 million.

Out[10]:

Visualization

Bar chart comparing standard vs factorized embedding parameters across ALBERT model sizes. — Embedding parameter reduction from factorization across model sizes. The savings grow with hidden dimension because the projection matrix ($E \times H$) grows linearly while the token matrix ($V \times E$) stays constant.