Safe and Efficient Multi-Epoch Embedding Training

Authors: Shruthi Hariharan | Machine Learning Engineer II; Konish Dutta | Machine Learning Engineer II; Abhinav Naikawadi | Sr. Machine Learning Engineer; Zhihua Zhang | Machine Learning Engineer II; Rakshita Nagalla | Sr. Machine Learning Engineer; Kousik Rajesh* | Machine Learning Engineer II; Matt Lawhon | Staff Machine Learning Engineer
Tags: Machine Learning, Infrastructure, Recommendation Systems
Content Track: Technical Discussion / Architecture of big components

* == work done while at Pinterest

Introduction

Large ID embedding tables underpin some of the largest modeling advances at Pinterest, including teacher-model scaling and cross-domain pretraining for ranking and retrieval. These tables are first pre-trained to learn broadly useful Pin representations and then fine-tuned as part of the downstream surface-specific ranking model. Here, we’ll focus on the Homefeed ranking model.

The strength of the embedding tables comes from their enormous capacity to absorb information about billions of Pins and learn rich representations. That power, however, is coupled with fragile training dynamics that have historically constrained fine-tuning to one epoch.

That compromise has become increasingly dissatisfactory; single-epoch fine-tuning prevents us from effectively optimizing our models over small, fresh datasets in retraining, and disproportionately harms effective representation learning of certain components of our ranking model which benefit from longer training, like Vision Language Models (VLMs). A one-epoch ceiling forces every component of our model to learn on the schedule of its most training-sensitive component.

When trained longer, two problems appear:

1. Model quality degrades sharply. In Homefeed ranking, we saw nearly a 6% decline in our offline save metric when extending training from one epoch to two epochs.
2. GPU waste from communication-bound training accumulates. Distributed embedding lookup and synchronization overhead created large idle periods in the training loop, leaving GPUs waiting on communication rather than performing useful computation.

Background

The “one-epoch” phenomenon of large ID embedding tables is well documented in recommendation systems literature. Their enormous capacity, combined with sparse and long-tailed updates, makes it easy for the model to “memorize” individual IDs associated with repeated examples rather than learn representations that generalize.

Prior work at Pinterest had already shown that multi-epoch ID embedding training does not fail in every setting. Taming the One-Epoch Phenomenon in Online Recommendation System by Two-stage Contrastive ID Pre-training demonstrated that embeddings themselves could be pretrained for multiple epochs using a contrastive-only pretraining objective before being transferred to downstream models for finetuning. We were interested in multi-epoch post-training: once the pretrained embeddings are part of the Homefeed ranking model, could we fine-tune them for multiple epochs without catastrophic overfitting?

Related work from Pinterest Ads explored strategies to train longer in post-training, including frequency-adaptive updates and larger learning rates; however, the ID embedding tables in this setting did not exhibit the same catastrophic multi-epoch degradation we observed in Homefeed.

Making multi-epoch training safe

Freezing as a method to identify the problem
To answer why we saw a nearly 6% decline in the offline save metric when extending training to two epochs, we isolated each component to determine whether the whole model was overfitting or whether the regression was concentrated in a particular module. The embedding table emerged as the dominant source of the regression: when we froze the table, the multi-epoch regression disappeared and the result beat the single-epoch trained model.

Over three epochs, we saw a +1.25% gain over the one-epoch model. Although freezing the embedding table immediately unlocked multi-epoch training, it prevented task-specific adaptation in the embedding table. Was task-specific adaptation possible given the fragility of embedding tables? Or were we constrained to only allowing tables to train during pretraining?

Isolating different gradient paths
We attempted to control gradient paths to prevent overfitting. The model uses three types of training objectives:

1. Head losses, which optimize per-action predictions such as save and click.
2. Pairwise losses, which optimize relative ranking between positive and negative items.
3. Contrastive loss, which trains next-item prediction as an auxiliary sequence-modeling objective. Next-item prediction is used only during training and is not an output of the production ranking model.

All three objectives can update the shared embedding table, but their gradients reach it through different paths:

The contrastive objective provides one gradient path to the table through its sequence-oriented next-prediction task. Separately, for the candidate Pin being scored, the head and pairwise objectives can send gradients to the table through two candidate-side paths:

1. Transformer path: The candidate embedding passes through cross-attention and subsequent transformer layers, creating a deep gradient path from the head and pairwise objectives back to the embedding table.
2. Feature cross path: The candidate embedding feeds directly into the MLP or DCN feature tower, creating a shorter path from the head and pairwise objectives back to the table.

We tested a series of gradient isolation strategies:

• Contrastive-only updates: We blocked gradients from the head and pairwise objectives from reaching the embedding table, while keeping the contrastive path active. The table learned only from its auxiliary next-item objective.
• Selective candidate gradients: We selectively re-enabled candidate-side gradients through the transformer path, the feature cross path, or both.

Below, we compare full co-training, frozen-embedding training, and the gradient isolation strategies vs. a one-epoch full trained model.

To summarize the results

• Full co-training allows the embedding table to train alongside the rest of the ranking model and leads to catastrophic overfitting. Contrastive + both is equivalent to full co-training.
• Frozen embedding table freezes the embedding table and allows the rest of the model to train throughout the training cycle; this resulted in the best gains.
• Contrastive-only training only allows gradients from contrastive loss to reach the embedding table and allows co-training; this also solved the overfitting issue but the results were worse than fully frozen embedding tables.
• Contrastive + transformer allows gradients from the contrastive loss and transformer route gradients from pairwise and head losses. This caused mild overfitting.
• Contrastive + feature cross allows gradients from contrastive loss and feature cross route gradients from pairwise and head losses. This resulted in catastrophic overfitting. One possible explanation is that the transformer’s additional structure and depth act as a form of implicit regularization.

Although gradient isolation prevented overfitting, we could not beat the model trained with a fully frozen embedding table through this technique.

Decoupling embedding optimization

We further explored decoupling embedding optimization in two settings with different practical goals:

1. The first asked whether multiple passes over a smaller training-data window could replace a single pass over a longer window.
2. The second asked whether we could give a newly added vision-language model more optimization time without allowing the embedding tables to overfit.

Because these investigations used different modeling and data configurations, we report their results separately. However, both pointed to the same pattern: freezing the embedding table early and adapting it later made multi-epoch training effective.

Reusing a smaller training-data window
The current Homefeed ranking model runs on a 3-week single-epoch training setup, fine tuning a pretrained foundation model backbone that is trained on data across 5 product surfaces including Homefeed. We explored adopting a shorter fine-tune window with multi-epoch training to trade in fresher fine-tune data for the ability to use a fresher pretrained backbone. Reducing the fine-tune window would also enable us to experiment faster on other workstreams that require feature logging.

The first investigation treated a 20-day, one-epoch full co-training run on the Homefeed ranking model as its control. Below, we compare the save offline metric lift of different training data windows, training horizons, and update strategies.

Initial experiments on doing full co-training saw regressions despite the fresher pretrained backbone, because of the embedding table overfitting. Freezing the embedding table prevented that sharp degradation. With a 10-day training-data window, the fully frozen configuration remained near the baseline after two epochs and slightly improved after three.

We then tested out freezing the embedding table for different training windows, and found the timing of embedding adaptation mattered. Freezing the table up to the last epoch produced a 0.49% improvement; on the other hand, co-training the table during the first epoch and freezing it for the rest of training produced a 0.80% regression. This delayed update strategy remained positive as we expanded the training window, but gains saturated around three epochs. The same pattern held with a 20-day training window, leading to a 1.16% lift. On three-epoch runs, we observed that freezing the embedding tables up to the last epoch performed much better than freezing them before.

Get Pinterest Engineering’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

We hypothesized that update timing mattered because the pretrained embedding table could adapt too quickly to repeated training examples. If the embeddings became trainable from the beginning, they could absorb the learning signal before the rest of the network had learned how to use the pretrained representations effectively. Keeping the table frozen initially might allow the surrounding network to learn more generalizable patterns, after which embeddings could safely adapt.

To test this hypothesis, we held the training-data window, training horizon, and duration of embedding updates fixed, while progressively shifting the one-epoch update window later within a two-epoch run:

The results showed a roughly linear, monotonic relationship between update timing and model performance that is consistent with our hypothesis. Based on this finding, we adopted the strategy of keeping the table frozen for one epoch and allowing co-training in the second epoch over a small training data window. The smoothness of this trend suggests that update timing may be a tunable optimization dimension.

Giving the VLM more optimization time

A parallel investigation was motivated by giving a newly added VLM more time to optimize. The VLM adapted more slowly than the established ranking components, but as with the one-epoch model, extending unrestricted full co-training caused overfitting. Similar to the above results, the best update strategy for the VLM-based model was the delayed scheme. The below results are compared against a one-epoch VLM model where the model and embedding table co-trained.

Hence, we discovered a recipe to allow task-specific adaptation in post-training for embedding tables that not only beat the one-epoch model, but also the fully frozen multi-epoch model.

Making longer training efficient

To make extended training practical, we had to reduce the amount of time spent waiting on sparse communication, synchronization, compilation overhead, and data delivery. The largest initial bottleneck was distributed embedding lookup.

Semi-synchronous training

Our embedding tables are sharded across GPUs. Looking up sparse IDs requires the system to redistribute IDs to the devices that own the corresponding table shards, perform the lookup, and return the resulting embeddings. This introduces cross-GPU all-to-all communication, whose cost grows as our embedding tables, batches, and context lengths scale.

For the production Homefeed ranking model, distributed embedding lookup alone could consume more than 15% of total training time. More importantly, much of that time appeared as a bubble in the step loop: GPUs waited while data moved between devices instead of executing model kernels.

We addressed this with a semi-synchronous training pipeline that overlaps embedding communication for the next batch with computation for the current batch. The pipeline has three main components:

• a lookahead dataloader that prepares the next batch’s sparse IDs early
• a prefetching forward wrapper that launches the next lookup immediately after the current forward pass; and
• an embedding prefetcher that coordinates work across CUDA streams and events.

Conceptually, the schedule is:

1. Run the forward pass for batch B − 1.
2. Launch ID redistribution and embedding lookup for batch B on a separate prefetch stream.
3. Continue the current step’s useful computation while that communication runs in the background.
4. Synchronize only when batch B needs its embeddings. In steady state, the lookup has already completed by that point.

TorchRec’s pipelining utilities did not integrate cleanly with our custom trainers and modules; a custom implementation gave us the flexibility we needed without forcing a larger rewrite of the training stack.

Bounded embedding staleness

The pipeline is semi-synchronous because the embeddings prefetched for batch B are read before the backward pass for batch B − 1 finishes. As a result, batch B uses embeddings that were last updated by batch B − 2: a bounded delay of one embedding update.

Empirically, we observed no model-quality degradation. Two factors likely explain why:

• Consecutive batches have little overlap in their sparse IDs, so most prefetched rows would not have been modified by the immediately preceding batch anyway.
• The update schedule is approximately analogous to gradient accumulation with a step size of two for the embedding table.

To preserve correctness and performance, we had to ensure that:

• Prefetch reads did not race with backward writes on another CUDA stream
• TorchRec autograd continued to update the intended sparse rows
• Embedding collectives did not collide with backward all-reduce operations, because NCCL calls could not be overlapped safely in our execution path
• The pipelined loop remained compatible with torch.compile rather than introducing new graph breaks

Embedding lookup optimizations

Overlapping communication hides latency, but reducing the amount of communication and kernel work is even better. We made several additional changes to the embedding path.

• Fused training-ID lookup: The prefetch path previously issued multiple sharded lookups per step, such as candidate signatures and deduplicated user-sequence signatures. Each lookup paid the fixed cost of collective setup, including split exchange and payload all-to-all operations. We now concatenate these IDs into one flat request, perform a single sharded lookup, and partition the returned embeddings afterward, reducing lookups by up to a factor of three.
• Specialized singleton-row lookup: Many embedding-bag operations were degenerate bags containing exactly one ID per row. The generic ragged-bag path still performed offset processing and reduction work that these inputs did not need. We added a specialized cuEmbed operation that uses a direct gather with a custom backward implementation. Registering a fake implementation also allows torch.compile to trace the operation instead of breaking the graph. Unit tests confirmed exact gradient parity with the generic path.
• Optimized small embedding table lookups: In addition to our large ID embedding tables, we also leverage smaller lookup tables in the model, which contain only tens of rows but can receive millions of indexed lookups per step. The default backward kernel used an atomic scatter-add, creating heavy contention on a small number of rows. We replaced it with an equivalent dense formulation: one_hot(indices)ᵀ @ output_grad. This turns the weight-gradient computation into a GEMM, which the GPU executes much more efficiently. The forward pass remains a standard gather, and tests verified bit-equivalent results.

Together, these changes reduced both collective latency and the number of kernels required for sparse training.

Amortizing optimizer communication

After the lookup and normalization bottlenecks moved, optimizer communication became a larger share of the step. Sharded optimizer state still required expensive broadcasts of updated parameters, leaving the GPUs idle while parameters were exchanged.

We amortized that cost with gradient accumulation, performing optimizer updates less frequently across multiple microbatches. TorchRec’s backward-fused embedding optimizer does not natively follow the same accumulation schedule, so we scaled the embedding learning rate to approximate the corresponding accumulated update.

Conclusion

The modeling and system insights allowed us to enable multi-epoch training in production, leading to +1.09% online saves. The best schedule we’ve discovered thus far allows the rest of the network to train first before enabling embedding adaptation. Future work could explore table-specific learning-rate schedules or treat embedding adaptation as a limited optimization budget, alternating between trainable and frozen phases to capture task-specific gains while controlling overfitting.

More broadly, these findings suggest that a single optimization schedule may not be appropriate for every model component, particularly when components differ in capacity, initialization, and convergence rate. The embedding table provided the clearest evidence of this. Whether a broader decoupling could produce additional benefits remains an open question.

Acknowledgements

Thanks to the teammates and reviewers who helped with the profiling, gradient-path investigations, experiment reviews, and repeated iterations that made both the multi-epoch and systems work possible.