C2AL: Cohort-Contrastive Auxiliary Learning For Large-Scale Recommendation Systems

## Abstract

Abstract Training large-scale recommendation models under a single global objective implicitly assumes homogeneity across user populations. However, real-world data are composites of heterogeneous cohorts with distinct conditional distributions. As models increase in scale and complexity and as more data is used for training, they become dominated by central distribution patterns, neglecting head and tail regions. This imbalance limits the model’s learning ability and can result in inactive attention weights or dead neurons. In this paper, we reveal how the attention mechanism can play a key role in factorization machines for shared embedding selection, and propose to address this challenge by analyzing the substructures in the dataset and exposing those with strong distributional contrast through auxiliary learning. Unlike previous research, which heuristically applies weighted labels or multi-task heads to mitigate such biases, we leverage partially conflicting auxiliary labels to regularize the shared representation. This approach customizes the learning process of attention layers to preserve mutual information with minority cohorts while improving global performance. We evaluated C2AL on massive production datasets with billions of data points each for six SOTA models. Experiments show that the factorization machine is able to capture fine-grained user–ad interactions using the proposed method, achieving up to a 0.16% reduction in normalized entropy overall and delivering gains exceeding 0.30% on targeted minority cohorts.

## 1 Introduction

Deep neural networks (DNNs) have demonstrated substantial advancements across a variety of application domains, including computer vision (He2016Deep), natural language processing (Vaswani2017Attention), and classification (krizhevsky_imagenet_2012) tasks. Recently, DNNs have served as the backbone of large-scale ads recommendation systems (cheng2016wide), especially in computational advertising, where they are trained to predict user actions—such as clicks or conversions—by optimizing a global objective across massive datasets.

Although substantial research efforts (caruana1997multitask; ma2018modeling; tang2020progressive; sener2018multi; liu2024conflict; yu2020gradient; jeong2025selective) have focused on improving model prediction accuracy under the assumption of homogeneous training data, real-world production datasets are typically composed of latent sub-distributions, or cohorts, which complicates the learning process. This challenge becomes even more pronounced as model and dataset scales increase; optimization tends to favor the high-density regions of the data, thereby prioritizing majority cohorts (crawshaw2020multi). As a result, models frequently underfit the distributional tails, failing to capture meaningful feature representations for minority cohorts. Particularly, in deep learning-based recommendation systems, when gradient updates are predominantly influenced by majority cohorts, the resulting attention mechanism (DBLP:journals/corr/abs-1708-04617) tends to overemphasize interactions that are common within these groups, leaving many potential interaction pathways underutilized. This dynamic introduces representation bias—a systematic limitation in the model’s ability to capture patterns unique to minority groups. Such bias results in poor generalization and miscalibration, ultimately reducing overall ad value and degrading the user experience.

To address this representation bias, we investigate auxiliary learning (caruana1997multitask), where secondary tasks are used during training to help improve the shared representation and are discarded at inference.
While empirically effective, the design of these tasks is often heuristic, and the mechanism by which they operate remains largely unanalyzed. We introduce Cohort-Contrastive Auxiliary Learning (C2AL)(^1^11The term “contrastive” refers to partially conflict traffic pattern between cohorts and is algorithmically unrelated to instance-discrimination objectives common in self-supervised learning (e.g., InfoNCE).), a method for constructing auxiliary tasks that improve the training process via auxiliary losses, leading to active interaction between features in models and robust shared parameters.

Our analysis of the resulting learning dynamics reveals how C2AL systematically reshapes the model’s factorization machine-based attention layer. Specifically, we demonstrate through gradient and parameter analysis that the C2AL objective compels the model’s attention mechanism within modern architectures like the Deep and Hierarchical Ensemble Network (DHEN) (zhang2022dhen) to learn a statistically denser and less concentrated weight distribution, which was heavily influenced by the majority cohort to collapse toward modeling only high-density regions of the input space previously. This establishes an interpretable between the training objective, attention mechanism, and the learned representation; C2AL induces a more expressive representation by promoting a richer set of feature interactions, thereby strengthening its ability to generalize across heterogeneous cohorts in the data distribution.

Our contributions are twofold:

- 1.
We introduce Cohort-Contrastive Auxiliary Learning (C2AL), a method for constructing targeted auxiliary tasks that mitigate representation bias in large-scale recommendation systems, leading to incremental overall performance improvement with no additional inference cost.
- 2.
We carried out theoretical analysis and architecture research that explains the interpretability of C2AL, which reshapes the model’s attention weights, resulting in denser and less concentrated feature interactions. By validating the interpretability on six production ads models with massive real-world data, C2AL demonstrating consistent pattern and significant improvements in predictive accuracy and normalized entropy.

## 2 Related Works

Our work uses the principles of Multi-Task Learning (MTL) as a framework for achieving mechanistic interpretability in large-scale recommendation models.

The central premise of MTL is to improve generalization by learning a shared representation across multiple tasks (caruana1997multitask). Our work focuses on a subset of MTL, auxiliary learning, where secondary training tasks serve as a regularization mechanism to shape a model’s inductive bias for a single primary objective (jaderberg2016reinforcement). A key challenge in this paradigm is managing the interplay between task signals. This has been addressed through architectural solutions that learn task-specific parameter sharing (e.g., MMOE (ma2018modeling), PLE (tang2020progressive)) and by framing MTL as a multi-objective optimization problem (sener2018multi), which has led to algorithms that manage conflicting gradients and task priorities (e.g., PCGrad (yu2020gradient), CAGrad (liu2024conflict), and selective task updates (jeong2025selective)). These methods, however, are primarily designed to optimize joint task performance.

Much of the work in auxiliary learning has focused on developing effective, but often heuristic, methods to determine when a task is beneficial. For instance, the cosine similarity between task gradients can be used as a dynamic signal to gate the auxiliary loss and prevent negative transfer (du2020adapting). The prevailing intuition is that such methods constrain optimization to a more structured solution space, forcing the model to learn more generalizable representations (nakhleh2024a; lippl2024inductive). This is particularly relevant for avoiding spurious correlations that fail on minority cohorts with different underlying causal patterns (hu2022improving; li2023removing). However, these explanations often lack a concrete, mechanistic link to specific architectural components; they posit that representations become “better” but rarely demonstrate how or where this improvement is realized. Our work moves beyond these accounts by providing a precise, mechanistic explanation, linking a targeted auxiliary objective to a specific architectural change.

## 3 Methodology

### 3.1 Problem Setup

We consider a standard supervised learning setup for a primary prediction task, such as click prediction. Given an input feature vector $\mathbf{x}\in\mathcal{X}$, the goal is to predict a label $y\in\{0,1\}$. Our model consists of two components: a shared representation encoder $f:\mathcal{X}\to\mathbb{R}^{d}$, parameterized by $\theta_{S}$, which maps the input to a $d$-dimensional embedding $\mathbf{h}=f(\mathbf{x};\theta_{S})$; and a primary prediction head $g_{\text{primary}}:\mathbb{R}^{d}\to[0,1]$, parameterized by $\theta_{H}$, which produces the final prediction $\hat{y}=g_{\text{primary}}(\mathbf{h};\theta_{H})$.

The baseline single-task objective is to find the optimal parameters $\theta_{S}^{*}$ and $\theta_{H}^{*}$ that minimize the expected loss over the data distribution $\mathcal{D}$:

$$ $\{\theta_{S}^{*},\theta_{H}^{*}\}=\arg\min_{\theta_{S},\theta_{H}}\mathbb{E}_{(\mathbf{x},y)\sim\mathcal{D}}\Big[\mathcal{L}(\hat{y},y)\Big]$ (1) $$

During back-propagation, both $\theta_{S}$ and $\theta_{H}$ will be updated according to the loss:

$$ $\theta^{(t+1)}=\theta^{(t)}-\alpha^{(t)}\nabla_{\theta}\mathcal{L}(\theta),\,\theta\in\{\theta_{S},\,\theta_{H}\}$ (2) $$

When optimized over heterogeneous data, this objective leads to the representation $\mathbf{h}$ becoming biased toward majority cohorts.

Figure: Figure 1: The Deep and Hierarchical Ensemble Network (DHEN) and its internal interaction layer.
Refer to caption: figures/DHEN.png

The baseline models are built on the state-of-the-art architecture of computational ads recommendation systems: Deep and Hierarchical Ensemble Network (DHEN) (zhang2022dhen). A key component of this architecture is the interaction layer, which, in our case, is a factorization machine (FM) based attention mechanism as shown in Fig. [1](https://arxiv.org/html/2510.02215v2#S3.F1). The FM-based models compute pairwise feature interactions before projecting them into a compressed, computationally tractable space.

For a mini-batch, let $\mathbf{X}\in\mathbb{R}^{d\times m}$ be the matrix whose columns are the $m$ active $d$-dimensional sparse embeddings. The attention mechanism uses a weight matrix $\mathbf{Y}\in\mathbb{R}^{d\times k}$ to produce a compressed interaction embedding $\mathbf{G}\in\mathbb{R}^{d\times k}$ via the bi-linear form:

$$ $\mathbf{G}=\mathbf{X}\mathbf{X}^{\top}\mathbf{Y}$ (3) $$

The term $\mathbf{X}\mathbf{X}^{\top}\in\mathbb{R}^{d\times d}$ represents the outer product of the batch’s sparse features, and the learned attention matrix $\mathbf{Y}$ projects these interactions into a $k$-dimensional compressed space.

We visualized the weight distribution generated by the attention MLP. As shown in Figure [3(a)](https://arxiv.org/html/2510.02215v2#S3.F3.sf1), the attention weights of the baseline exhibit a light-tailed distribution: only a small subset of tokens receive meaningful values, while the majority of values are concentrated near zero. This observation indicates inefficient interactions within the input features/embedding, leading to representation bias which hurts model’s performance, especially for under-represented segments.

Figure: (a) Comparison of attention weight sparsity between two representative production models. The baseline model demonstrates pronounced sparsity in attention weights, indicating that feature selection is primarily driven by majority cohorts. Consequently, features relevant to minority cohorts are frequently underrepresented or ignored.
Refer to caption: x1.png

### 3.2 The C2AL Framework

C2AL is a framework for mitigating this representation bias through two stages: (1) data-driven discovery of cohorts where a baseline model exhibits high predictive contrast, and (2) formulation of cohort-specific auxiliary tasks to regularize the shared parameters $\theta_{S}$.

#### 3.2.1 Contrastive Cohort Discovery

We segment the data along interpretable semantic axes (e.g., user value, age) into disjoint cohorts $\{\mathcal{C}_{1},\ldots,\mathcal{C}_{N}\}$. For each axis, we use a baseline model’s predictions to quantify pairwise divergence between cohort distributions, employing metrics such as KL divergence, Cosine Similarity, Jensen-Shannon distance, and Wasserstein distance.

We denote the pair of cohorts with maximal distributional disparity as $\mathcal{C}_{\text{head}}$ and $\mathcal{C}_{\text{tail}}$, and use them for auxiliary task construction. Figure [3](https://arxiv.org/html/2510.02215v2#S3.F3) illustrates an example of the distributional divergence between these cohorts. Distributional divergence provides a principled, though not exclusive, criterion for cohort discovery; practical factors like cohort size and causal structure may also influence selection.

Figure: Figure 3: Probability density functions of the baseline Model A’s (optimize for click objective on Instagram) predictions for the head and tail cohorts along a selected semantic axis, illustrating distributional divergence.
Refer to caption: x1.png

#### 3.2.2 Contrastive Auxiliary Learning

Given the identified cohorts, we construct two auxiliary binary classification tasks designed to inject targeted gradient signals into the shared encoder. We define two auxiliary labels, $y_{\text{head}}$ and $y_{\text{tail}}$, which are positive only for samples that are both positive for the primary task ($y=1$) and belong to the corresponding cohort. Using the indicator function $\mathbb{I}(\cdot)$, this is expressed as:

$$ $y_{\text{head}}=y\cdot\mathbb{I}(\mathbf{x}\in\mathcal{C}_{\text{head}})\qquad\text{and}\qquad y_{\text{tail}}=y\cdot\mathbb{I}(\mathbf{x}\in\mathcal{C}_{\text{tail}})$ $$

We augment the model with two auxiliary heads, $g_{\text{head}}$ and $g_{\text{tail}}$, parameterized by $\theta_{\text{head}}$ and $\theta_{\text{tail}}$, which take the shared representation $\mathbf{h}$ as input.
The complete C2AL objective, optimized during training, is:

$$ $\mathcal{L}_{\text{C2AL}}=\underbrace{\mathcal{L}(g_{\text{primary}}(\theta_{S};\theta_{\text{H}}),y)}_{\text{Primary Task Loss}}+\underbrace{\lambda_{\text{head}}\mathcal{L}\left(g_{\text{head}}(\theta_{S};\theta_{\text{head}}),y_{\text{head}}\right)+\lambda_{\text{tail}}\mathcal{L}\left(g_{\text{tail}}(\theta_{S};\theta_{\text{tail}}),y_{\text{tail}}\right)}_{\text{Cohort-Contrast Losses}}$ (4) $$

Importantly, this objective is used only during training. At inference time, the auxiliary heads and their parameters $\{\theta_{\text{head}},\theta_{\text{tail}}\}$ are discarded. The model reverts to the single-task architecture defined in Section [3.1](https://arxiv.org/html/2510.02215v2#S3.SS1) and is evaluated using Equation [1](https://arxiv.org/html/2510.02215v2#S3.E1), incurring no additional computational cost or architectural changes at serving.

##### Learning Dynamics.

To analyze the learning dynamics, we simplify the C2AL training objective to a primary loss and a single weighted auxiliary loss, both functions of the embedding $\mathbf{G}$:

$$ $\mathcal{L}_{\text{C2AL}}=\mathcal{L}_{\text{primary}}(\mathbf{G},y)+\lambda_{\text{aux}}\mathcal{L}_{\text{aux}}(\mathbf{G},y_{\text{aux}})$ (5) $$

Since the FM-based attention mechanism plays a critical role in terms of embedding interaction and representation quality, we present how C2AL affects the attention module by taking the partial derivative with respect to the attention matrix $\mathbf{Y}$. Backpropagating to the attention matrix $\mathbf{Y}$, the gradient is given by:

$$ $\nabla_{\mathbf{Y}}\mathcal{L}_{\text{C2AL}}=(\mathbf{X}\mathbf{X}^{\top})\left(\nabla_{\mathbf{G}}\mathcal{L}_{\text{primary}}+\lambda_{\text{aux}}\nabla_{\mathbf{G}}\mathcal{L}_{\text{aux}}\right)$ (6) $$

This equation provides the central mechanistic insight. When trained solely on the main objective (the baseline case, where $\lambda_{\text{aux}}=0$), the gradient $\nabla_{\mathbf{Y}}\mathcal{L}_{\text{C2AL}}$ is dominated by majority cohorts, causing $\mathbf{Y}$ to converge to a sparse, concentrated state that captures only globally predictive feature interactions. The auxiliary term, however, injects cohort-specific gradient signals ($\nabla_{\mathbf{G}}\mathcal{L}_{\text{aux}}$) directly into the update for $\mathbf{Y}$. Since the compressed embedding $\mathbf{G}$ depends linearly on $\mathbf{Y}$, these targeted changes to the attention matrix map directly to changes in the representation:

$$ $\Delta\mathbf{G}=(\mathbf{X}\mathbf{X}^{\top})\Delta\mathbf{Y}$ (7) $$

Figure: (a) Weight distributions across various network layers. Left: two layers preceding the attention MLP. Right: attention weights (top) and post-attention weights (bottom).

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 4
presentation: 2
questions: The paper motivates cohort-level contrastive learning as a way to capture heterogeneity among different user groups, but the experiments mainly focus on long-tail splits or data imbalance.
Could the authors clarify why there is this mismatch between the motivation and the experimental design?
Why not explore more cohort dimensions (e.g., user interests, regions, device types) to demonstrate the method’s effectiveness under broader forms of heterogeneity?
soundness: 2
strengths: The paper starts from a reasonable motivation — improving the DHEN framework for multi-task recommendation — and the idea is clear.

Introducing cohort-level contrastive auxiliary learning is a natural and practically meaningful direction to enhance representation learning.

The method is simple, easy to implement, and can be smoothly integrated into existing systems.

The gradient decomposition analysis helps explain the intuition of reducing task interference.
summary: This paper builds on the DHEN architecture for multi-task recommendation and adds a cohort-level contrastive auxiliary learning idea. Basically, it introduces cohort-based auxiliary heads and a contrastive regularization term to make the shared representation more discriminative across different user or ad groups. The idea is quite straightforward — using group-level contrast instead of instance-level contrast — and it’s a reasonable extension of DHEN. The paper also gives some gradient analysis to explain how the auxiliary loss affects shared features. Overall, it’s a simple and practical modification with clear motivation, though not a big conceptual jump.
weaknesses: The novelty is limited. The idea of cohort-level contrast or group-based regularization has already appeared in prior work (e.g., fairness or group-robust learning). Here it feels more like an extension on top of DHEN rather than a new concept.

The so-called contrastive loss is closer to an auxiliary classification objective and doesn’t really exploit the core property of contrastive learning.

Experiments are done on large industrial data but lack strong baselines (e.g., PLE, MMoE, GroupDRO, PCGrad, etc.), making it hard to tell whether the improvement is substantial.

No ablation studies are provided to show how different components (e.g., cohort definitions or loss weights) contribute to the result.

The reported improvements (0.05%–0.16%) are very small, and there’s no statistical significance analysis, which weakens the empirical claims.

The theoretical part is more interpretive than rigorous, without real theoretical innovation.

---

confidence: 4
contribution: 3
initial_rating: 4
presentation: 3
questions: 1. In online advertising, multiple tasks are often trained jointly rather than optimized independently; in such cases, how to design labels for auxiliary tasks remains unclear.
2. Is there any analysis of user embedding of long-tail users? 
3. Did you consider any multi-embedding method[1,2,3] to avoid conflict representions?
4. Is the method still suitable for transformer-based recommendation model, e.g., [4]? 

[1] STEM: unleashing the power of embeddings for multi-task recommendation. AAAI 2024. 
[2] Ads Recommendation in a Collapsed and Entangled World. KDD 2024. 
[3] On the Embedding Collapse when Scaling up Recommendation Models. ICML 2024.
[4] RankMixer: Scaling Up Ranking Models in Industrial Recommenders.
soundness: 2
strengths: - The research problem is valuable, as it focuses on the heterogeneous distribution of users.
- C2AL is applied only during the training stage, resulting in minimal additional cost.
- The effectiveness of the method is validated through online A/B testing.
- The paper provides visualization analyses, examining the attention weight distribution to demonstrate the effectiveness of  C2AL.
summary: This work argues that user data in  RS exhibit heterogeneous distributions. A single learning objective may cause the model to focus mainly on the behavioral representations of majority users while neglecting long-tail users, leading to poorer accuracy for the latter.
To address this, the authors propose a contrastive auxiliary task during training. Experiments on industrial-scale datasets and online A/B testing demonstrate that the proposed C2AL method is effective.
weaknesses: - The proposed method relies on manually defined priors to label and divide user cohorts, rather than an automated approach, which makes the manually set hyperparameters potentially influential.
- Even among long-tail users, not all samples are truly “unique,” so simply partitioning users into groups may be an oversimplified and potentially unreasonable assumption.
- The paper does not consider the optimization conflict between the auxiliary and primary tasks, and the hyperparameter λ still requires manual tuning, which may introduce risks.
- Common advertising metrics such as AUC are not reported, even though I understand that NE is an important metric.
- The work lacks validation results on public datasets.

---

confidence: 3
contribution: 1
initial_rating: 2
presentation: 2
questions: Please refer to Weaknesses.
soundness: 2
strengths: 1. C2AL addresses the problem of model bias toward the majority cohort.
2. Gradient analysis demonstrates how the auxiliary loss injects cohort-specific gradient signals into the attention matrix.
3. Experiments are conducted on a large-scale dataset.
summary: This paper addresses the problem of data heterogeneity in recommender systems leading to model bias toward the majority cohort and proposes a cohort contrastive auxiliary learning framework, C2AL. This framework first automatically discovers "head" and "tail" cohorts with high predictive distribution disparity and then constructs cohort-specific auxiliary classification tasks. The gradient signals from these auxiliary tasks serve as a functional regularizer, forcing the model's attention mechanism to learn denser and more diverse weight distributions.
weaknesses: 1. Although Figure 2 demonstrates the sparsity of attention weights, this does not prove that the selection of these features is driven by the majority cohort. A more detailed division of features into majority and minority cohorts should be performed.
2. The motivation for constructing an auxiliary task on the pair of cohorts with maximal distributional disparity is unclear.
3. The paper indicates the use of multiple distributional divergence metrics (such as KL divergence and Cosine similarity) to quantify the difference, but does not report how these metrics were chosen in the experiments.
4. The performance improvement is not significant. Although the authors emphasize the significance of this improvement, no significance test is performed.
5. The lack of an experimental setup makes the reproducibility of the proposed algorithm unconvincing.
6. The lack of sensitivity analysis, for example, how the hyperparameters introduced in Equation 10 affect performance.
7. The lack of ablation experiments, for example, to analyze the performance of using only one auxiliary task.
8. The paper targets large-scale recommender systems, but there is no specific design for such systems; it appears to be just one application.

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 3
questions: 1. The paper does not provide a detailed analysis of how the auxiliary loss weights (λ_head and λ_tail) affect model behavior or performance. These hyperparameters likely play a crucial role in balancing the influence of primary and auxiliary gradients, especially under partially conflicting objectives. 

2. It remains unclear whether the observed improvements genuinely arise from using the two most divergent cohorts or simply from the introduction of auxiliary learning itself. If cohorts with smaller or random divergence were used, would the model still exhibit similar attention weight diversification? This ambiguity raises questions about whether the gains stem from cohort-contrastive design or merely from additional regularization.
soundness: 3
strengths: 1. It introduces a novel cohort-contrastive auxiliary learning approach that bridges multi-task learning and bias correction in recommendation systems.
2. The paper provides clear gradient-based analysis linking auxiliary tasks to improved attention diversity and representation learning and offers interpretable insights through weight distribution and training dynamics visualization, enhancing model transparency.
3. Applies easily to large-scale models with no additional inference cost, making it highly practical for industrial deployment.
summary: This paper proposes Cohort-Contrastive Auxiliary Learning (C2AL), a framework to mitigate representation bias in large-scale recommender systems. By identifying user cohorts with high predictive divergence and introducing contrastive auxiliary tasks during training, C2AL enhances attention diversity and improves minority cohort representation.
weaknesses: 1. A key limitation of C2AL is its reliance on manually defined semantic axes for cohort segmentation (e.g., user age or advertiser size), without a clear or principled selection procedure. This manual choice introduces subjectivity and may lead to inconsistent results across datasets or platforms. Moreover, the method only leverages the two cohorts with the largest distributional divergence (“head” and “tail”), but the rationale for using exactly two is not well justified. Real-world data often exhibit more complex subgroup structures, and restricting to a binary contrast may oversimplify the heterogeneity. 

2. While the performance gains achieved by C2AL are statistically significant, the improvements are relatively small in absolute terms (around 0.1% in normalized entropy). 

3. In Figure 2, it is unclear whether the left and right plots correspond to experiments on different datasets or different model settings. While the visualization idea in Figures 2 and 4 is appealing, the baseline (blue) curves are largely obscured by the red ones, making it difficult to distinguish peak density and distribution width.

4.The innovation needs further clarification. According to Equation (4), the method mainly separates head and tail data for modeling and introduces a regularization term in the latent space. This is a common approach for handling data imbalance. The authors should more explicitly articulate how C2AL differs from or improves upon existing latent-space alignment and auxiliary learning methods.


Decision: Reject