Technical Summary: Self-Balancing Gradient Allocation for Heterogeneity-Aware Feature Generation in Click-Through Rate Prediction

1. Problem Definition: Generative Difficulty Imbalance

The paper identifies a fundamental limitation in existing generative Click-Through Rate (CTR) prediction methods, such as DGenCTR and SGCTR. While these methods successfully employ discrete diffusion models to provide dense reconstruction supervision across all feature fields (unlike discriminative models that only supervise the final binary label), they treat feature generation as a homogeneous task.

In practice, CTR feature spaces are highly heterogeneous, comprising:

• High-cardinality ID fields (e.g., User ID, Item ID) with millions of unique values.
• Sparse categorical attributes with moderate cardinality.
• Dense numerical features.
• Variable-length behavioral sequences.

The authors argue that these fields possess vastly different intrinsic reconstruction difficulties. High-cardinality ID fields and sequences are significantly harder to reconstruct than low-cardinality categorical or numerical fields. Existing methods assign equal loss weights to all fields during training. This leads to Generative Difficulty Imbalance:

1. Gradient Dominance: Easy-to-reconstruct fields converge early and dominate the training gradients, pulling shared model parameters toward representations that prioritize these simple fields.
2. Chronic Underfitting: The most informative fields for downstream CTR tasks (high-cardinality IDs and sequences) remain underfit because they receive insufficient gradient signal relative to their difficulty.
3. Attention Imbalance: Within the denoising network, converged easy fields tend to dominate the attention mechanism, suppressing cross-field information flow toward the harder, high-signal fields.

Existing multi-task learning techniques (e.g., GradNorm, PCGrad) are designed for discrete task boundaries and do not directly apply to the continuous per-field difficulty variations in feature generation.

2. Methodology: HeteGenCTR

The authors propose HeteGenCTR, a heterogeneous generative framework built on discrete diffusion. The core innovation is a unified per-field difficulty signal that simultaneously drives two coordinated mechanisms without introducing additional hyperparameters.

Core Component: Learnable Log-Difficulty Parameters

The model learns a single scalar parameter $s_i$ for each feature field $i$, representing the log-variance of the field's reconstruction uncertainty (homoscedastic difficulty). These parameters are jointly trained with the denoising network.

Mechanism 1: Self-Balancing Loss

Derived from a multi-task maximum likelihood estimation framework (inspired by Kendall et al.), the total generation loss is reformulated to include per-field uncertainty weights.

• Formulation: The objective minimizes the negative log-likelihood:
  $$\mathcal{L}_{bal} = \sum_{i=1}^{N} \left[ \exp(-s_i) \cdot \ell_i + \frac{s_i}{2} \right]$$
  where $\ell_i$ is the per-field reconstruction loss.
• Dynamic Rebalancing: The weight $\exp(-s_i)$ is inversely proportional to the reconstruction difficulty.
  • If a field is hard to reconstruct (high $\ell_i$), the equilibrium condition drives $s_i$ to be smaller, resulting in a larger loss weight.
  • If a field is easy (low $\ell_i$), $s_i$ grows, reducing its weight.
• Stability: The authors prove that for fixed network parameters, this equilibrium is a unique, strictly convex global minimum. The mechanism creates a negative feedback loop that automatically reallocates the gradient budget toward fields that currently need it most.

Mechanism 2: Difficulty-Guided Attention

To address the imbalance at the representation level, the learned $s_i$ parameters modulate the attention mechanism within the HSTU (Hierarchical Sequential Transducer Unit) denoising network.

• Query Scaling: The attention query for field $i$ is scaled by $\exp(-s_i/2)$:
  $$\tilde{q}_i = \exp\left(-\frac{s_i}{2}\right) \cdot q_i$$
• Effect:
  • Hard fields (small $s_i$) receive larger query norms, leading to sharper attention distributions that aggressively gather relevant context.
  • Easy fields (large $s_i$) receive smaller query norms, resulting in softer, more uniform attention that reduces their influence on the pooled representation.
• Alignment: The scaling factor $\exp(-s_i/2)$ is mathematically derived to ensure the attention-level weighting aligns with the loss-level weighting ($\exp(-s_i)$), ensuring both mechanisms reinforce the same difficulty signal.

Training Pipeline

HeteGenCTR follows a coupled two-stage paradigm:

1. Pre-training: The model reconstructs all feature fields (including the click label treated as an additional field) using the self-balancing loss and difficulty-guided attention.
2. Fine-tuning: The pre-trained scoring function is directly inherited and fine-tuned for CTR prediction using standard binary cross-entropy, with no architectural changes.

3. Key Contributions

1. Problem Formalization: The paper formally identifies and characterizes the "generative difficulty imbalance," demonstrating that uniform treatment of feature fields causes easy fields to suppress learning for high-signal ID and sequence fields.
2. Unified Framework (HeteGenCTR): Proposes a framework that uses a single set of learnable difficulty parameters to simultaneously drive:
   • Self-balancing loss aggregation (loss-level gradient rebalancing).
   • Difficulty-guided attention modulation (representation-level information flow control).
3. Theoretical Guarantees: Provides a stability analysis showing the self-balancing equilibrium is a strict local minimum and derives the attention scaling coefficients to ensure consistency with the loss weighting.
4. Parameter Efficiency: The solution introduces no new hyperparameters to tune and only adds $N$ scalar parameters (one per field) to the model.

4. Experimental Results

The authors evaluated HeteGenCTR on five datasets (Criteo, Avazu, KDD12, Amazon, and a proprietary Industrial dataset) and conducted an online A/B test.

• Performance: HeteGenCTR consistently outperformed state-of-the-art discriminative (DeepFM, DCN, HSTU) and generative (GenCTR, DGenCTR, SGCTR) baselines across all datasets. Improvements were most significant on datasets with high feature heterogeneity (Amazon, Industrial).
• Ablation Studies:
  • Removing the self-balancing loss (keeping only attention modulation) resulted in significant performance drops, confirming loss reweighting is the primary driver.
  • Removing attention modulation (keeping only loss reweighting) also caused performance drops, confirming that gradient rebalancing alone is insufficient if the attention mechanism remains biased toward easy fields.
• Difficulty Parameter Analysis:
  • Learned parameters $s_i$ evolved as expected: values for easy fields (numerical, low-cardinality) increased rapidly (reducing their weights), while values for hard fields (IDs, sequences) increased slowly, maintaining high weights.
  • Converged weights showed ID and sequence fields received significantly higher effective loss weights than other fields.
• Cold-Start Impact: The model provided disproportionate gains for cold-start users (fewer than 10 historical clicks) and long-tail items, where ID embeddings are sparsest and hardest to reconstruct.
• Online A/B Test: Deployed on a large-scale e-commerce platform, HeteGenCTR achieved a 4.7% relative improvement in CTR with no increase in serving latency (as the generative component is purely offline).

5. Significance and Claims

The paper claims that heterogeneity-aware generation is a critical, previously overlooked dimension of generative CTR modeling. By resolving the generative difficulty imbalance, HeteGenCTR ensures that the most informative features (IDs and sequences) receive adequate supervision during pre-training.

The authors emphasize that their solution is holistic and parameter-free, deriving both loss reweighting and attention modulation from a single learned signal. They assert that this approach not only improves ranking metrics (AUC) but also enhances probability calibration (LogLoss) and robustness in sparse data regimes, addressing a fundamental flaw in how current generative models aggregate multi-field objectives.