Question about `ate_inference()`

[victor5as]

Hello,

I have a question regarding the standard error computation in ate_inference() when X is not None.

For a linear CATE model, $\tau(x) = \psi(x)'\beta$, the ATE is $\theta = \Psi'\beta$, where $\Psi = E[\psi(X)]$.

The point estimate $\hat{\theta} = \hat{\Psi}'\hat{\beta}$ (where $\hat{\Psi} = \frac{1}{n}\sum_{i=1}^n \psi(X_i)$ and $\hat{\beta}$ are the final model coefficients) aligns with ate_inference().

However, it appears that ate_inference() calculates the standard error as $\sqrt{\hat{\Psi}' \hat{V} \hat{\Psi}}$, treating $\hat{\Psi}$ as a constant and disregarding its sampling variation.

Wouldn't it be preferable to use the Delta Method for standard error calculation in this case?

I'd appreciate your insights on this matter.

Thank you!

[vsyrgkanis]

I think that’s a great point worth seriously considering @victor5as !

Currently the ate inference has the semantics of inference on the “fixed design” or “in sample” population of X’s provided by the user.

But I agree maybe it’s better to have the “population” inference semantics as default and maybe have a flag for “in-sample” that allows for the two variants.

@kbattocchi this might be a top priority issue that @victor5as is raising. To the least we should document in the docstring of ate inference that this is fixed design or in-sample ate inference.

@victor5as do you have any appetite for implementing and contributing the population correction?

[vsyrgkanis]

The reason we went with this approach is that the ate_inference api accepts arbitrary population X as input.

This complicates a bit things as it is not clear what the cross term in the influence function would look like. I believe, one would need to know whether the X’s provided here are the same as the X’s used during training to do an exact variance calculation as the cross term requires such knowledge when the final model is a best linear projection and not a well specified linear model.

[vsyrgkanis]

Most probably the best would be to use a different api, similar to ate_inference_ that we have in causalforestdml, which performs ate inference on the training data, hence the above ambiguities are not an issue. This can be calculated at training time, so that the training data dont need to be stored.

[victor5as]

Hi @vsyrgkanis,

Thank you so much for your quick reply!

That makes sense about the "fixed design" framework and about using different X.

I think when X_train (used to fit $\hat{\beta}$) and X_eval (used to fit $\hat{\Psi}$) are independent, the off-diagonal blocks of the var-cov matrix are 0, so you can estimate the variance of the ATE (or the "predicted projection" if the model is misspecified) with:
$$N_{train}^{-1} \hat{\Psi}' \hat{\Sigma}_\beta \hat{\Psi} + N_{eval}^{-1} \hat{\beta}' \hat{\Sigma}_\Psi \hat{\beta}$$
but I need to double check that.

And I agree that things are complicated if you want to evaluate on an X that overlaps but is not exactly the same as the ones in training (e.g. if you want to get an ATT by averaging only over T==1).

[victor5as]

Here is the code that I used to get the errors manually (averaging over the same X used to fit the final model).

It's not optimized and doesn't support multiple treatments, sample weights, etc. but hopefully it will be helpful if you end up deciding to implement the ate_inference_:

from statsmodels.api import add_constant

def manual_ate_inference(estimator: Union[LinearDML, LinearDRLearner]) -> tuple[float, float]:
    """
    Performs manual average treatment effect (ATE) inference for a fitted LinearDML or LinearDRLearner estimator.
    
    This function calculates the ATE point estimate and standard error on the training sample in which the estimator was fitted. 
    Inference is based on the delta method.
    
    Parameters
    ----------
    estimator : Union[LinearDML, LinearDRLearner]
        A fitted estimator object that must have cache_values=True during fitting.
    
    Returns
    -------
    tuple[float, float]
        ate_point : float
            The point estimate of the average treatment effect
        ate_se : float
            The standard error of the ATE estimate
    """
    
    if estimator._cached_values is None:
        raise ValueError("The estimator needs to be a fitted LinearDML or LinearDRLearner with cache_values=True.")

    X = estimator._cached_values.X
    n = estimator._cached_values.Y.shape[0]
    
    # Get psi(X) (transformed features):
    if estimator.featurizer_ is not None:
        psi = estimator.featurizer_.transform(X) 
    elif estimator.fit_cate_intercept:
        if X is not None:
            psi = add_constant(X)
        else:
            psi = np.ones((n, 1))
    else:
        psi = X

    psi_mean = np.mean(psi, axis=0)

    # Re-do final linear regression:
    if isinstance(estimator, LinearDML):
        Yres, Dres = estimator._cached_values.nuisances
        lhs = Yres
        rhs = Dres.reshape(-1, 1) * psi
    elif isinstance(estimator, LinearDRLearner):
        Y1_DR, Y0_DR = (estimator._cached_values.nuisances[0][:, 1], 
                        estimator._cached_values.nuisances[0][:, 0])
        lhs = Y1_DR - Y0_DR
        rhs = psi
    else:
        raise NotImplementedError("The estimator needs to be a fitted LinearDML or LinearDRLearner.")

    XX = np.matmul(rhs.T, rhs)
    XY = np.matmul(rhs.T, lhs)
    beta, _, _, _ = np.linalg.lstsq(XX, XY, rcond=None)

    ate_point = np.dot(psi_mean, beta)
    
    # Joint influence function for (beta, psi_mean):
    influence_beta = np.einsum('ij,kj,k->ki', np.linalg.pinv(XX / n), rhs, lhs - np.matmul(rhs, beta))
    influence_psi = psi - psi_mean
    influence_both = np.hstack([influence_beta, influence_psi])

    # Delta method:
    variance = np.cov(influence_both, rowvar=False)
    jacobian = np.hstack([psi_mean, beta]).T
    ate_se = np.sqrt(np.dot(jacobian, np.matmul(variance, jacobian)) / n)

    return ate_point, ate_se

The results coincide with ate_inference when X=None as expected, and it seems to be capturing the variance correctly in simulations.

[vsyrgkanis]

One option is to implement the ate_ and att_ api (with the underscore) that we implemented in causalforestdml, but for every linearmodelfinal method and use the influence approach as you use. This would provide inference on ate and att over the training data, as what you want here.

This needs to be done in a more general purpose manner inside LinearModelFinalInference, for any estimator with a linear final model that admits an asymptotically linear representation.

For the ate without the underscore api, maybe we have rhe flag insample, with default True and if False, we declare that rhe method assumes the sample of X’s is independent of the training sample, when calculating standard errors and CIs