CLAX: Fast and Flexible Neural Click Models in JAX

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Hager, Philipp, Zoeter, Onno, de Rijke, Maarten
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866908631264395264
author Hager, Philipp
Zoeter, Onno
de Rijke, Maarten
author_facet Hager, Philipp
Zoeter, Onno
de Rijke, Maarten
contents 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 $\approx$ 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
format Preprint
id arxiv_https___arxiv_org_abs_2511_03620
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle CLAX: Fast and Flexible Neural Click Models in JAX
Hager, Philipp
Zoeter, Onno
de Rijke, Maarten
Information Retrieval
Machine Learning
Software Engineering
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 $\approx$ 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
title CLAX: Fast and Flexible Neural Click Models in JAX
topic Information Retrieval
Machine Learning
Software Engineering
url https://arxiv.org/abs/2511.03620