Criteo Boosts Ad Performance with Latest AI Advancements

Abstract. Everyone talks about AI today, and it seems AI has become an off-the-shelf tool providing plug-and-play capabilities. However, when addressing the hardest problems, such as those arising in the Commerce Media and AdTech industry, things are far more complex and require advanced AI (more specifically, machine learning) expertise. In this post, we share the story of how we continuously improve Criteo AI. We outline the challenges posed by production constraints (scale, latency, reliability, etc.), detailing the many iterations and techniques required on both the modeling and platform sides.

Introduction

Meeting shoppers wherever they are and delivering performance with AI-powered targeted ads is in Criteo’s DNA. As part of a multi-pronged effort to be at the forefront of AI (see the recent article on How RecSys & LLMs Will Converge: Architecture of Hybrid RecoAgents and the scientific publications of the Criteo AI Lab), we have recently delivered crucial improvements to the core of Criteo AI to predict an ad’s performance even more accurately.

How RecSys & LLMs Will Converge: Architecture of Hybrid RecoAgents

Authors: Flavian Vasile, Martin Bompaire

medium.com

Criteo’s clients pursue a variety of marketing objectives, ranging from increasing brand awareness to maximizing the number and/or value of sales among active customers. These objectives can be formulated as classification and regression tasks for AI models to solve, leveraging users’ past interactions with ads as training data. Accurately predicting an ad’s performance enables the selection of the most relevant ads and their optimal pricing.

The advertising problem space presents numerous challenges that must be addressed to maximize the value we deliver to our clients:

billion-scale training dataset with hundreds of features per sample (including many high-cardinality categorical features, and with frequently missing data) encoding the shopper’s intent as well as advertiser and contextual data;
a variety of optimization goals (related to, e.g., views, clicks, visits, or sales) used as training objectives for the AI model, with a huge imbalance in the frequency of the associated events;
the massive scale and stringent latency constraints of real-time bidding in online advertising (up to a billion inferences per second with a serving latency under 10 ms).

To address these needs, we’ve introduced new AI models following the blueprint established by the multi-year deep learning program launched with the creation of the Criteo AI Lab in 2018. The phased rollout of these models has already delivered a 3% uplift in Return on Ad Spend (ROAS) for participating advertisers, with significant further gains projected, marking an important milestone in Criteo AI.

In the rest of this blog post, we will shed some light on these AI models and their underlying technologies.

The new AI models

Previous models involved significant feature engineering and had limitations when optimizing against marketers’ multiple goals: indeed, the models for predicting views, clicks, visits, and sales were distinct and trained separately. To address these issues and maximize performance, we’ve introduced new deep learning models that are continuously retrained to maintain freshness and performance. These models employ a three-tier architecture consisting of embedding layers, a backbone deep architecture, and a multi-head output layer for multi-task predictions. The specific purpose of each architectural component is detailed below:

Embedding layer: to maintain backward compatibility with previous models and simplify the handling of missing data, we first discretize all numerical features (assigning a dedicated category for missing values) and treat them as categorical features. All categorical features are passed through an embedding layer, while multi-valued categorical features are processed using an embedding-bag layer. Subsequent projection of the concatenated feature embeddings into a lower-dimensional space has been found to reduce the model size without any degradation of its performance.
Backbone architecture: we tested many state-of-the-art architectures from the scientific literature on CTR prediction (see the Criteo 1TB Click Logs Dataset) and on deep tabular learning. Well-tuned, modern multi-layer perceptrons (MLP) with residual connections, inspired by the feed-forward layers of recent large language models, were found to be very competitive against many more complex alternatives, including the latest Deep & Cross Networks (DCN). Scale matters: the optimal architecture differs significantly when training on 1 million versus 1 billion rows. We observed considerable differences in offline metrics between the various architectures, despite applying similar tuning efforts.
Multi-task: the backbone architecture is followed by multiple heads, each of which predicts a user action within the conversion funnel. As users progress through their shopping journey, later events become increasingly scarce in the data. Sharing the embedding layer across event types allows rarer events to leverage representations learned from more frequent ones. For example, sales events can benefit from embeddings trained on much more abundant click data. This enables training models on a more holistic view of the user’s shopping journey, and means that we can support multiple optimization goals within a single, unified model. In practice, the sequential nature of the shopping events is used in our modeling by training a chain of dependent predictors. For example, the cross-entropy loss term associated with the presence of a sale is conditionally added to the final loss depending on the presence of an add-to-cart event (a sale always being preceded by it). Special care is also required for the balancing of the various tasks: in practice, we have explored multiple ways to weight the loss terms and adjusted the regularization of the prediction heads depending on the task.
Continual learning: to account for the ever-shifting landscape of user interests and advertising campaigns, the AI models of Criteo’s engine are retrained at least once a day. However, retraining from scratch is infeasible, as training on weeks of Criteo-scale data would take so long that the data would lose its freshness by completion. Therefore, instead of full retraining, we employ a continuous fine-tuning procedure that is lightweight and efficiently updates the model using the latest data.

To maximize GPU utilization and speed up model training, we use large mini-batches distributed with Distributed Data Parallel over multiple high-end NVIDIA GPUs colocated on the same machine and connected via NVLink. We use the largest batch size possible, which is typically O(100k). In line with Shallue et al, 2018, we have found out that it becomes “increasingly difficult to find good values for the metaparameters with larger batch sizes”.

Model training in a continual learning setup is nontrivial given that the feature categories change over time — for example, new web pages appear and new advertising campaigns are launched every day. This means that, when fine-tuning a model on new data, embedding tables are partly pre-trained (for the “old” categories) and partly initialized from scratch (for the “new” ones). This situation requires special care when tuning the learning rate and its schedule, and the regularization parameters. To further mitigate the risk of a performance drop when continuously fine-tuning over many iterations, we periodically retrain the models from scratch.

Rolling out such models to production requires dedicated strategies to cope with the so-called offline/online distribution shift. Indeed, the models are trained on historical data, for ad traffic acquired based on the predictions of the previous models. When A/B testing new models, Criteo starts acquiring a different type of traffic, on which the models are not calibrated. We tested multiple strategies to mitigate this issue, including the use of the predictions from the previous models as a loss term when training, and a post-training calibration step on the subset of the traffic acquired by the newly A/B tested model. The latter solution has shown good results and enables scaling up the A/B tests more quickly without a significant impact on the performance.

AI technology stack

Leveraging billion-scale training datasets and doing inference of deep neural networks at scale and with strict latency constraints can only be done with a bleeding-edge AI tech stack. Over the past years, Criteo has significantly invested in it to scale and optimize across the board.

Training

On the training side, we are using

PyTorch as deep learning framework;
Ray, deployed on Criteo’s on-premise Kubernetes cluster, to scale data loading and processing (Ray Data) and do distributed training (Ray Train).

Removing all the bottlenecks and doing all the optimizations required to optimally use the most recent NVIDIA GPUs (based on the Hopper or Blackwell architectures) was a big challenge. Here are some of the improvements that gradually reduced training time:

using automatic mixed precision;
leveraging fused Adam;
optimizing the embedding layer for multi-valued categorical features;
materializing the dataset in object store memory at startup instead of continually reading it from remote storage;
doing asynchronous checkpointing;
leveraging memory pinning;
doing the preprocessing with Ray actors instead of tasks (to initialize the process only once).

The use of profiling tools at various levels (see Ray Profiling) proved crucial in helping us identify the bottlenecks and the parts worth optimizing.

Inference

After training the PyTorch model, we export it to ONNX. Some preprocessing steps, such as string manipulation, are performed outside of PyTorch. We have implemented them in ONNX with ONNX Script. We then fuse the PyTorch-exported ONNX graph with the preprocessing ONNX graph. We quickly realized that having a large number of input nodes was negatively affecting inference performance, and thus decided to stack all the input features into one tensor per feature type, thus solving the issue.

We use Dynamo-Triton (with the ONNX Runtime backend) as high-performance inference server, running on GPU to fully benefit from dynamic batching. We spent time tuning the Dynamo-Triton model configuration and found throughput improvements when increasing the number of gRPC inference handler threads and when switching from malloc to tcmalloc. We also noticed a 25% speed-up when using the TensorRT execution provider compared with the CUDA one.

In the future, we will closely follow the evolution of the ahead-of-time (AOT) compilation of PyTorch models and their optimized serving.

The latest updates to Criteo AI delivers a significant boost in ad performance. These milestones mark another step in Criteo’s AI journey and reflect the continuous evolution of our technology stack and machine learning expertise. Training and serving multi-task deep learning models on structured data at such massive scale is uniquely challenging and has required multiple iterations on the model’s architecture and training setup.

Building upon this, many more improvements are in store to further improve ad performance across the marketing funnel going forward. This includes incremental changes (such as improving the encoding of numerical features and testing architectures such as DHEN and Wukong) as well as more structural ones, in particular with the large-scale modeling of the sequences of user events with Transformers. This will enable an even more powerful modeling of the full shopper journey.

Stay tuned for upcoming posts that share Criteo’s AI journey!