Negative binomial regression and inference using a pre-trained transformer

2025-08-06

AI Summary

In large-scale comparative studies, negative binomial regression faces severe computational bottlenecks due to repeated maximum likelihood estimation (MLE), especially in million-scale hypothesis testing. To address this, we propose a pre-trained Transformer-based framework for rapid parameter inference. Our method trains a Transformer on synthetic count data to directly infer the mean and dispersion parameters of the negative binomial distribution—bypassing numerical optimization—and introduces calibrated method-of-moments estimation as a lightweight alternative. Experiments demonstrate that Transformer-based inference achieves 20× speedup over MLE while improving estimation accuracy; moment estimation further accelerates inference by 1,000×, retaining excellent calibration and statistical power. This framework constitutes the first scalable solution for high-throughput, overdispersed count data analysis that jointly ensures high accuracy, computational efficiency, and robustness.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Relational Probabilistic Models
• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems

Application Category

• User Modeling, Personalization and Recommendation: Practical large-scale studies of user experience
• Web Mining and Content Analysis: Large pretrained models with web data
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs

Abstract

Negative binomial regression is essential for analyzing over-dispersed count data in in comparative studies, but parameter estimation becomes computationally challenging in large screens requiring millions of comparisons. We investigate using a pre-trained transformer to produce estimates of negative binomial regression parameters from observed count data, trained through synthetic data generation to learn to invert the process of generating counts from parameters. The transformer method achieved better parameter accuracy than maximum likelihood optimization while being 20 times faster. However, comparisons unexpectedly revealed that method of moment estimates performed as well as maximum likelihood optimization in accuracy, while being 1,000 times faster and producing better-calibrated and more powerful tests, making it the most efficient solution for this application.

Problem

• Estimating negative binomial regression parameters efficiently
• Comparing transformer-based and traditional estimation methods
• Identifying fastest and most accurate parameter estimation approach

Innovation

• Pre-trained transformer for parameter estimation
• Synthetic data generation for model training
• Method of moments outperforms maximum likelihood