The R package balnet fits regularized logistic propensity-score models, but replaces the ordinary Bernoulli likelihood with a covariate-balancing loss. The key engineering trick is that the loss can be presented to adelie as a custom GLM family, so Adelie’s pathwise group elastic-net solver can do the hard optimization work.
This memo sketches the optimization problem and a Python proof of concept in aipyw:
aipyw/balnet_adelie.py
The implementation uses a Python subclass of Adelie’s GlmBase64 to define the balnet calibration loss, then calls adelie.grpnet(...) for the regularization path.
2 One-arm balancing loss
For one arm, let \(y_i \in \{0,1\}\) indicate membership in the arm being modeled. Let
With intercept handling, this is equivalent to balancing a propensity-score weighted arm against the target arm in the calibration-loss geometry. With the lasso penalty, the KKT conditions become approximate balance:
\[
\left| X_j' r \right| \leq \lambda
\]
for inactive coordinates, and equality with sign corrections for active coordinates. This is the sense in which \(\lambda\) is a direct imbalance budget for standardized covariates.
For the ATE, balnet fits two one-arm models:
a treated model with \(y=W\), producing treated weights \(W_i / \hat e_1(X_i)\);
a control model with \(y=1-W\), producing control weights \((1-W_i)/(1-\hat e_0(X_i))\).
The Python proof of concept follows that structure.
4 How Adelie is used
The prototype defines:
class CBPSGlm(ad.glm.glm_base, ad.glm.GlmBase64):def gradient(self, eta, out): ...def hessian(self, eta, grad, out): ...def loss(self, eta): ...def inv_link(self, eta, out): ...
Then it calls:
state = ad.grpnet( X_standardized, CBPSGlm(y, weights=sample_weights), lmda_path=lambdas, alpha=alpha, intercept=True,)
The returned Adelie state contains a sparse coefficient path and intercept path. The wrapper unstandardizes coefficients back to the original covariate scale and provides balancing weights.
5 Sanity check: Adelie vs direct SciPy optimization
For a small one-arm lasso problem, solve the same objective once with Adelie and once with a generic SciPy optimizer. The generic optimizer is slow and not pathwise, but it is useful as a correctness check.
6 Path experiment: balance and effective sample size
Now fit the two-arm ATE path on a confounded simulated dataset. The diagnostic is the maximum absolute standardized mean difference against the full-sample covariate mean. Lower is better. Effective sample size tracks the variance cost of aggressive balancing.
This is a proof of concept, not a drop-in port yet.
What is working:
custom Python Adelie GLM for the balnet calibration loss;
lasso path through adelie.grpnet;
coefficient unstandardization;
two-arm ATE weights;
basic balance and ESS diagnostics;
tests against direct SciPy optimization and simulated balance reduction.
What still needs work for a faithful package-level port:
interpolation at arbitrary lambda values;
grouped penalties and path printing matching balnet;
ATT target scaling and diagnostics audited against the R package;
CV helpers using balance loss;
stronger numerical guards for poor overlap, where the calibration loss can diverge as \(\lambda \to 0\);
side-by-side tests against R balnet on fixed seeds.
The main takeaway is positive: Adelie’s Python API can solve the same mathematical problem by subclassing GlmBase64, so the Python port does not need to reimplement the coordinate-descent path solver.
