# CLAX: Fast and Flexible Neural Click Models in JAX

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 13 of its 64 references that are in the map.
- First posted on arXiv 5 Nov 2025.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [An overview of gradient descent optimization algorithms](https://mlsanity.com/paper/1609.04747/) (2016)
- [DCN V2: Improved Deep & Cross Network and Practical Lessons for Web-scale Learning to Rank Systems](https://mlsanity.com/paper/2008.13535/) (2020)
- [On Neural Differential Equations](https://mlsanity.com/paper/2202.02435/) (2022)
- [Compositional Embeddings Using Complementary Partitions for Memory-Efficient Recommendation Systems](https://mlsanity.com/paper/1909.02107/) (2019)
- [BlackJAX: Composable Bayesian inference in JAX](https://mlsanity.com/paper/2402.10797/) (2024)
- [Improving Deep Learning For Airbnb Search](https://mlsanity.com/paper/2002.05515/) (2020)
- [A Graph-Enhanced Click Model for Web Search](https://mlsanity.com/paper/2206.08621/) (2022)
- [A Large Scale Search Dataset for Unbiased Learning to Rank](https://mlsanity.com/paper/2207.03051/) (2022)
- [Unbiased Learning to Rank Meets Reality: Lessons from Baidu's Large-Scale Search Dataset](https://mlsanity.com/paper/2404.02543/) (2024)
- [Multi-Feature Integration for Perception-Dependent Examination-Bias Estimation](https://mlsanity.com/paper/2302.13756/) (2023)
- [SARDINE: A Simulator for Automated Recommendation in Dynamic and Interactive Environments](https://mlsanity.com/paper/2311.16586/) (2023)
- [Pretraining De-Biased Language Model with Large-scale Click Logs for Document Ranking](https://mlsanity.com/paper/2302.13498/) (2023)

## Abstract (arXiv)

CLAX is a JAX-based library that implements classic click models using modern gradient-based optimization. While neural click models have emerged over the past decade, complex click models based on probabilistic graphical models (PGMs) have not systematically adopted gradient-based optimization, preventing practitioners from leveraging modern deep learning frameworks while preserving the interpretability of classic models. CLAX addresses this gap by replacing EM-based optimization with direct gradient-based optimization in a numerically stable manner. The framework's modular design enables the integration of any component, from embeddings and deep networks to custom modules, into classic click models for end-to-end optimization. We demonstrate CLAX's efficiency by running experiments on the full Baidu-ULTR dataset comprising over a billion user sessions in ≈ 2 hours on a single GPU, orders of magnitude faster than traditional EM approaches. CLAX implements ten classic click models, serving both industry practitioners seeking to understand user behavior and improve ranking performance at scale and researchers developing new click models. CLAX is available at: https://github.com/philipphager/clax