Sparse Models and Distributed Embedding Tables

"I hold a vector for every user, every product, and every word anyone ever typed. Each training step asks me for nine of them. The other billion just sit there, judging."

An Embedding Table With Mostly Idle Rows

Most of the parameters in a large recommendation or ranking model are not in its dense layers; they are in its embedding tables. A categorical feature such as "user id" has no natural numeric meaning, so we cannot feed the raw integer into a neural network and expect it to learn. Instead we give each distinct value its own short, dense vector and let training discover what that vector should be. A platform with two hundred million users, fifty million items, and a vocabulary of hashed cross-features can easily carry tens of billions of such vectors. The dense part of the same model, the multilayer perceptron that consumes the looked-up vectors, might be only a few million parameters. The embedding tables are the elephant, and they are sparse: any one training example references only the few entities it actually involves.

This section explains the sparse structure precisely, shows why the parameter-server primitives of this chapter are the natural home for it, and builds a tiny embedding table from scratch so you can watch the sparsity in action: only the rows in a batch change, and each row carries its own optimizer state that accumulates independently of every other row. We then tie embeddings forward to the retrieval and recommendation systems that consume them at serving time.

1. Embeddings: A Learned Vector per Entity Beginner

An embedding layer is, at heart, a lookup table. Let a categorical feature take values in a vocabulary of size $V$ (for example, $V$ distinct user ids). We allocate a matrix $E \in \mathbb{R}^{V \times d}$, where $d$ is a small embedding dimension, often between 8 and 256. The embedding of entity $i$ is simply row $i$ of that matrix, $E_i \in \mathbb{R}^{d}$. There is no matrix multiply in the lookup itself: feeding the one-hot vector $e_i$ through $E$ would give $e_i^\top E = E_i$, but materializing a $V$-dimensional one-hot and multiplying it by the whole table would be absurd when $V$ is in the billions. The lookup is a memory access, an indexing operation,

A real example carries several categorical features, and a feature may be multi-valued (a user's list of recently clicked items). The model gathers the relevant rows and pools them, typically by summing or averaging, before handing the result to its dense layers. For a feature whose example holds a set $S$ of active keys, the pooled embedding is $\frac{1}{|S|}\sum_{i \in S} E_i$. The crucial property for distribution is that $|S|$ is tiny next to $V$: a user clicked a few dozen items, not a few billion. The number of rows any one example reads is bounded by the number of features it actually has, which is a property of the data, not of the table size.

2. The Sparse Gradient: Only Used Rows Move Intermediate

The forward sparsity carries straight through to the backward pass, and this is what makes the parameter-server fit exact rather than approximate. Suppose the loss $\mathcal{L}$ depends on the table only through the rows that were looked up. For a row $i$ that was not in the example, the loss does not depend on $E_i$ at all, so its gradient is exactly the zero vector. For a row that was looked up, the gradient flows back through the pooling and into that row alone:

The full gradient $\partial \mathcal{L} / \partial E$ is therefore a $V \times d$ matrix with at most $|S|$ nonzero rows. Representing it densely would waste billions of zeros; instead the framework returns a sparse gradient, a short list of (row index, gradient vector) pairs. This is precisely the shape that the push/pull interface of Section 11.2 consumes: pull the keyed rows you need, compute, push back the keyed gradients for exactly those rows. The keys that index the parameter server are the entity ids, and the sparsity of the gradient is what keeps each push small no matter how large the table grows.

3. Sparse Optimizers: Per-Row State Intermediate

Adaptive optimizers such as Adagrad and Adam keep a running statistic per parameter: a sum of squared gradients for Adagrad, first and second moment estimates for Adam. For a dense layer those statistics live alongside the weights and are updated every step. For an embedding table, updating the optimizer state of every row every step would defeat the entire point of sparsity, since most rows received a zero gradient and their state should not change. A sparse optimizer therefore updates the per-row state only for the rows that appeared in the batch. Adagrad on row $i$ accumulates

applied elementwise, and only for $i$ in the batch. A frequently seen entity (a popular item) accumulates a large $G_i$ and so takes progressively smaller steps, while a rarely seen entity keeps a small $G_i$ and a relatively large effective learning rate. Each row's state evolves on its own clock, set by how often that entity appears, which is exactly the behavior you want when entity frequencies are wildly skewed. The next code makes a tiny table and shows both the sparsity of the update and this independent per-row accumulation directly.

import numpy as np

rng = np.random.default_rng(7)
V, D = 1_000_000, 8          # one million rows (entities), 8-dim vectors each
table = np.zeros((V, D))     # the embedding table; row v is entity v's vector
G = np.zeros((V, D))         # per-row Adagrad accumulator (sparse optimizer state)
lr, eps = 0.1, 1e-8
touched = set()              # which rows have ever been pulled/updated

def step(batch_rows, target):
    # PULL: fetch only the rows this batch references (keyed sparse access).
    rows = np.array(sorted(set(batch_rows)))
    vecs = table[rows]                          # tiny slice, not the whole table
    # toy gradient: pull each used row toward a shared target vector.
    grad = vecs - target                        # gradient only for used rows
    G[rows] += grad * grad                      # accumulate per-row Adagrad state
    table[rows] -= lr * grad / (np.sqrt(G[rows]) + eps)   # per-row scaled update
    touched.update(rows.tolist())
    return rows

target = np.ones(D)
# Three batches, each referencing only a handful of the million rows.
b1 = step([3, 17, 17, 42, 256], target)
b2 = step([42, 99, 1000], target)
b3 = step([42, 256], target)   # row 42 seen in all three; 256 in two; 99 in one

rows_changed = np.flatnonzero(np.any(table != 0.0, axis=1))
print("table shape            :", table.shape)
print("rows ever touched      :", sorted(touched))
print("rows actually changed  :", rows_changed.tolist())
print("untouched fraction     :", f"{1 - len(touched)/V:.6f}")
print()
# Per-row optimizer state accumulates independently per row's visit count.
for r in [42, 256, 99, 1000]:
    print(f"row {r:>4}: visits-implied G[0]={G[r,0]:.4f}  vec[0]={table[r,0]:+.4f}")
print()
# Row 7 was never in any batch: its vector and state remain exactly zero.
print("row    7 (never used)  : vec norm =", np.linalg.norm(table[7]),
      " G norm =", np.linalg.norm(G[7]))

table shape            : (1000000, 8)
rows ever touched      : [3, 17, 42, 99, 256, 1000]
rows actually changed  : [3, 17, 42, 99, 256, 1000]
untouched fraction     : 0.999994

row   42: visits-implied G[0]=2.5041  vec[0]=+0.2195
row  256: visits-implied G[0]=1.8100  vec[0]=+0.1669
row   99: visits-implied G[0]=1.0000  vec[0]=+0.1000
row 1000: visits-implied G[0]=1.0000  vec[0]=+0.1000

row    7 (never used)  : vec norm = 0.0  G norm = 0.0

The output is the section in miniature. The table has a million rows, but the set of rows that ever changed is the union of the batch keys and nothing more, so the untouched fraction is 0.999994. The accumulator values rank exactly by how many batches each row appeared in, which is the per-row clock the optimizer runs on. A real system replaces the in-process array with a sharded parameter server, the slice indexing with keyed pull and push over the network, and the toy gradient with the backward pass of a deep model, but the access pattern is precisely the one shown here.

4. Hashing Tricks and Feature Collisions Intermediate

Even a sparse table has a ceiling: you must allocate one row per distinct entity, and for open vocabularies (raw search queries, arbitrary cross-features, a feature store that grows every day) the number of distinct keys is unbounded and unknown in advance. The hashing trick caps the table at a fixed number of rows $B$ by mapping each entity through a hash function $h$ into a bucket, $\text{row}(i) = h(i) \bmod B$. Now the table size is a budget you choose, not a quantity the data dictates, and unseen entities map to some bucket without any vocabulary lookup. The cost is collisions: two distinct entities that hash to the same bucket share a single embedding and become indistinguishable to the model. With $B$ buckets and $n$ active entities, the expected number of collided pairs grows like $n^2 / (2B)$, so you trade memory against the blurring of rare entities.

In practice collisions are tolerable because they fall hardest on infrequent keys, whose embeddings were poorly trained anyway, while frequent keys, although they can still collide, dominate their bucket's gradient and so steer it toward their own meaning. Multi-hashing (the "hashing trick" generalized to several independent hashes whose buckets are summed) reduces the chance that two entities collide in every hash at once, recovering much of the lost capacity at a fixed memory budget. Choosing $B$ is the same kind of resource decision as choosing the number of workers in Chapter 10: bigger is more faithful but costs memory and network, and the right size is measured, not guessed.

# pip install torchrec
from torchrec import EmbeddingBagConfig, EmbeddingBagCollection
from torchrec.optim import RowWiseAdagrad   # sparse, per-row optimizer state

cfg = EmbeddingBagConfig(
    name="user_id", embedding_dim=64,
    num_embeddings=200_000_000,             # 200M rows, sharded for you
    feature_names=["user_id"],
)
ebc = EmbeddingBagCollection(tables=[cfg], device="meta")
# DistributedModelParallel later shards `ebc` across the cluster and routes
# each batch's keys to the shard that owns them; you never index rows by hand.

5. Where These Vectors Go Next Beginner

Embeddings are not an end in themselves; they are the input currency of retrieval and recommendation. Once trained, the item embeddings become a searchable index: to recommend for a user, you take the user's embedding and find the nearest item vectors, which is approximate nearest-neighbor (ANN) search rather than a scan over billions of items. The classical-ML treatment of ANN structures lives in Chapter 12, and the distributed vector databases that serve these lookups at scale, sharding the index the same way this chapter shards the table, are the subject of Chapter 25. The full recommendation pipeline, where the embedding table is trained as in this section and queried as in Chapter 25, is built end to end in the case study of Chapter 38.

This forward arc is also why the partitioning ideas of this chapter recur: the embedding table you shard for training in Section 11.2 becomes the vector index you shard for serving in Chapter 25, and the same key-routing logic decides which machine answers a given lookup. The next section pushes the table size to its limit, examining the terabyte-scale embeddings of industrial recommendation systems and the memory hierarchy (HBM, host RAM, SSD) that holds them when even the sharded servers run short of fast memory.

6. When Sparse Tables Earn Their Complexity Intermediate

A parameter server for embeddings is not free: it adds a routing layer, a sharding scheme, and a fault-tolerance story (the table is now spread across machines that can fail, taken up in Section 11.8). It earns that complexity only when two conditions hold together: the table is too large to replicate on every worker, and the per-step access is genuinely sparse. If the table is small enough to fit on each device, replicating it and using all-reduce is simpler and often faster, because you avoid the routing round trip. If the access is dense (every step reads most rows), the parameter server loses its advantage, since it would move nearly the whole table anyway. The interesting regime, and the one industrial recommenders live in, is large-and-sparse, where the table cannot be replicated and any single step touches a vanishing fraction of it.

This is the same "match the remedy to the ceiling" discipline that the book has insisted on since Section 1.1: the binding constraint here is model memory (the table does not fit) combined with a communication structure (the gradient is sparse), and the parameter server is the remedy fitted to exactly that pair. Get either condition wrong and a simpler design wins. The skill is recognizing the large-and-sparse regime when you are in it, and the embedding table is its purest example.