genriesz: A Python Package for Automatic Debiased Machine Learning with Generalized Riesz Regression

Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.

34,643 characters · 21 sections · 58 citation commands

Rendered from LaTeX for readability, not typeset faithfully. Citation keys are highlighted; maths is left as source; figures, tables and equation environments are summarised rather than reproduced; unrecognised commands are greyed out so nothing is silently dropped. Email addresses are removed.

abstractEfficient estimation of causal and structural parameters can be automated using the Riesz representation theorem and debiased machine learning (DML). We present genriesz, an open-source Python package that implements automatic DML and generalized Riesz regression, a unified framework for estimating Riesz representers by minimizing empirical Bregman divergences. This framework includes covariate balancing, nearest-neighbor matching, calibrated estimation, and density ratio estimation as special cases. A key design principle of the package is automatic regressor balancing (ARB): given a Bregman generator $g$ and a representer model class, genriesz automatically constructs a compatible link function so that the generalized Riesz regression estimator satisfies balancing (moment-matching) optimality conditions in a user-chosen basis. The package provides a modular interface for specifying (i) the target linear functional via a black-box evaluation oracle, (ii) the representer model via basis functions (polynomial, RKHS approximations, random forest leaf encodings, neural embeddings, and a nearest-neighbor catchment basis), and (iii) the Bregman generator, with optional user-supplied derivatives. It returns regression adjustment (RA), Riesz weighting (RW), augmented Riesz weighting (ARW), and TMLE-style estimators with cross-fitting, confidence intervals, and $p$-values. We highlight representative workflows for estimation problems such as the average treatment effect (ATE), ATE on treated (ATT), and average marginal effect estimation. The Python package is available at \url{https://github.com/MasaKat0/genriesz} and on PyPI.

Introduction

Many targets in causal inference and econometrics can be expressed as linear functionals of an unknown regression function. Prominent examples include the average treatment effect (ATE), the average treatment effect on the treated (ATT), and average marginal effects (AME). Automatic debiased machine learning (ADML) provides general tools for valid inference on such targets Chernozhukov2022automaticdebiased.

The ADML workflow separates the problem into two parts. First, one estimates nuisance components, typically a regression function and a Riesz representer. Second, one plugs these estimates into a Neyman orthogonal score and averages the resulting score over the sample. A Donsker condition or cross-fitting then yields valid inference under weak rate conditions Chernozhukov2018doubledebiased,Klaassen1987consistentestimation.

Various methods have been proposed for Riesz representer estimation, such as Riesz regression Chernozhukov2021automaticdebiased,Chen2015sievesemiparametric, covariate balancing Imai2013covariatebalancing, and density ratio estimation Sugiyama2012densityratio. Kato2026unifiedframework,Kato2025directbias demonstrates a unifying perspective, Riesz representer fitting under Bregman divergences, which is also referred to as generalized Riesz regression, Bregman-Riesz regression, or generalized covariate balancing. Hereafter, we refer to this method as generalized Riesz regression, a name that emphasizes the connection between Riesz representer fitting and balancing weights through the choice of link function. In this framework, one can recover the various methods listed above by specifying particular forms of the Bregman divergence.

Our genriesz package implements generalized Riesz regression for ADML. Users specify the estimand through an evaluation oracle and choose a Riesz representer model and a Bregman generator. The package then constructs a generator-induced link function to deliver automatic regressor balancing, estimates the representer by convex optimization with $\ell_p$ regularization, and reports regression adjustment (RA), Riesz weighting (RW), augmented Riesz weighting (ARW), and targeted maximum likelihood estimation (TMLE)-style estimates with confidence intervals, optionally using cross-fitting Bang2005doublyrobust,vanderLaan2006targetedmaximum.

The flowchart of this automatic procedure is below, where each object will be defined in the subsequent sections (Figure (ref)):

itemize• the user specifies the parameter functional $m\left(W,\gamma_0\right)$ for the parameter of interest $\theta_0 \coloneqq {\mathbb{E}}\left[m\left(W,\gamma_0\right)\right]$, the Bregman generator $g$ for Riesz representer estimation, and the basis functions $\phi(X)$. • the package automatically computes the link function $\zeta\left(X,\phi(X)^\top\beta\right)$ that yields regressor balancing, and estimates the Riesz representer $\alpha_0$ and, when needed, the regression function $\gamma_0$. • the package outputs RA, RW, ARW, and TMLE-style estimators with standard errors and confidence intervals.

During this process, the user does not need to specify the analytic form of the Riesz representer $\alpha_0$ or the link function $\zeta$. These are constructed from the chosen basis and generator.

figure[figure omitted — 177 chars of source]

\paragraph{Relation to existing software.} The genriesz package complements established DML libraries such as DoubleML Bach2022doublemlpython and EconML Battocchi2019econml, as well as broader causal inference toolkits such as DoWhy Sharma2020dowhy and CausalML Chen2020causalml. These libraries offer rich sets of estimators for canonical causal models and heterogeneous treatment effects. In contrast, genriesz is estimand-and-balancing-centric: the user provides the functional $m$, the representer model, and the specific form of the Bregman divergence, while the package constructs generalized Riesz representers with a link function that yields balancing weights and returns the corresponding estimators. This focus also makes explicit the connection between Riesz representer fitting and balancing weights, including stable weights Zubizarreta2015stableweights and entropy balancing Hainmueller2012entropybalancing, as well as density ratio estimation methods available in specialized packages such as densratio Makiyama2019densratio.

Key Ingredients of Generalized Riesz Regression

Generalized Riesz regression connects Riesz representer estimation, covariate balancing, and debiased estimation through Bregman divergence minimization. This section summarizes the methodological and theoretical foundations that underlie the software design. Full technical details, proofs, and extensions are provided in Kato2026unifiedframework.

Linear Functionals, Riesz Representers, and Orthogonal Scores

Let $W \coloneqq (X,Y)\sim P$, where $X\in{\mathbb{R}}^d$ is a regressor and $Y\in{\mathbb{R}}$ is an outcome. Let $\gamma_0(x)\coloneqq{\mathbb{E}}\left[Y\mid X=x\right]$. We consider targets of the form

align[align omitted — 99 chars of source]

where $m\left(W,\gamma\right)$ is linear in $\gamma$. In genriesz, users supply $m$ as a callable that evaluates $\gamma$ at modified inputs, for example, by switching a treatment component.

Under standard conditions, there exists a Riesz representer $\alpha_0$ such that

align[align omitted — 163 chars of source]

The Riesz representer enters the Neyman orthogonal score

align[align omitted — 148 chars of source]

which plays a central role in debiased estimation and valid inference under weak conditions, with cross-fitting Chernozhukov2018doubledebiased.

table*[table* omitted — 3,911 chars of source]

Bregman-Riesz Objectives

Let $g\left(x,\alpha\right)$ be convex in the scalar $\alpha$ for each fixed $x$. The pointwise Bregman divergence is

align[align omitted — 203 chars of source]

Generalized Riesz regression estimates $\alpha_0$ by minimizing an empirical Bregman-Riesz objective of the form

align[align omitted — 344 chars of source]

where $\Omega$ is a regularizer, for example an $\ell_p$ penalty on a coefficient vector. Squared-loss generators recover Riesz regression Chernozhukov2021automaticdebiased and series Riesz representer Chen2015sievesemiparametric, while KL-type generators recover tailored-loss, calibrated estimation formulations Zhao2019covariatebalancing,Tan2019regularizedcalbrated and entropy balancing Hainmueller2012entropybalancing. The package provides built-in generators, squared distance (SQ), unnormalized KL (UKL) divergence, binary KL (BKL) divergence, Basu's power (BP) divergence, and PU families, and it also supports user-defined generators.

\paragraph{Dual coordinate and conjugate objective.} A key identity behind the implementation is the conjugacy relation

align[align omitted — 285 chars of source]

where $\mathcal{A}(x)$ is the domain of $\alpha$ given $x$. For differentiable generators, $\alpha^*\left(x,v\right)=(\partial_\alpha g)^{-1}\left(x,v\right)$. Using $v(x)\coloneqq \partial_\alpha g\left(x,\alpha(x)\right)$ and $g^*\left(x,v(x)\right)=\alpha(x)v(x)-g\left(x,\alpha(x)\right)$, the objective ((ref)) is equivalent, up to constants, to

align[align omitted — 192 chars of source]

This dual view is convenient for optimization and for deriving the balancing optimality conditions.

Automatic Regressor Balancing via Generator-Induced Links

To fit $\alpha_0$ in a model class, genriesz focuses on GLM-style parameterizations

align[align omitted — 125 chars of source]

where $\phi$ is a user-chosen basis and $\zeta$ is a link. A key choice is to set the link to the derivative of the generator,

align[align omitted — 183 chars of source]

so that the dual coordinate $v(x)=\partial_\alpha g\left(x,\alpha(x)\right)$ is linear in $\beta$. This is the package's notion of ARB.

In genriesz, the model is estimated by solving a convex program in $\beta$:

align[align omitted — 260 chars of source]

After estimating $\widehat{\beta}$, the fitted Riesz representer is $\widehat{\alpha}(x)=\zeta^{-1}\left(x,f_{\widehat{\beta}}(x)\right)$.

proposition[ARB implies balancing optimality conditions] Consider the model ((ref))--((ref)). Fix $q\in\left[1,\infty\right]$ and set $\Omega(\beta)=\frac{1}{q}\|\beta\|_q^q$. Let $\widehat{\beta}$ be any empirical minimizer of ((ref)) and define $\widehat{\alpha}(x)=\alpha_{\widehat{\beta}}(x)$. Under mild regularity conditions, the KKT conditions imply that there exist scalars $s_1,\dots,s_p$ such that \begin{align} \frac{1}{n}\sum^n_{i=1}\Big(\widehat{\alpha}(X_i)\phi_j(X_i)-m\left(W_i,\phi_j\right)\Big) + \lambda s_j = 0,\qquad s_j\in\partial\left(\frac{1}{q}\left|\beta_j\right|^q\right)|_{\beta_j=\widehat{\beta}_j}, \qquad j=1,\dots,p. \end{align} Consequently, the implied balancing condition takes the following explicit form: \begin{itemize} • If $q=1$, then $\left|s_j\right|\le 1$ and hence \begin{align} \left| \frac{1}{n}\sum^n_{i=1}\Big(\widehat{\alpha}(X_i)\phi_j(X_i)-m\left(W_i,\phi_j\right)\Big)\right| \le \lambda,\qquad j=1,\dots,p. \end{align} • If $q>1$, then $s_j=\operatorname{sign}\left(\widehat{\beta}_j\right)\left|\widehat{\beta}_j\right|^{q-1}$ and hence \begin{align} \left| \frac{1}{n}\sum^n_{i=1}\Big(\widehat{\alpha}(X_i)\phi_j(X_i)-m\left(W_i,\phi_j\right)\Big)\right| = \lambda\left|\widehat{\beta}_j\right|^{q-1},\qquad j=1,\dots,p. \end{align} \end{itemize} In particular, when $\lambda=0$ and the constraints are feasible, ((ref)) yields exact sample balancing.

Proposition (ref) is a software-relevant consequence of the duality theory in Kato2026unifiedframework. ARB ensures that the solver automatically enforces the correct balancing equations for the user-specified basis, even when the primal model ((ref)) is nonlinear in $\beta$, for example under KL-type links.

remark[Automatic link construction in genriesz] Users can supply analytic derivatives $\partial_\alpha g$ and $(\partial_\alpha g)^{-1}$. If they do not, genriesz approximates $\partial_\alpha g$ by finite differences and computes $(\partial_\alpha g)^{-1}$ by root finding. This allows rapid prototyping of new Bregman generators, at the cost of additional numerical care, such as domain constraints and branch selection.

remark{$\ell_1$ penalty and sparsity} The $\ell_1$ choice is useful because ((ref)) directly controls the maximum absolute moment imbalance by the single tuning parameter $\lambda$. In genriesz, this role is primarily about feasibility relaxation and stability of representer fitting, rather than recovering a sparse coefficient vector. In particular, using $\ell_1$ here should be understood as imposing an interpretable slack on balancing equations, not as a sparsity assumption on the true representer model.

Special Cases and Connections to Balancing Weights

The ARB construction ((ref)) makes explicit how generalized Riesz regression recovers classical balancing-weight estimators as dual solutions. Intuitively, the dual variable associated with the linear moment conditions corresponds to per-sample weights, and different generators $g$ induce different weight regularizers. Table (ref) summarizes common choices implemented in genriesz, see Kato2026unifiedframework for precise statements.

\paragraph{Domain constraints and branch selection.} Some generators require constraints such as $\left|\alpha\right|>C$ and $\alpha\in\left(0,1\right)$. genriesz exposes a branch_fn interface that selects a valid branch, for example treated versus control, so that the fitted representer respects the generator domain by construction.

Estimators: RA, RW, ARW, and TMLE

Given nuisance estimates $\widehat{\gamma}$ and $\widehat{\alpha}$, optionally cross-fitted, genriesz outputs four plug-in estimators, regression adjustment (RA), Riesz weighting (RW), augmented Riesz weighting (ARW), and targeted maximum likelihood estimation (TMLE)-style estimators, defined as follows:

align[align omitted — 522 chars of source]

In TMLE, $\widehat{\gamma}^{(1)}$ is a one-dimensional fluctuation update of $\widehat{\gamma}$ along the direction $\widehat{\alpha}$. We consider two likelihood choices, Gaussian and Bernoulli with a logit link, which lead to different update maps for $\widehat{\gamma}^{(1)}$. See Appendix (ref).

proposition[Asymptotic normality] Suppose $\widehat{\alpha}$ and $\widehat{\gamma}$ are obtained by cross-fitting and satisfy mean-square-error rate conditions such as $\|\widehat{\alpha}-\alpha_0\|_{L_2(P)}\|\widehat{\gamma}-\gamma_0\|_{L_2(P)}=o_p\left(n^{-1/2}\right)$ along with mild moment conditions. Then the ARW estimator in ((ref)) is asymptotically linear with influence function $\psi\left(W;\theta_0,\gamma_0,\alpha_0\right)$ in ((ref)), so that \[ \sqrt{n}\left(\widehat{\theta}^{\mathrm{ARW}}-\theta_0\right)\xrightarrow{{\mathrm{d}}}\mathcal{N}\left(0,\mathbb{V}\left[\psi\left(W;\theta_0,\gamma_0,\alpha_0\right)\right]\right)\right)}}. \] Analogous statements hold for the TMLE-style estimator.

\paragraph{Inference.} For each estimator, genriesz computes Wald-type standard errors from the empirical variance of the corresponding estimated influence-function scores. Cross-fitting is recommended in high-capacity settings Chernozhukov2018doubledebiased.

API Design and Software Architecture

The design goal of genriesz is to separate statistical intent, the functional $m$ and the representer class, from numerical implementation, basis matrices, generator derivatives, solvers, and cross-fitting.

Core Abstractions

The main entry point is grr_functional:

lstlisting[lstlisting omitted — 406 chars of source]

The returned result object stores point estimates, standard errors, confidence intervals, $p$-values, and optional out-of-fold nuisance predictions.

\paragraph{Estimators, cross-fitting, and outcome models.} The estimators argument selects which plug-in estimators to report. The built-in names follow the genriesz naming convention, "ra" for regression adjustment, "rw" for Riesz weighting, "arw" for augmented Riesz weighting, and "tmle" for the TMLE-style update. Cross-fitting is enabled by \texttt{cross_fit=True} with \texttt{folds} controlling the number of folds.

For RA, ARW, and TMLE, the package needs an outcome regression model $\widehat{\gamma}$. Users can control its construction by outcome_models, for example, "shared" fits a linear model using the same basis and regularization interface as the Riesz model, while "separate" fits the same outcome regression on a user-supplied outcome basis outcome_basis. Setting outcome_models="none" skips outcome modeling, then only RW is available.

Built-In Functionals and Wrappers

For common causal estimands with a binary treatment indicator $D$ stored in a column of $X$, the package provides convenience wrappers that call grr_functional with predefined $m$:

itemize• grr_ate: ATE for $X=\left[D,Z\right]$ with $m\left(W,\gamma\right)=\gamma\left(1,Z\right)-\gamma\left(0,Z\right)$. • grr_att: ATT for $X=\left[D,Z\right]$. One convenient linear functional is $m\left(W,\gamma\right)=\frac{D}{\pi_1}\left(\gamma\left(1,Z\right)-\gamma\left(0,Z\right)\right)a\left(0,Z\right)}$ with $\pi_1\coloneqq {\mathbb{E}}\left[D\right]$. In the wrapper, $\pi_1$ is estimated by $\widehat{\pi}_1\coloneqq \frac{1}{n}\sum^n_{i=1}D_i$. • grr_did: panel difference-in-difference (DID) implemented as ATT on $\Delta Y \coloneqq Y_1-Y_0$, where $Y_0$ and $Y_1$ are pre and post outcomes for the same units. • grr_ame: AME via derivatives, requiring a basis that implements $\frac{\partial \phi(x)}{\partial x_k}$.

These wrappers reduce boilerplate for standard workflows while still exposing basis and generator choices.

Basis Functions and Extensibility

The genriesz package treats the representer model as linear in a user-specified feature map $\phi(x)$. The core module includes polynomial features and treatment-interaction features for ATE and ATT workflows, as well as RKHS-style approximations via random Fourier features and Nystr\"om features.

Two optional modules expand the basis library: (i) genriesz.sklearn_basis provides a random forest leaf one-hot basis that turns a fitted ensemble into a sparse feature map, and (ii) genriesz.torch_basis wraps a PyTorch embedding network as a frozen feature map.

Finally, genriesz provides a kNN catchment basis KNNCatchmentBasis and matching utilities in genriesz.matching, which connect nearest-neighbor matching to squared-loss Riesz regression Kato2025nearestneighbor,Lin2023estimationbased.

\paragraph{Feature maps.} Keeping the solver linear in $\beta$ in the dual coordinate preserves convexity and the exact ARB optimality conditions in Proposition (ref). This suggests a practical pattern for neural pipelines, learn an embedding, freeze it, and then fit the representer in the induced feature space.

Optimization and Regularization

The GLM solver supports $\ell_p$ penalties for any $p\ge 1$.

\paragraph{Penalty interface.} For the Riesz model, set riesz_penalty="l2" for ridge, riesz_penalty="l1" for lasso, and riesz_penalty="lp" with riesz_p_norm=p for general $p\ge 1$. A shorthand such as riesz_penalty="l1.5" is also supported. When using the default linear outcome regression, the same interface is available via outcome_penalty and \texttt{outcome_p_norm}.

For $p\ge 1$, genriesz uses L-BFGS-B via SciPy, with a smooth approximation of the $\ell_1$ subgradient when $p=1$. The solver exposes iteration limits and tolerances, but defaults aim to work out-of-the-box for moderate feature dimensions.

algorithm[algorithm omitted — 771 chars of source]

Assumptions on Input Data

genriesz is designed around a simple data interface: X is a NumPy array of shape $\left(n,d\right)$ and Y is a vector of shape $\left(n,\right)$. The meaning of columns of X is left to the user and to the functional object.

\paragraph{Binary treatment conventions.} For wrappers for ATE, ATT, and DID estimation, the treatment indicator $D$ is assumed to be binary and stored at a known column index of X. The user can freely choose the remaining covariates $Z$.

\paragraph{Panel DID.} For grr_did, the package expects two outcome vectors Y0 and Y1 representing pre and post outcomes for the same units and computes $\Delta Y \coloneqq Y_1-Y_0$ internally.

\paragraph{Sampling and inference.} The default inference uses an i.i.d. approximation and Wald-type confidence intervals. As with standard DML software, clustered or dependent data require user-supplied adaptations, for example cluster-robust variance estimators, which are not yet part of the core release.

Examples and Reproducibility

The repository includes runnable scripts and Jupyter notebooks illustrating the workflows below.

ATE with Polynomial Features and UKL Generator

Let $X=\left[D,Z\right]$ and choose a polynomial basis with treatment interactions:

lstlisting[lstlisting omitted — 518 chars of source]

ATT and Panel DID

ATT and panel DID are available via wrapper functions:

lstlisting[lstlisting omitted — 559 chars of source]

Average Marginal Effects

Average marginal effects can be computed when the chosen basis implements derivatives:

lstlisting[lstlisting omitted — 337 chars of source]

Nearest-Neighbor Matching

Nearest-neighbor matching weights can be computed using the built-in matching Riesz method:

lstlisting[lstlisting omitted — 380 chars of source]

Project Development and Dependencies

\paragraph{Availability and installation.} The package is available on PyPI and can be installed by pip install genriesz. The package is also available on GitHub \url{https://github.com/MasaKat0/genriesz}. Optional extras enable scikit-learn-based and PyTorch-based bases. The documentation is hosted at \url{https://genriesz.readthedocs.io}.

\paragraph{Dependencies.} The core package depends only on NumPy and SciPy Harris2020arrayprogramming,Virtanen2020scipy10. Optional extras provide scikit-learn-based tree feature maps and PyTorch-based neural feature maps Pedregosa2011scikitlearn,Buitinck2013apidesign,Paszke2019pytorchan.

\paragraph{Quality control.} The project includes unit tests, type hints, and continuous integration to catch regressions. The software is released under the GNU General Public License v3.0 (GPL-3.0).

Comparison to Related Software

\paragraph{Debiased ML and causal ML libraries.} DoubleML Bach2022doublemlpython,Bach2024doublemlr and EconML Battocchi2019econml provide production-quality implementations of canonical DML estimators, with a focus on treatment effect estimation and, in EconML, heterogeneous effects. DoWhy Sharma2020dowhy provides an end-to-end causal inference interface centered on causal graphs. CausalML Chen2020causalml offers a broad suite of uplift modeling and heterogeneous effect estimators.

genriesz is complementary. It targets low-dimensional linear functionals with valid inference via orthogonalization and emphasizes representer estimation and balancing. This estimand-centric API makes it convenient to prototype new targets by defining only the oracle $m$, while keeping the representer fitting problem explicit.

\paragraph{Balancing-weight and density ratio toolkits.} Several toolkits implement specific balancing-weight estimators, such as entropy balancing or stable weights, for ATE-like problems. Density ratio estimation packages such as densratio Makiyama2019densratio implement methods such as uLSIF, RuLSIF, and KLIEP. genriesz connects these approaches to Riesz representer estimation and DML-style inference through a single interface.

Conclusion

The genriesz package provides an estimand-and-balancing-centric implementation of ADML via generalized Riesz regression under Bregman divergences. By unifying Riesz representer fitting, automatic regressor balancing, and debiased estimation behind modular abstractions, the package aims to make ADML practical for a wide range of causal and structural parameter estimation problems.