The exact contents of citations.db main_text.text for this paper — one flattened LaTeX string, title through conclusion, appendix excluded, unmodified except for removing email addresses. This is what our citation measures are computed over.
34,643 characters
genriesz: A Python Package for Automatic Debiased Machine Learning with Generalized Riesz Regression
\maketitle
\begin{abstract}
Efficient estimation of causal and structural parameters can be automated using the Riesz representation theorem and debiased machine learning (DML). We present \texttt{genriesz}, an open-source Python package that implements automatic DML and \emph{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 \emph{automatic regressor balancing} (ARB): given a Bregman generator $g$ and a representer model class, \texttt{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.
\end{abstract}
\section{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 \citep{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 \citep{Chernozhukov2018doubledebiased,Klaassen1987consistentestimation}.
Various methods have been proposed for Riesz representer estimation, such as Riesz regression \citep{Chernozhukov2021automaticdebiased,Chen2015sievesemiparametric}, covariate balancing \citep{Imai2013covariatebalancing}, and density ratio estimation \citep{Sugiyama2012densityratio}. \citet{Kato2026unifiedframework,Kato2025directbias} demonstrates a unifying perspective, \emph{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 \texttt{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 \citep{Bang2005doublyrobust,vanderLaan2006targetedmaximum}.
The flowchart of this automatic procedure is below, where each object will be defined in the subsequent sections (Figure~\ref{fig:concept}):
\begin{itemize}
\item[(i)] 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)$.
\item[(ii)] 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$.
\item[(iii)] the package outputs RA, RW, ARW, and TMLE-style estimators with standard errors and confidence intervals.
\end{itemize}
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.
\begin{figure}
\centering
\includegraphics[width=0.9\linewidth]{Figure/concept-eps-converted-to.pdf}
\caption{Flowchart of \texttt{genriesz}.}
\label{fig:concept}
\end{figure}
\paragraph{Relation to existing software.}
The \texttt{genriesz} package complements established DML libraries such as DoubleML \citep{Bach2022doublemlpython} and EconML \citep{Battocchi2019econml}, as well as broader causal inference toolkits such as DoWhy \citep{Sharma2020dowhy} and CausalML \citep{Chen2020causalml}.
These libraries offer rich sets of estimators for canonical causal models and heterogeneous treatment effects.
In contrast, \texttt{genriesz} is \emph{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 \citep{Zubizarreta2015stableweights} and entropy balancing \citep{Hainmueller2012entropybalancing}, as well as density ratio estimation methods available in specialized packages such as \texttt{densratio} \citep{Makiyama2019densratio}.
\section{Key Ingredients of Generalized Riesz Regression}
\label{sec:key}
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 \citet{Kato2026unifiedframework}.
\subsection{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
\begin{align}
\theta_0 \coloneqq {\mathbb{E}}\left[m\left(W,\gamma_0\right)\right],
\label{eq:theta}
\end{align}
where $m\left(W,\gamma\right)$ is linear in $\gamma$.
In \texttt{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
\begin{align}
{\mathbb{E}}\left[m\left(W,\gamma\right)\right] = {\mathbb{E}}\left[\alpha_0(X)\gamma(X)\right]
\qquad\text{for all suitable }\gamma.
\label{eq:riesz}
\end{align}
The Riesz representer enters the Neyman orthogonal score
\begin{align}
\psi\left(W;\theta,\gamma,\alpha\right)
\coloneqq
m\left(W,\gamma\right) + \alpha(X)\left(Y-\gamma(X)\right) - \theta,
\label{eq:score}
\end{align}
which plays a central role in debiased estimation and valid inference under weak conditions, with cross-fitting \citep{Chernozhukov2018doubledebiased}.
\begin{table*}[!t]
\caption{Correspondence among Bregman divergence losses, density ratio (DR) estimation methods, and Riesz representer (RR) estimation. RR estimation for ATE includes propensity score estimation and covariate balancing weights. In the table, $C \in {\mathbb{R}}$ denotes a constant that is determined by the problem and the loss function. The package also supports user-defined generators.}
\label{tab:special}
\begin{center}
\resizebox{0.9\linewidth}{!}{
\begin{tabular}{lll}
\hline
$g(\alpha)$ & \textbf{DR estimation} & \textbf{RR estimation} \\
\hline
\multirow{2}{*}{$\left(\alpha-C\right)^2$} & LSIF & SQ-Riesz regression \\
& \citep{Kanamori2009aleastsquares} & (\texttt{genriesz}) \\
& KuLSIF & Riesz regression (RieszNet and ForestRiesz) \\
& \citep{Kanamori2012statisticalanalysis} & \citep{Chernozhukov2021automaticdebiased,Chernozhukov2022riesznet} \\
& Hyv\"arinen score matching & RieszBoost \\
& \citep{Hyvarinen2005estimationof} & \citep{Lee2025rieszboost} \\
& & KRRR \\
& & \citep{Singh2024kernelridge} \\
& & Nearest neighbor matching \\
& & \citep{Lin2023estimationbased} \\
& & Causal forest and generalized random forest \\
& & \citep{Wager2018estimationinference,Athey2019generalizedrandom} \\
\multicolumn{3}{c}{\textbf{Dual solution with a linear link function}} \\
& Kernel mean matching & Sieve Riesz representer\\
& \citep{Gretton2009covariateshift} & \citep{Chen2015sievesemiparametric,Chen2015sievewald} \\
& & Stable balancing weights \\
& & \citep{Zubizarreta2015stableweights,BrunsSmith2025augmentedbalancing} \\
& & Approximate residual balancing\\
& & \citep{Athey2018approximateresidual} \\
& & Covariate balancing by SVM \\
& & \citep{Tarr2025estimatingaverage} \\
& & Distributional balancing \\
& & \citep{Santra2026distributionalbalancing} \\
\hline
\multirow{2}{*}{$\big(\left(\left|\alpha\right|-C\right)\big)\log \big(\left(\left|\alpha\right|-C\right)\big) - \left|\alpha\right|$} & UKL divergence minimization & UKL-Riesz regression\\
& \citep{Nguyen2010estimatingdivergence} & (\texttt{genriesz})\\
& & Tailored loss minimization ($\alpha=\beta=-1$)\\
& & \citep{Zhao2019covariatebalancing} \\
& & Calibrated estimation \\
& & \citep{Tan2019regularizedcalbrated} \\
\multicolumn{3}{c}{\textbf{Dual solution with a logistic or log link function}} \\
& KLIEP & Entropy balancing weights \\
& \citep{Sugiyama2008directimportance} & \citep{Hainmueller2012entropybalancing} \\
\hline
$\left(\left|\alpha\right|-C\right)\log \big(\left(\left|\alpha\right|-C\right)\big) - \left(\left|\alpha\right|+C\right)\log \big(\left(\left|\alpha\right|+C\right)\big)$ & BKL divergence minimization & BKL-Riesz regression \\
& \citep{Qin1998inferencesfor} & (\texttt{genriesz}) \\
& TRE & MLE of the propensity score \\
& \citep{Rhodes2020telescopingdensityratio} & (Standard approach) \\
& & Tailored loss minimization ($\alpha=\beta=0$)\\
& & \citep{Zhao2019covariatebalancing}\\
\hline
$\frac{\big(\left(\left|\alpha\right|-C\right)\big)^{1 + \omega} - \big(\left(\left|\alpha\right|-C\right)\big)}{\omega} - \left(\left|\alpha\right|-C\right)$ & BP divergence minimization & BP-Riesz regression\\
for some $\omega \in \left(0,\infty\right)$ & \citep{Sugiyama2011densityratio} & (\texttt{genriesz})\\
\hline
$C\log\left(1-\alpha\right) + C\alpha\left(\log\left(\alpha\right)-\log\left(1-\alpha\right)\right)t(1-\alpha\right)}$ & PU learning & PU-Riesz regression\\
for $\alpha \in \left(0,1\right)$ & \citep{duPlessis2015convexformulation} & (\texttt{genriesz})\\
& Nonnegative PU learning & \\
& \citep{Kiryo2017positiveunlabeledlearning} & \\
\hline
General formulation by Bregman & Density-ratio matching & Generalized Riesz regression \\
divergence minimization & \citep{Sugiyama2011densityratio} & (\texttt{genriesz}) \\
& D3RE & \\
& \citep{Kato2021nonnegativebregman} & \\
\hline
\end{tabular}}
\vspace{-5mm}
\end{center}
\end{table*}
\subsection{Bregman-Riesz Objectives}
Let $g\left(x,\alpha\right)$ be convex in the scalar $\alpha$ for each fixed $x$.
The pointwise Bregman divergence is
\begin{align}
\text{BD}_g\left(\alpha_0(x)\| \alpha(x)\right)
\coloneqq
g\left(x,\alpha_0(x)\right)-g\left(x,\alpha(x)\right)-\partial_\alpha g\left(x,\alpha(x)\right)\left(\alpha_0(x) - \alpha(x)\right).
\end{align}
Generalized Riesz regression estimates $\alpha_0$ by minimizing an empirical Bregman-Riesz objective of the form
\begin{align}
\widehat{L}\left(\alpha\right)
\coloneqq
\frac{1}{n}\sum^n_{i=1}\Big(
- g\left(X_i,\alpha(X_i)\right) + \partial_\alpha g\left(X_i,\alpha(X_i)\right)\alpha(X_i)
- m\left(W_i, \partial_\alpha g\left(\cdot,\alpha\left(\cdot\right)\right)t\right)}\right)\Big)\right)t\right)}}
}
+
\lambda \Omega\left(\alpha\right),
\label{eq:grr_obj}
\end{align}
where $\Omega$ is a regularizer, for example an $\ell_p$ penalty on a coefficient vector.
Squared-loss generators recover Riesz regression \citep[primal,][]{Chernozhukov2021automaticdebiased} and series Riesz representer \citep[dual,][]{Chen2015sievesemiparametric},
while KL-type generators recover tailored-loss, calibrated estimation formulations \citep[primal,][]{Zhao2019covariatebalancing,Tan2019regularizedcalbrated} and entropy balancing \citep[dual,][]{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
\begin{align}
g^*\left(x,v\right)\coloneqq \sup_{\alpha\in \mathcal{A}(x)}\left\{\alpha v - g\left(x,\alpha\right)\right\},
\qquad
\alpha^*\left(x,v\right)\coloneqq \operatorname*{arg\,max}_{\alpha\in \mathcal{A}(x)}\left\{\alpha v - g\left(x,\alpha\right)\right\},
\label{eq:conjugate}
\end{align}
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{eq:grr_obj}) is equivalent, up to constants, to
\begin{align}
\widehat{L}^*\left(v\right)
\coloneqq
\frac{1}{n}\sum^n_{i=1}\Big(
g^*\left(X_i,v(X_i)\right) - m\left(W_i,v\right)
\Big)
+
\lambda \Omega^*\left(v\right).
\label{eq:grr_obj_dual}
\end{align}
This dual view is convenient for optimization and for deriving the balancing optimality conditions.
\subsection{Automatic Regressor Balancing via Generator-Induced Links}
To fit $\alpha_0$ in a model class, \texttt{genriesz} focuses on GLM-style parameterizations
\begin{align}
\alpha_\beta(x)=\zeta^{-1}\left(x, f_\beta(x)\right),
\qquad
f_\beta(x)=\phi(x)^\top \beta,
\label{eq:glm_alpha}
\end{align}
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,
\begin{align}
\zeta\left(x,\alpha\right)=\partial_\alpha g\left(x,\alpha\right),
\qquad
\zeta^{-1}\left(x,\cdot\right)=(\partial_\alpha g\left(x,\cdot\right))^{-1},
\label{eq:arb_link}
\end{align}
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 \texttt{genriesz}, the model is estimated by solving a convex program in $\beta$:
\begin{align}
\widehat{\beta}\coloneqq \operatorname*{arg\,min}_{\beta\in{\mathbb{R}}^p}\left\{
\frac{1}{n}\sum^n_{i=1}\Big( g^*\left(X_i,f_\beta(X_i)\right) - m\left(W_i, f_\beta\right) \Big)
+
\lambda \Omega\left(\beta\right)
\right\}.
\label{eq:grr_obj_beta}
\end{align}
After estimating $\widehat{\beta}$, the fitted Riesz representer is $\widehat{\alpha}(x)=\zeta^{-1}\left(x,f_{\widehat{\beta}}(x)\right)$.
\begin{proposition}[ARB implies balancing optimality conditions]
\label{prop:arb}
Consider the model (\ref{eq:glm_alpha})--(\ref{eq:arb_link}).
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{eq:grr_obj_beta}) 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.
\label{eq:kkt_balance_general}
\end{align}
Consequently, the implied balancing condition takes the following explicit form:
\begin{itemize}
\item 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.
\label{eq:kkt_balance}
\end{align}
\item 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.
\label{eq:kkt_balance_lq}
\end{align}
\end{itemize}
In particular, when $\lambda=0$ and the constraints are feasible, (\ref{eq:kkt_balance_general}) yields exact sample balancing.
\end{proposition}
\noindent Proposition~\ref{prop:arb} is a software-relevant consequence of the duality theory in \citet{Kato2026unifiedframework}.
ARB ensures that the solver automatically enforces the correct balancing equations for the user-specified basis, even when the primal model (\ref{eq:glm_alpha}) is nonlinear in $\beta$, for example under KL-type links.
\begin{remark}[Automatic link construction in \texttt{genriesz}]
Users can supply analytic derivatives $\partial_\alpha g$ and $(\partial_\alpha g)^{-1}$.
If they do not, \texttt{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.
\end{remark}
\begin{remark}{$\ell_1$ penalty and sparsity}
The $\ell_1$ choice is useful because (\ref{eq:kkt_balance}) directly controls the maximum absolute moment imbalance by the single tuning parameter $\lambda$.
In \texttt{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.
\end{remark}
\subsection{Special Cases and Connections to Balancing Weights}
\label{sec:special}
The ARB construction (\ref{eq:arb_link}) 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{tab:special} summarizes common choices implemented in \texttt{genriesz}, see \citet{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)$.
\texttt{genriesz} exposes a \texttt{branch\_fn} interface that selects a valid branch, for example treated versus control, so that the fitted representer respects the generator domain by construction.
\subsection{Estimators: RA, RW, ARW, and TMLE}
\label{sec:estimators}
Given nuisance estimates $\widehat{\gamma}$ and $\widehat{\alpha}$, optionally cross-fitted, \texttt{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:
\begin{align}
\widehat{\theta}^{\mathrm{RA}} &\coloneqq \frac{1}{n}\sum^n_{i=1}m\left(W_i, \widehat{\gamma}\right), \\
\widehat{\theta}^{\mathrm{RW}} &\coloneqq \frac{1}{n}\sum^n_{i=1}\widehat{\alpha}(X_i)Y_i,\\
\widehat{\theta}^{\mathrm{ARW}} &\coloneqq \frac{1}{n}\sum^n_{i=1}\Big(\widehat{\alpha}(X_i)\left(Y_i-\widehat{\gamma}(X_i)\right) + m\left(W_i,\widehat{\gamma}\right)\Big),\\
\widehat{\theta}^{\mathrm{TMLE}}
&\coloneqq
\frac{1}{n}\sum^n_{i=1}m\left(W_i,\widehat{\gamma}^{(1)}\right),
\label{eq:estimators_main}
\end{align}
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{appdx:tmle}.
\begin{proposition}[Asymptotic normality]
\label{prop:an}
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{eq:estimators_main}) is asymptotically linear with influence function $\psi\left(W;\theta_0,\gamma_0,\alpha_0\right)$ in (\ref{eq:score}), 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.
\end{proposition}
\paragraph{Inference.}
For each estimator, \texttt{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 \citep{Chernozhukov2018doubledebiased}.
\section{API Design and Software Architecture}
\label{sec:api}
The design goal of \texttt{genriesz} is to separate \emph{statistical intent}, the functional $m$ and the representer class, from \emph{numerical implementation}, basis matrices, generator derivatives, solvers, and cross-fitting.
\subsection{Core Abstractions}
The main entry point is \texttt{grr\_functional}:
\begin{lstlisting}[language=Python]
from genriesz import grr_functional
res = grr_functional(
X=X, Y=Y,
m=m, # Functional object
basis=basis, # Feature map phi(x)
generator=gen, # BregmanGenerator (or g=..., grad_g=..., inv_grad_g=..., grad2_g=...)
cross_fit=True, folds=5,
estimators=("ra","rw","arw","tmle"),
)
print(res.summary_text())
\end{lstlisting}
\noindent 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 \texttt{estimators} argument selects which plug-in estimators to report.
The built-in names follow the \texttt{genriesz} naming convention, \texttt{"ra"} for regression adjustment, \texttt{"rw"} for Riesz weighting, \texttt{"arw"} for augmented Riesz weighting, and \texttt{"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 \texttt{outcome\_models}, for example, \texttt{"shared"} fits a linear model using the same basis and regularization interface as the Riesz model, while \texttt{"separate"} fits the same outcome regression on a user-supplied outcome basis \texttt{outcome\_basis}.
Setting \texttt{outcome\_models="none"} skips outcome modeling, then only RW is available.
\subsection{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 \texttt{grr\_functional} with predefined $m$:
\begin{itemize}
\item \texttt{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)$.
\item \texttt{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$.
\item \texttt{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.
\item \texttt{grr\_ame}: AME via derivatives, requiring a basis that implements $\frac{\partial \phi(x)}{\partial x_k}$.
\end{itemize}
These wrappers reduce boilerplate for standard workflows while still exposing basis and generator choices.
\subsection{Basis Functions and Extensibility}
The \texttt{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) \texttt{genriesz.sklearn\_basis} provides a random forest leaf one-hot basis that turns a fitted ensemble into a sparse feature map,
and (ii) \texttt{genriesz.torch\_basis} wraps a PyTorch embedding network as a frozen feature map.
Finally, \texttt{genriesz} provides a kNN catchment basis \texttt{KNNCatchmentBasis} and matching utilities in \texttt{genriesz.matching}, which connect nearest-neighbor matching to squared-loss Riesz regression \citep{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{prop:arb}.
This suggests a practical pattern for neural pipelines, learn an embedding, freeze it, and then fit the representer in the induced feature space.
\subsection{Optimization and Regularization}
The GLM solver supports $\ell_p$ penalties for any $p\ge 1$.
\paragraph{Penalty interface.}
For the Riesz model, set \texttt{riesz\_penalty="l2"} for ridge, \texttt{riesz\_penalty="l1"} for lasso, and \texttt{riesz\_penalty="lp"} with \texttt{riesz\_p\_norm=p} for general $p\ge 1$.
A shorthand such as \texttt{riesz\_penalty="l1.5"} is also supported.
When using the default linear outcome regression, the same interface is available via \texttt{outcome\_penalty} and \texttt{outcome\_p\_norm}.
For $p\ge 1$, \texttt{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.
\begin{algorithm}[t]
\caption{Simplified workflow implemented by \texttt{grr\_functional}}
\label{alg:workflow}
\begin{algorithmic}[1]
\REQUIRE Data $(X_i,Y_i)_{i=1}^n$, functional $m$, basis $\phi$, generator $g$, folds $K$.
\FOR{$k=1$ to $K$ (if cross-fitting; else $K=1$)}
\STATE Fit representer model $\widehat{\alpha}^{(-k)}$ on training fold via generalized Riesz regression.
\STATE Fit outcome model $\widehat{\gamma}^{(-k)}$ on training fold if RA, ARW, or TMLE is requested.
\STATE Predict $\widehat{\alpha}_i$ and $\widehat{\gamma}_i$ on held-out fold.
\ENDFOR
\STATE Compute RA, RW, ARW, and TMLE estimators.
\STATE Compute standard errors and confidence intervals from influence-function scores.
\RETURN Estimates and inference summary.
\end{algorithmic}
\end{algorithm}
\section{Assumptions on Input Data}
\label{sec:assumptions}
\texttt{genriesz} is designed around a simple data interface: \texttt{X} is a NumPy array of shape $\left(n,d\right)$ and \texttt{Y} is a vector of shape $\left(n,\right)$.
The meaning of columns of \texttt{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 \texttt{X}.
The user can freely choose the remaining covariates $Z$.
\paragraph{Panel DID.}
For \texttt{grr\_did}, the package expects two outcome vectors \texttt{Y0} and \texttt{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.
\section{Examples and Reproducibility}
\label{sec:examples}
The repository includes runnable scripts and Jupyter notebooks illustrating the workflows below.
\subsection{ATE with Polynomial Features and UKL Generator}
Let $X=\left[D,Z\right]$ and choose a polynomial basis with treatment interactions:
\begin{lstlisting}[language=Python]
from genriesz import (
grr_ate,
PolynomialBasis, TreatmentInteractionBasis,
UKLGenerator
)
psi = PolynomialBasis(degree=2, include_bias=True)
phi = TreatmentInteractionBasis(base_basis=psi)
gen = UKLGenerator(C=1.0, branch_fn=lambda x: int(x[0] == 1.0)).as_generator()
res = grr_ate(
X=X, Y=Y,
basis=phi,
generator=gen,
cross_fit=True, folds=5,
estimators=("ra","rw","arw","tmle"),
riesz_penalty="l2", riesz_lam=1e-3,
)
print(res.summary_text())
\end{lstlisting}
\subsection{ATT and Panel DID}
ATT and panel DID are available via wrapper functions:
\begin{lstlisting}[language=Python]
from genriesz import (
grr_att, grr_did,
PolynomialBasis, TreatmentInteractionBasis,
SquaredGenerator
)
psi = PolynomialBasis(degree=2, include_bias=True)
phi = TreatmentInteractionBasis(base_basis=psi)
gen = SquaredGenerator().as_generator()
res_att = grr_att(
X=X, Y=Y,
basis=phi,
generator=gen,
cross_fit=True, folds=5,
estimators=("arw","tmle"),
)
res_did = grr_did(
X=X, Y0=Y0, Y1=Y1,
basis=phi,
generator=gen,
cross_fit=True, folds=5,
estimators=("arw","tmle"),
)
\end{lstlisting}
\subsection{Average Marginal Effects}
Average marginal effects can be computed when the chosen basis implements derivatives:
\begin{lstlisting}[language=Python]
from genriesz import grr_ame, PolynomialBasis, SquaredGenerator
phi = PolynomialBasis(degree=2, include_bias=True)
res_ame = grr_ame(
X=X, Y=Y,
coordinate=2,
basis=phi,
generator=SquaredGenerator().as_generator(),
cross_fit=True, folds=5,
estimators=("ra","rw","arw","tmle"),
)
\end{lstlisting}
\subsection{Nearest-Neighbor Matching}
Nearest-neighbor matching weights can be computed using the built-in matching Riesz method:
\begin{lstlisting}[language=Python]
from genriesz import grr_ate, PolynomialBasis
# A basis is required by the API. When outcome_models="none", it is not used.
basis = PolynomialBasis(degree=1, include_bias=True)
res_match = grr_ate(
X=X, Y=Y,
basis=basis,
riesz_method="nn_matching",
M=1,
cross_fit=False,
outcome_models="none",
estimators=("rw",),
)
\end{lstlisting}
\section{Project Development and Dependencies}
\label{sec:dev}
\paragraph{Availability and installation.}
The package is available on PyPI and can be installed by \texttt{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 \citep{Harris2020arrayprogramming,Virtanen2020scipy10}.
Optional extras provide scikit-learn-based tree feature maps and PyTorch-based neural feature maps \citep{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).
\section{Comparison to Related Software}
\label{sec:related}
\paragraph{Debiased ML and causal ML libraries.}
\texttt{DoubleML} \citep{Bach2022doublemlpython,Bach2024doublemlr} and \texttt{EconML} \citep{Battocchi2019econml} provide production-quality implementations of canonical DML estimators, with a focus on treatment effect estimation and, in \texttt{EconML}, heterogeneous effects.
\texttt{DoWhy} \citep{Sharma2020dowhy} provides an end-to-end causal inference interface centered on causal graphs.
\texttt{CausalML} \citep{Chen2020causalml} offers a broad suite of uplift modeling and heterogeneous effect estimators.
\texttt{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 \texttt{densratio} \citep{Makiyama2019densratio} implement methods such as uLSIF, RuLSIF, and KLIEP.
\texttt{genriesz} connects these approaches to Riesz representer estimation and DML-style inference through a single interface.
\section{Conclusion}
The \texttt{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.
\bibliography{arXiv2.bbl}
\bibliographystyle{tmlr}
\newpage,