Genriesz: Debiased ML via Riesz Regression

• Genriesz is an open-source Python package for automatic debiased machine learning that implements generalized Riesz regression to estimate causal and structural parameters.
• It leverages the Riesz representation theorem and empirical Bregman divergence minimization to construct cross-fitted estimators with confidence intervals and p-values.
• The package features Automatic Regressor Balancing (ARB) and a modular architecture that simplifies basis selection, generator configuration, and estimator calculation.

genriesz is an open-source Python package for automatic debiased machine learning (DML) with generalized Riesz regression, designed for efficient estimation of causal and structural parameters using the Riesz representation theorem and empirical Bregman-divergence minimization (Kato, 19 Feb 2026). It targets parameters expressible as linear functionals of a regression function, automates construction of compatible representer estimators, and returns regression adjustment, Riesz weighting, augmented Riesz weighting, and TMLE-style estimators with cross-fitting, confidence intervals, and $p$-values. A central design principle is Automatic Regressor Balancing (ARB): given a Bregman generator $g$ and a representer model class, the package automatically constructs a compatible link function so that the generalized Riesz regression estimator satisfies balancing, or moment-matching, optimality conditions in a user-chosen basis (Kato, 19 Feb 2026).

1. Statistical target and Riesz-representer formulation

The package operates in an i.i.d. setting with data $W=(X,Y)\sim P$, feature $X\in\mathbb{R}^d$, outcome $Y\in\mathbb{R}$, and regression function $\gamma_0(x)=\mathbb{E}[Y\mid X=x]$ (Kato, 19 Feb 2026). The estimand is assumed to have the form

$\theta_0=\mathbb{E}\bigl[m(W,\gamma_0)\bigr],$

where the user-supplied oracle $m(W,\gamma)$ is linear in $\gamma$ (Kato, 19 Feb 2026).

By the Riesz representation theorem in $L_2(P_X)$, there exists a unique Riesz representer $g$0 such that

$g$1

for every square-integrable $g$2 (Kato, 19 Feb 2026). This representation is the organizing device of the package: once $g$3 is estimated, inference proceeds through a Neyman-orthogonal score,

$g$4

The score is orthogonal, or locally insensitive, to small perturbations of the nuisance functions $g$5 and $g$6, and averaging plug-in scores yields a root-$g$7–consistent, asymptotically normal estimator under weak rates plus cross-fitting (Kato, 19 Feb 2026).

This setup places genriesz within the DML tradition while centering estimation of the representer itself. A plausible implication is that the package is especially useful when the main technical bottleneck is not outcome modeling alone, but accurate recovery of the linear functional’s representer under structured balance constraints.

2. Generalized Riesz regression via Bregman divergences

The package implements generalized Riesz regression as estimation of $g$8 through a pointwise convex generator $g$9, convex in $W=(X,Y)\sim P$0 for each $W=(X,Y)\sim P$1, and its associated Bregman divergence

$W=(X,Y)\sim P$2

The empirical primal loss is, up to additive constants,

$W=(X,Y)\sim P$3

with $W=(X,Y)\sim P$4 a convex penalty, for example an $W=(X,Y)\sim P$5 penalty on coefficients (Kato, 19 Feb 2026).

The dual formulation is expressed in the coordinate $W=(X,Y)\sim P$6, using the convex conjugate

$W=(X,Y)\sim P$7

Then the dual loss is

$W=(X,Y)\sim P$8

which the package describes as often more convenient to optimize (Kato, 19 Feb 2026).

For GLM-style estimation, genriesz approximates the dual coordinate by a linear predictor in a user-specified feature map $W=(X,Y)\sim P$9: $X\in\mathbb{R}^d$0 The fitted coefficient vector solves

$X\in\mathbb{R}^d$1

after which

$X\in\mathbb{R}^d$2

is recovered (Kato, 19 Feb 2026).

This formulation unifies several procedures that the paper explicitly lists as special cases of the framework: covariate balancing, nearest-neighbor matching, calibrated estimation, and density ratio estimation (Kato, 19 Feb 2026).

3. Automatic Regressor Balancing

ARB is the package’s defining methodological principle. The key observation is that if the link is chosen as

$X\in\mathbb{R}^d$3

then the dual coordinate $X\in\mathbb{R}^d$4 is exactly the linear predictor $X\in\mathbb{R}^d$5, so no manual link construction is needed (Kato, 19 Feb 2026).

The balancing property arises from the KKT conditions. If

$X\in\mathbb{R}^d$6

then any minimizer $X\in\mathbb{R}^d$7 satisfies, for each feature $X\in\mathbb{R}^d$8,

$X\in\mathbb{R}^d$9

For $Y\in\mathbb{R}$0, corresponding to the lasso penalty, this yields

$Y\in\mathbb{R}$1

For $Y\in\mathbb{R}$2, the package description states that one gets exact matching up to the multiplier $Y\in\mathbb{R}$3. In particular, setting $Y\in\mathbb{R}$4 recovers exact balancing in the chosen basis (Kato, 19 Feb 2026).

This makes ARB more than an implementation convenience. It converts the choice of generator $Y\in\mathbb{R}$5, basis $Y\in\mathbb{R}$6, and penalty $Y\in\mathbb{R}$7 into an explicit balance geometry. A plausible implication is that the package can be interpreted as a convex balance engine whose estimand-specific content is supplied through the oracle $Y\in\mathbb{R}$8.

4. Software design and representer-model specification

The package exposes a core API centered on a callable target functional, a basis object, and a Bregman generator. The basic pattern is: $\gamma$2 (Kato, 19 Feb 2026)

The target functional is provided as a Python function or “oracle” of the form $\gamma$3 and all inner products with the basis $Y\in\mathbb{R}$9 are handled automatically (Kato, 19 Feb 2026).

The built-in representer model classes include polynomial bases, treatment-interaction bases, random-forest leaf encodings, neural embeddings, a nearest-neighbor catchment basis, and kernel-feature approximations via random Fourier or Nyström features (Kato, 19 Feb 2026). The specific classes listed are:

• PolynomialBasis(degree=k, include_bias=True)
• TreatmentInteractionBasis(base_basis)
• genriesz.sklearn_basis.RandomForestLeafBasis(forest_model)
• genriesz.torch_basis.TorchEmbeddingBasis(torch_nn_module)
• KNNCatchmentBasis(n_neighbors=M)
• kernel-feature approximations via random Fourier or Nyström features

The implemented Bregman generators include SquaredGenerator, UKLGenerator, BKLGenerator, BPGenerator, and PUGenerator; user-defined generators may be supplied by providing g(x,α), grad_g(x,α), inv_grad_g(x,v), and optionally grad2_g (Kato, 19 Feb 2026). The package summary also states that SquaredGenerator corresponds to $\gamma_0(x)=\mathbb{E}[Y\mid X=x]$0 (Kato, 19 Feb 2026).

The package is available on GitHub and on PyPI, and installation is via $\gamma$4 (Kato, 19 Feb 2026)

5. Estimators, orthogonalization, and inference

Given cross-fitted estimates $\gamma_0(x)=\mathbb{E}[Y\mid X=x]$1 and $\gamma_0(x)=\mathbb{E}[Y\mid X=x]$2, genriesz reports four estimators (Kato, 19 Feb 2026).

Regression Adjustment (RA) is

$\gamma_0(x)=\mathbb{E}[Y\mid X=x]$3

Riesz Weighting (RW) is

$\gamma_0(x)=\mathbb{E}[Y\mid X=x]$4

Augmented Riesz Weighting (ARW) is

$\gamma_0(x)=\mathbb{E}[Y\mid X=x]$5

The TMLE-style procedure updates $\gamma_0(x)=\mathbb{E}[Y\mid X=x]$6 by a one-dimensional fluctuation along $\gamma_0(x)=\mathbb{E}[Y\mid X=x]$7, then plugs in the updated regression (Kato, 19 Feb 2026). For Gaussian outcomes,

$\gamma_0(x)=\mathbb{E}[Y\mid X=x]$8

For binary outcomes, the package description states that one solves a small logistic score equation in $\gamma_0(x)=\mathbb{E}[Y\mid X=x]$9 (Kato, 19 Feb 2026).

Inference is based on $\theta_0=\mathbb{E}\bigl[m(W,\gamma_0)\bigr],$0-fold cross-fitting. The workflow described in the paper is:

$\gamma$5

Under the rate condition

$\theta_0=\mathbb{E}\bigl[m(W,\gamma_0)\bigr],$1

the ARW and TMLE estimators are asymptotically normal (Kato, 19 Feb 2026).

This combination of orthogonal scoring and automatic representer estimation suggests that genriesz is built to support both causal and structural functionals without requiring a bespoke orthogonal-score derivation for each new estimand, provided the estimand can be encoded through the oracle $\theta_0=\mathbb{E}\bigl[m(W,\gamma_0)\bigr],$2.

6. Canonical workflows and empirical examples

The package highlights representative workflows for the average treatment effect (ATE), average treatment effect on the treated (ATT), panel difference-in-differences, average marginal effects (AME), and nearest-neighbor matching (Kato, 19 Feb 2026).

For ATE with polynomial features and an unnormalized-KL generator, the documented pattern is: $\gamma$6 (Kato, 19 Feb 2026)

For ATT and panel DID, the package provides grr_att and grr_did with SquaredGenerator (Kato, 19 Feb 2026). For AME, grr_ame accepts a coordinate index for the derivative target: $\gamma$7 (Kato, 19 Feb 2026)

The nearest-neighbor matching mode uses riesz_method="nn_matching" and KNNCatchmentBasis, with the note that the basis is unused when outcome_models="none" (Kato, 19 Feb 2026).

These examples illustrate the intended division of labor. The estimand enters through a task-specific wrapper or oracle; the representer model is chosen through a basis; and the balance geometry is controlled by the generator and penalty. This suggests a deliberately modular architecture rather than separate implementations for each estimand.

7. Hyperparameters, scaling, and practical usage

The package description gives specific practical guidance on $\theta_0=\mathbb{E}\bigl[m(W,\gamma_0)\bigr],$3, basis design, and computational scaling (Kato, 19 Feb 2026). In $\theta_0=\mathbb{E}\bigl[m(W,\gamma_0)\bigr],$4-penalized fitting, $\theta_0=\mathbb{E}\bigl[m(W,\gamma_0)\bigr],$5 directly controls the maximum sample imbalance in each basis moment; setting $\theta_0=\mathbb{E}\bigl[m(W,\gamma_0)\bigr],$6 enforces near exact balancing. In $\theta_0=\mathbb{E}\bigl[m(W,\gamma_0)\bigr],$7-penalized fitting, $\theta_0=\mathbb{E}\bigl[m(W,\gamma_0)\bigr],$8 stabilizes the system of equations when the basis is high-dimensional or collinear (Kato, 19 Feb 2026).

For basis selection, low-degree polynomial or random Fourier features are said often to suffice in moderate dimensions. For very large $\theta_0=\mathbb{E}\bigl[m(W,\gamma_0)\bigr],$9, one can embed through a random-forest leaf encoder or a frozen neural network via genriesz.torch_basis; the final Riesz solver remains a convex $m(W,\gamma)$0-problem in $m(W,\gamma)$1 (Kato, 19 Feb 2026).

On scaling, the core relies on SciPy’s L-BFGS-B for smooth penalties and a smooth approximation of the $m(W,\gamma)$2 penalty, and the package description states that it comfortably handles thousands of parameters (Kato, 19 Feb 2026). For very large $m(W,\gamma)$3 or $m(W,\gamma)$4, suggested mitigations are subsampling, sparse features, and parallel cross-fitting (Kato, 19 Feb 2026).

The paper also includes empirical guidance from usage experiments: if results are unstable, increase $m(W,\gamma)$5 or reduce basis dimension; when the true representer is believed sparse in $m(W,\gamma)$6, $m(W,\gamma)$7 penalty plus small $m(W,\gamma)$8 can recover approximate sparsity; and the default squared generator is described as robust and often a good first try, while KL-type generators are suggested when entropy balancing or positivity constraints are desired (Kato, 19 Feb 2026).

In summary, genriesz presents a single coherent interface for defining an estimand $m(W,\gamma)$9, choosing a basis $\gamma$0, picking a Bregman generator $\gamma$1, and solving the resulting convex generalized Riesz regression problem (Kato, 19 Feb 2026). Its distinguishing contribution is the automation of representer estimation through ARB, coupled with cross-fitted DML-style inference for a family of causal and structural targets.