Regularization Of Treatment-Effect Interactions

Why Regularize Treatment-Effect Interactions?

A Q-model can include many ways in which a patient’s history modifies the effect of treatment: separate interactions for clinical measurements, indicators for categories, and nonlinear basis terms. With limited data, fitting all these terms without shrinkage can produce unstable treatment-effect contrasts. A small change in the sample may then change the estimated contrast enough to change the recommended action.

Regularization adds a penalty to the Q-model fitting criterion. A penalty can pull interaction estimates toward zero, or set some of them exactly to zero, while leaving the fitted baseline comparatively unrestricted. The objective is not simply a shorter list of coefficients: it is to stabilize the estimated treatment-effect function and, through it, the estimated decision. Shrinkage does not guarantee that the decision improves, and a penalized coefficient is not automatically a causal finding.

This parameterization separates the baseline from the terms of main interest: the treatment main effect \(\theta_0\) and the treatment-by-history interactions \(\theta\). One common choice is to penalize \(\theta\) while leaving the baseline parameters \(\beta\) and the treatment main effect \(\theta_0\) unpenalized. That choice targets effect modification rather than shrinking every component of the outcome model. Other penalty choices are possible, but they imply different shrinkage and must be stated explicitly.

At one stage of Q-learning, let \(Y_i^*\) denote the response being regressed on history and treatment; in a terminal stage it may be the observed outcome, while at an earlier stage it may be the pseudo-outcome constructed in backward Q-learning. With \(n\) observations, a common penalized criterion is

Here \(\lambda\geq 0\) controls the penalty. Ridge regression uses \(P(\theta)=\frac{1}{2}\sum_j\theta_j^2\); the lasso uses \(P(\theta)=\sum_j|\theta_j|\). Ridge generally shrinks coefficients toward zero without setting them exactly to zero. The lasso can set some estimates exactly to zero. The penalty is applied to fitted interaction coefficients, not directly to a patient’s estimated blip or decision.

What the Penalty Does: Two Derivations

The formulas are especially transparent when the interaction columns have been centered and scaled, and are orthogonal after accounting for the unpenalized columns. In this setting, the unpenalized least-squares estimates of the interaction coefficients are independent coordinate-wise summaries of the data. The derivations below show exactly how ridge and lasso transform those summaries.

Proof for ridge. Orthogonality separates the part of the least-squares loss involving the interaction coefficients into a sum of coordinate-wise terms. For coordinate \(j\), its loss, up to a constant not involving \(\theta_j\), is \(\frac{1}{2}(\theta_j-\tilde\theta_j)^2\). Adding the ridge penalty gives \(g(\theta_j)=\frac{1}{2}(\theta_j-\tilde\theta_j)^2+\frac{\lambda}{2}\theta_j^2\). Differentiating and setting the derivative to zero yields

Since \(g''(\theta_j)=1+\lambda>0\), this stationary point is the unique minimizer. Thus each coefficient is multiplied by the same factor, which lies between zero and one for \(\lambda\geq0\). This derivation assumes the stated orthogonality and scaling; with correlated columns, ridge estimates are generally coupled rather than coordinate-wise rescalings.

Proof for the lasso. For one coordinate, minimize \(g(\theta_j)=\frac{1}{2}(\theta_j-\tilde\theta_j)^2+\lambda|\theta_j|\). If \(\tilde\theta_j>\lambda\), a positive minimizer has derivative \(\theta_j-\tilde\theta_j+\lambda=0\), so \(\theta_j=\tilde\theta_j-\lambda>0\). If \(\tilde\theta_j<-\lambda\), a negative minimizer has derivative \(\theta_j-\tilde\theta_j-\lambda=0\), so \(\theta_j=\tilde\theta_j+\lambda<0\). For \(|\tilde\theta_j|\leq\lambda\), the subgradient of \(|\theta_j|\) at zero is the interval \([-1,1]\). The subgradient condition at zero is \(0\in-\tilde\theta_j+\lambda[-1,1]\), which holds exactly when \(|\tilde\theta_j|\leq\lambda\). The loss is strictly convex because of its quadratic term, so these solutions are the unique minimizers. Combining the cases gives the stated soft-threshold formula. \(\square\)

Reading Penalized Decisions Carefully

For any fitted penalized Q-model, the estimated decision is obtained by evaluating its fitted blip at the patient’s history and choosing the action with the larger fitted Q-value. The penalty acts upstream: it changes the coefficient estimates, which changes the contrast and may change the action. It does not directly encode the clinical cost of an action, nor does it make a decision robust merely because the coefficient vector is shorter.

The next example uses a simple orthogonal-design calculation to make this distinction concrete. Assume the treatment main effect is unpenalized and the two interaction columns are standardized and orthogonal, with unpenalized estimates \(0.50\) and \(-0.30\). Ridge with \(\lambda=0.25\) multiplies each interaction estimate by \(1/1.25=0.8\), giving \(0.40\) and \(-0.24\). For the history \(z=(1,1)\), if the fitted treatment main effect is \(0.10\), the unpenalized blip is \(0.10+0.50-0.30=0.30\), while the ridge blip is \(0.10+0.40-0.24=0.26\). Both fitted models select treatment. For \(z=(-1,1)\), the respective blips are \(0.10-0.50-0.30=-0.70\) and \(0.10-0.40-0.24=-0.54\); both select control. In these two histories shrinkage changes the contrast magnitude but not the action. Combined with the previous example, this illustrates that coefficient shrinkage can alter either the estimated contrast alone or the estimated decision as well.

Assumptions, Diagnostics, and Alternatives

For the stated least-squares optimization, the observations used in the regression must have well-defined, finite responses and predictors. The ridge and lasso derivations above require their additional orthogonality and scaling assumptions. More generally, the penalized criterion can be minimized with correlated features, but its coefficient solution does not have those simple formulas. The choice of penalty, the features included, and whether the baseline or treatment main effect is penalized should be specified before interpreting the resulting rule.

A fitted rule is intended to represent a causal treatment decision only under the relevant causal assumptions for Q-learning. In particular, the observed treatment and outcome data must be consistent with the specified treatment options, treatment assignment must be sequentially exchangeable given the measured history, and there must be positivity for those treatment options at histories in the target population. The Q-model must also provide a suitable approximation to the conditional mean being modeled. Regularization does not correct unmeasured confounding, a lack of treatment overlap, or a seriously misspecified Q-model. In a backward analysis, the pseudo-outcome at a stage is itself estimated, so the quality of the downstream fit matters as well.

Useful diagnostics focus on both prediction and decisions. Examine whether interaction columns have been centered and scaled appropriately; inspect the number and identity of retained interactions across plausible penalty values; and check whether estimated blips and decisions change substantially across those values. Compare fitted contrasts with the histories actually represented in the data, paying particular attention to regions with few observations or limited treatment overlap. A stable coefficient path is not a guarantee of a stable or valuable treatment rule, and a good average prediction score does not by itself establish good decisions. Tuning procedures should match the scientific goal: prediction-focused tuning does not automatically select the penalty that best estimates a decision rule.

The main alternatives have different trade-offs. An unpenalized Q-model avoids penalty-induced bias when the model is well specified and adequately supported by data, but can have high variance with many interactions. Ridge is useful when many correlated modifiers may each contribute modestly; it shrinks them together but usually retains them all. Lasso gives a sparse representation, but correlated features can compete, so which one is retained may be sensitive to the sample. Elastic net combines an absolute-value and a squared-coefficient penalty, potentially reducing that competition while retaining some sparsity. These methods regularize a chosen model; they do not replace careful model specification or causal identification.

Common Pitfalls and Limitations

• Penalizing every term by default. If the target is effect modification, penalizing baseline terms along with treatment interactions changes more than the modifiers. State exactly which coefficients enter the penalty.
• Reading a lasso zero as scientific absence. A zero is a property of the fitted penalized model at a chosen penalty level, not a test establishing no effect modification.
• Ignoring feature scale. An absolute-value penalty depends on coefficient scale. Rescaling a predictor changes the coefficient needed to represent the same fitted contrast, and hence changes the penalty’s effective strength.
• Assuming shrinkage must improve decisions. Shrinkage can reduce variance, but it can also introduce bias and move a decision boundary. Patients near that boundary can receive a different estimated action.
• Treating regularization as a cure for weak support. A penalty may yield finite coefficients when treatment data are sparse, but it cannot create information about treatment effects in unsupported histories.

The practical limitation is that the estimated rule inherits both the assumptions of Q-learning and the choices made in regularization. The selected features, penalty type, tuning value, and treatment-support pattern all influence the fitted blip. Report these choices and assess whether conclusions about the rule are sensitive to reasonable alternatives.