A new Python package, genriesz, automates debiased machine-learning estimation of causal parameters using the Riesz representation and Bregman-divergence minimization; it delivers regression-adjustment, weighting, and TMLE-style estimators with cross-fitting, confidence intervals and flexible model bases. The tool unifies covariate balancing, matching, density-ratio estimation and other approaches, making principled causal estimation more accessible for applied researchers.

Main Finding

genriesz is an open-source Python package that automates debiased machine learning for linear functionals (e.g., ATE, ATT, AME) via a unifying framework called generalized Riesz regression. The package (i) estimates Riesz representers by minimizing empirical Bregman divergences, (ii) constructs generator-induced link functions to automatically enforce regressor balancing (ARB), and (iii) returns a suite of debiased estimators (RA, RW, ARW, TMLE-style) with cross-fitting and Wald inference. genriesz unifies and operationalizes many existing approaches (covariate balancing, density-ratio estimation, entropy/stable balancing, matching) under a single, extensible API.

Key Points

• Conceptual contribution
  • Generalized Riesz regression: estimate the Riesz representer α0 by minimizing pointwise Bregman divergences; different Bregman generators recover many existing methods (squared loss, KL variants, BP, PU, etc.).
  • Automatic regressor balancing (ARB): set the GLM-style link ζ to the derivative ∂αg of the Bregman generator so that the dual coordinate (v = ∂αg(·, α(·))) is linear in basis coefficients; this yields automatic sample/basis balancing optimality conditions via KKT.
• Software contribution
  • genriesz implements generalized Riesz regression with a modular interface: user supplies the target functional m (as an oracle), a representer basis ϕ(X), and a Bregman generator (with optional analytic derivatives).
  • Provides built-in generators (SQ, UKL, BKL, BP, PU), and supports user-defined generators (finite-difference and inversion fallback when analytic derivatives not provided).
  • Supports multiple representer bases: polynomial, RKHS approximations (RFF/Nyström), random-forest leaf encodings, frozen neural embeddings (PyTorch), and a kNN catchment basis (matching).
  • Returns estimators: regression adjustment (RA), Riesz weighting (RW), augmented Riesz weighting (ARW), and TMLE-style updates, with cross-fitting, standard errors and CIs computed from estimated influence functions.
• Theoretical results implemented
  • KKT/duality shows ARB implies explicit balancing (moment-matching) conditions for the chosen basis; ℓ1 regularization yields a uniform slack bound on imbalances.
  • Asymptotic normality of ARW/TMLE-style estimators under usual product-rate mean-square-error conditions: ∥bα − α0∥2 · ∥bγ − γ0∥2 = oP(n−1/2).
• Practical details & defaults
  • Optimization: convex program in β (dual view uses g*), supports ℓp penalties (p ≥ 1), L-BFGS-B solver with smooth ℓ1 approximations where needed.
  • Convenience wrappers for common estimands: grr_ate, grr_att, grr_did, grr_ame.
  • Reproducible examples included (notebooks, scripts); package on PyPI and GitHub: https://github.com/MasaKat0/genriesz (pip install genriesz).

Data & Methods

• Target parameterization
  • Works with targets θ0 = E[m(W, γ0)] where m is linear in γ (user supplies m as an evaluation oracle that may modify X, e.g., switch treatment).
  • Riesz representer α0 satisfies E[m(W, γ)] = E[α0(X)γ(X)] for all suitable γ; α0 enters the Neyman-orthogonal score ψ(W; θ, γ, α) = m(W, γ) + α(X)(Y − γ(X)) − θ.
• Bregman-Riesz objective
  • Pointwise Bregman divergence BDg(α0‖α) with convex generator g(x, α).
  • Empirical objective (primal) bL(α) = (1/n)∑[−g(Xi, α(Xi)) + ∂αg(Xi, α(Xi))·α(Xi) − m(Wi, ∂αg(·, α(·)))] + λΩ(α).
  • Dual formulation: optimize over v(x) = ∂αg(x, α(x)) using conjugate g*(x, v), yielding a convex program in the linear predictor fβ = ϕ(x)⊤β when ζ = ∂αg and v = fβ.
• Model parameterization and ARB
  • GLM-style parametrization αβ(x) = ζ−1(x, ϕ(x)⊤β) with ζ = ∂αg (generator-induced link).
  • Solves convex program in β: bβ = argminβ (1/n)∑ [g*(Xi, fβ(Xi)) − m(Wi, fβ)] + λΩ(β); fitted α̂ = ζ−1(x, f̂β(x)).
  • KKT conditions yield balancing moment equations (exact when λ = 0 and feasible); ℓ1 penalty gives an interpretable slack bound on maximum imbalance.
• Estimators & inference
  • RA: plug-in mean of m(W, γ̂).
  • RW: sample average of α̂(Xi)Yi.
  • ARW: doubly robust-style estimator = average of α̂(Xi)(Yi − γ̂(Xi)) + m(Wi, γ̂).
  • TMLE-style: one-dimensional fluctuation of γ̂ along α̂ direction; package supports Gaussian and Bernoulli likelihood choices.
  • Cross-fitting recommended; asymptotic linearity under standard DML regularity (product-rate) yields consistent Wald CIs based on empirical influence-function variance.
• Implementation notes & limitations
  • Input data: X (n×d NumPy), Y (n). Binary treatment D assumed in a known column for ATE/ATT wrappers. DID expects paired pre/post outcomes.
  • Domain constraints/branch selection for generators with limited domains (e.g., α∈(0,1)) handled via branch functions.
  • Clustered/dependent data inference not built-in (user must supply adaptations).
  • Numerical care: if user omits analytic derivatives, finite differences/root finding are used; branch selection and domain constraints require attention.

Implications for AI Economics

• Low-friction causal estimation for economic applications of ML
  • genriesz automates a key, often bespoke, step in debiased ML pipelines: estimating the Riesz representer (weight or influence function) while enforcing balance on user-chosen bases. This reduces implementation barriers for economists who want principled, debiased inference for policy or mechanism evaluations that are linear functionals of a regression.
• Unifies and exposes modeling choices relevant to applied economics
  • The Bregman-generator abstraction makes explicit the tradeoffs between different weighting paradigms (squared-loss → matching/least-squares reweighting; KL/entropy → entropy balancing/tailored-loss; BP/PU → other robust estimators). Analysts can therefore select a divergence that reflects desired weight regularity (stability vs. exactness), with clear consequences for balancing constraints.
• Flexible basis choices enable domain-specific embeddings
  • Support for random-forest leaf encodings and frozen neural embeddings lets practitioners incorporate complex feature transforms or learned representations (e.g., marketplace features, textual embeddings, high-dimensional policy variables) while retaining convexity and balancing guarantees at the second-stage fit. This is useful when using ML-derived features but wanting reliable, debiased scalar estimates.
• Facilitates reproducible policy evaluation and model comparison
  • By standardizing the functional interface (user supplies m), genriesz makes it straightforward to compare alternative representer estimators, generators, and bases for the same economic estimand (ATE, ATT, AME, DID), aiding robustness checks and model selection in empirical work.
• Caveats and considerations for empirical economists
  • Inference assumes i.i.d. sampling; cluster-robust or dependent-data variance estimation is not built into the core release and must be supplied by users for many economic data settings.
  • Numerical stability and domain constraints: nonstandard generators (or omitted analytic derivatives) require careful branch selection and root-finding; regularization tuning interacts with balancing slack and variance, so empirical tuning and diagnostics are important.
  • License: GPL-3.0 may constrain embedding genriesz into some proprietary pipelines.
• Research and policy directions
  • genriesz lowers the cost of experimenting with alternative balancing/divergence choices in empirical economic research (e.g., comparing entropy balancing vs. SQ- or BP-based representers for policy evaluation).
  • Enables study of how representation learning (frozen embeddings) combined with automatic balancing affects bias/variance in causal estimates of algorithmic interventions—relevant to evaluations of recommender systems, pricing algorithms, and AI-driven market mechanisms.

Practical pointers - Use cross-fitting and ARW/TMLE for robust inference when using flexible representer and outcome learners. - Prefer ℓ1 penalty when you want interpretable control on maximum sample imbalance (λ governs slack). - Check generator domains and provide analytic ∂αg/(∂αg)−1 when possible for numerical stability. - For clustered data, compute cluster-robust influence-function variances externally.

Package & docs - GitHub: https://github.com/MasaKat0/genriesz - PyPI: pip install genriesz - Docs: https://genriesz.readthedocs.io

If you want, I can (a) extract the minimal code pattern for an ATE workflow with cross-fitting and recommended defaults, (b) summarize how each built-in generator maps to common econometrics estimators in one compact table, or (c) prepare a short checklist for empirical use (tuning, diagnostics, and cluster-robust inference). Which would be most useful?