# Self-Balancing Gradient Allocation for Heterogeneity-Aware Feature Generation in Click-Through Rate Prediction

arXiv: https://arxiv.org/abs/2605.24986 · HTML: https://mlsanity.com/paper/2605.24986/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 18 of its 36 references that are in the map.
- First posted on arXiv 24 May 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 18 of 18)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Wide & Deep Learning for Recommender Systems](https://mlsanity.com/paper/1606.07792/) (2016)
- [DeepFM: A Factorization-Machine based Neural Network for CTR Prediction](https://mlsanity.com/paper/1703.04247/) (2017)
- [Conflict-Averse Gradient Descent for Multi-task Learning](https://mlsanity.com/paper/2110.14048/) (2021)
- [DCN V2: Improved Deep & Cross Network and Practical Lessons for Web-scale Learning to Rank Systems](https://mlsanity.com/paper/2008.13535/) (2020)
- [AutoInt: Automatic Feature Interaction Learning via Self-Attentive Neural Networks](https://mlsanity.com/paper/1810.11921/) (2018)
- [xDeepFM: Combining Explicit and Implicit Feature Interactions for Recommender Systems](https://mlsanity.com/paper/1803.05170/) (2018)
- [FiBiNET: Combining Feature Importance and Bilinear feature Interaction for Click-Through Rate Prediction](https://mlsanity.com/paper/1905.09433/) (2019)
- [PEPNet: Parameter and Embedding Personalized Network for Infusing with Personalized Prior Information](https://mlsanity.com/paper/2302.01115/) (2023)
- [MaskNet: Introducing Feature-Wise Multiplication to CTR Ranking Models by Instance-Guided Mask](https://mlsanity.com/paper/2102.07619/) (2021)
- [Deep Learning for Click-Through Rate Estimation](https://mlsanity.com/paper/2104.10584/) (2021)
- [On the Embedding Collapse when Scaling up Recommendation Models](https://mlsanity.com/paper/2310.04400/) (2023)
- [Hierarchically Modeling Micro and Macro Behaviors via Multi-Task Learning for Conversion Rate Prediction](https://mlsanity.com/paper/2104.09713/) (2021)
- [Towards Deeper, Lighter and Interpretable Cross Network for CTR Prediction](https://mlsanity.com/paper/2311.04635/) (2023)

## Abstract (arXiv)

Generative pre-training via discrete diffusion provides dense reconstruction supervision across all feature fields simultaneously, mitigating representation collapse from data sparsity in CTR prediction. However, all existing generative CTR methods share a fundamental limitation: the reconstruction objective assigns equal training weight to every feature field, ignoring the profound heterogeneity of reconstruction difficulty across high-cardinality ID fields, sparse categorical attributes, numerical values, and behavioral sequences. This causes easy fields to dominate training gradients while the hardest but most informative fields remain chronically underfit, a problem we term the generative difficulty imbalance.We propose HeteGenCTR, which resolves this imbalance through per-field learnable difficulty parameters jointly trained with the denoising network. This unified signal drives two coordinated components without additional hyperparameters: a self-balancing loss that automatically reallocates gradient budget toward harder fields with a provably stable equilibrium, and a difficulty-guided attention mechanism that suppresses the influence of already-converged easy fields while amplifying cross-field information flow toward hard fields. Both components share the same learned signal and remain mutually consistent throughout training. Experiments on five CTR benchmarks and a seven-day online A/B test demonstrate consistent, statistically significant improvements over state-of-the-art baselines, with disproportionate gains for cold-start and long-tail users.