Rosenman 2020 - Unbiased Risk Estimate and Common Shrinkage

Summary

Generalising Stein’s unbiased risk estimate to heteroscedastic, weighted loss gives the risk of any estimator $\kappa =Z+\Sigma g$ in closed form. Minimising the URE over a single shrinkage factor yields ${\kappa }_1={\hat{\tau }}_r-{\lambda }_1^{URE}({\hat{\tau }}_r-{\hat{\tau }}_o)$ with ${\lambda }_1^{URE}=\text{Tr}({\Sigma }_{r}D)/\{({\hat{\tau }}_o-{\hat{\tau }}_r{)}^{T}D({\hat{\tau }}_o-{\hat{\tau }}_r)\}$. It dominates the RCT-only estimator whenever $4{\max }_k{d}_k{\sigma }_{rk}^2\le {\sum }_k{d}_k{\sigma }_{rk}^2$ (so $K\ge 4$), and its positive-part version is asymptotically optimal among common-factor shrinkers as $K\to \infty$.

Overview

Procedure used for every estimator in the paper:

1. Posit a structure for the shrinkage estimator.
2. Derive the shrinkage factor by optimising the URE assuming the factor is known (“base” estimator).
3. Optionally correct for the extra risk from estimating the factor from the data.

Main Content

Theorem 1 (estimator risk; generalisation of Strawderman et al. 2003, Thm 3.1)

Let $Z\sim N(\theta ,\Sigma )$ with $\Sigma =\text{diag}({\sigma }_1^2,\dots ,{\sigma }_K^2)$, $Y$ a random vector, $D=\frac{1}{K}\text{diag}(d_1,\dots ,d_K)$, $L(\theta ,v)=(v-\theta {)}^{T}D(v-\theta )$ and $\kappa (Z,Y)=Z+\Sigma g(Z,Y)$ with $g$ differentiable in $Z$ and $E\Vert g{\Vert }^2<\infty$. Then

$$R(\theta ,\kappa )=\text{Tr}(\Sigma D)+E\left[\sum _k{\sigma }_k^4{D}_{kk}\left(g_k^2(Z,Y)+2\frac{\partial g_k(Z,Y)}{\partial Z_k}\right)\right],D_{kk}=d_k/K.$$

Proof. Fix $Y=y$ and apply Theorem 3.1 of Strawderman et al. (2003) to $g^{(y)}(Z)=g(Z,y)$; then use the tower rule, $E[R(\theta ,\kappa (Z,Y))]=E[E(R(\theta ,\kappa (Z,y))|Y=y)]$. (In the paper the $1/K$ is displayed outside; it is absorbed into $D$ here.)

Dropping the expectation gives the generalised Stein unbiased risk estimate

$$\text{URE}(\theta ,\kappa (Z,Y))=\text{Tr}(\Sigma D)+\sum _k{\sigma }_k^4{D}_{kk}\left(g_k^2+2\partial g_k/\partial Z_k\right).$$

The estimator ${\kappa }_1$ (common shrinkage, Sec. 4.2)

Structure: $\kappa (\lambda ,{\hat{\tau }}_r,{\hat{\tau }}_o)={\hat{\tau }}_r-\lambda ({\hat{\tau }}_r-{\hat{\tau }}_o)$. For fixed $\lambda$ (Eq. 1):

$$\text{URE}=\text{Tr}({\Sigma }_{r}D)+{\lambda }^2({\hat{\tau }}_o-{\hat{\tau }}_r{)}^{T}D({\hat{\tau }}_o-{\hat{\tau }}_r)-2\lambda \text{Tr}({\Sigma }_{r}D).$$

Strictly convex in $\lambda$ if ${\hat{\tau }}_r\ne {\hat{\tau }}_o$; minimiser

${\kappa }_1$ and ${\kappa }_{1+}$

$${\lambda }_1^{URE}=\frac{\text{Tr}({\Sigma }_{r}D)}{({\hat{\tau }}_o-{\hat{\tau }}_r{)}^{T}D({\hat{\tau }}_o-{\hat{\tau }}_r)},{\kappa }_1={\hat{\tau }}_r-{\lambda }_1^{URE}({\hat{\tau }}_r-{\hat{\tau }}_o),{\kappa }_{1+}={\hat{\tau }}_o+{\left(1-{\lambda }_1^{URE}\right)}_{+}({\hat{\tau }}_r-{\hat{\tau }}_o)$$

(the positive part keeps the weight on ${\hat{\tau }}_r$ between 0 and 1, i.e. $\lambda \le 1$: Baranchik 1964). ${\kappa }_1$ generalises Green & Strawderman (1991) to the heteroscedastic, weighted-loss case.

Oracle. The weighted-MSE-optimal shrinkage factor is

$${\lambda }_{opt}=\frac{\text{Tr}({\Sigma }_{r}D)}{\text{Tr}({\Sigma }_{r}D)+\text{Tr}({\Sigma }_{o}D)+{\xi }^T{D}^2\xi }.$$

The denominator cannot be estimated because $\xi$ is unknown, but it equals $E[({\hat{\tau }}_o-{\hat{\tau }}_r{)}^{T}D({\hat{\tau }}_o-{\hat{\tau }}_r)]$, so ${\kappa }_1$ substitutes that “best guess”. Interpretation: the weight moved from ${\hat{\tau }}_r$ toward ${\hat{\tau }}_o$ is the RCT variance share of total discrepancy.

Lemma 1 (testable dominance)

If $4{\max }_k{d}_k{\sigma }_{rk}^2\le {\sum }_k{d}_k{\sigma }_{rk}^2$ then ${\kappa }_1$ dominates ${\hat{\tau }}_r$ under the loss.

Proof. By Theorem 1, $R(\tau ,{\hat{\tau }}_r)-R(\tau ,{\kappa }_1)=E\left[-\frac{\text{Tr}({\Sigma }_{r}D{)}^2}{Q}+\frac{4\text{Tr}({\Sigma }_{r}D){\Delta }^T{D}^2{\Sigma }_r\Delta }{Q^2}\right]$ with $\Delta ={\hat{\tau }}_o-{\hat{\tau }}_r$, $Q={\Delta }^{T}D\Delta$. Bound ${\Delta }^T{D}^2{\Sigma }_r\Delta \le {\max }_k{d}_k{\sigma }_{rk}^2\cdot Q$ (up to the $1/K$ in $D$) to get a bound of the form $\text{Tr}({\Sigma }_{r}D)E\left[\frac{4{\max }_k{\sigma }_{rk}^2-\text{Tr}({\Sigma }_{r}D)\cdot (\text{normalisation})}{Q}\right]$ whose numerator is non-positive exactly under the stated condition, so the risk difference is $\ge 0$ in favour of ${\kappa }_1$ (the paper’s displayed chain carries the $d_k$ and $1/K$ factors, which the PDF extraction scrambles).

Remark: the condition forces $K\ge 4$ at least (the paper says this is at least the dimension needed for classical James-Stein shrinkage toward a grand mean, Efron 2012), more if variances/weights are unequal. In a multi-arm trial that means at least four distinct treatments; for heterogeneous effects at least four subgroups.

Theorem 2 ( ${\kappa }_{1+}$ asymptotic risk)

Under mild conditions, as $K\to \infty$ ${\kappa }_{1+}$ has the lowest risk among all estimators with a shared shrinkage factor across components.

Proof (Appendix 8, modelled on Xie, Kou & Brown 2012).

• Lemma 4. Assume $\mathrm{limsup}\frac{1}{K}{\sum }_k{d}_k^2{\sigma }_{rk}^2{\xi }_k^2<\infty$, $\mathrm{limsup}\frac{1}{K}{\sum }_k{d}_k^2{\sigma }_{rk}^2{\sigma }_{ok}^2<\infty$, $\mathrm{limsup}\frac{1}{K}{\sum }_k{d}_k^2{\sigma }_{rk}^4<\infty$. Then ${\sup }_{0\le \lambda \le 1}|\text{URE}(\tau ,\kappa (\lambda ))-L(\tau ,\kappa (\lambda ))|\to 0$ in $L^2$ and in probability. The difference splits into $\frac{1}{K}{\sum }_k{d}_k[{\sigma }_{rk}^2-({\hat{\tau }}_{rk}-{\tau }_k{)}^2]$ (variance $\propto \frac{2}{K^2}\sum d_k^2{\sigma }_{rk}^4\to 0$) and $\frac{2\lambda }{K}{\sum }_k{d}_k[{\sigma }_{rk}^2-({\hat{\tau }}_{rk}-{\tau }_k)({\hat{\tau }}_{rk}-{\hat{\tau }}_{ok})]$; for the latter bound ${\sup }_{\lambda }$ by ${\sup }_{1\ge c_1\ge \cdots \ge c_K\ge 0}$, apply Lemma 2.1 of Li et al. (1986) to reduce to ${\max }_j|\frac{2}{K}{\sum }_{k\le j}\cdots |$, note the partial sums $M_j$ form a martingale (mean zero), and use the $L^p$ maximal inequality $E{\max }_j{M}_j^2\le 4E{M}_K^2=4{\sum }_k{d}_k^2{\sigma }_{rk}^2({\xi }_k^2+{\sigma }_{ok}^2+2{\sigma }_{rk}^2)$.
• Lemma 5. Under those conditions ${\lim }_{K\to \infty }R(\kappa ({\lambda }_1^{URE}))-R(\kappa (\lambda ))\le 0$ for every fixed $\lambda$. Write $L({\lambda }^{URE})-L(\lambda )=[L-\text{URE}]({\lambda }^{URE})+[\text{URE}({\lambda }^{URE})-\text{URE}(\lambda )]+[\text{URE}-L](\lambda )$; the middle term is $\le 0$ by minimality, the others are bounded by $2\sup |\text{URE}-L|\to 0$.

Corrected version ${\kappa }_1^{\ast }$ (optional step 3)

Because ${\lambda }_1^{URE}$ is estimated from the data, the risk of ${\kappa }_1$ is not the URE at fixed $\lambda$; one should shrink by less. Scale by $a$:

$$\text{URE}(\kappa (a{\lambda }_1^{URE}))=\text{Tr}({\Sigma }_{r}D)+\frac{(a^2-2a)\text{Tr}({\Sigma }_{r}D{)}^2}{{\Delta }^{T}D\Delta }+\frac{4a\text{Tr}({\Sigma }_{r}D){\Delta }^T{D}^2{\Sigma }_r\Delta }{({\Delta }^{T}D\Delta {)}^2},$$

minimised at

$$a_1^{\ast }=1-\frac{2{\Delta }^T{D}^2{\Sigma }_r\Delta }{{\Delta }^{T}D\Delta }\cdot \frac{1}{\text{Tr}({\Sigma }_{r}D)},{\kappa }_1^{\ast }={\hat{\tau }}_r-a_1^{\ast }\frac{\text{Tr}({\Sigma }_{r}D)}{{\Delta }^{T}D\Delta }({\hat{\tau }}_r-{\hat{\tau }}_o).$$

Iterating would continue “ad infinitum”; the paper stops after one step.

Lemma 2 ( ${\kappa }_1^{\ast }$ vs ${\kappa }_1$)

${\kappa }_1^{\ast }$ has risk no greater than ${\kappa }_1$ if ${\max }_k{d}_k^2{\sigma }_{rk}^4\le \frac{3}{2}{\left({\min }_k{\sigma }_{rk}^2{d}_k\right)}^2$ (reconstructed from the proof, which requires $8{\max }_k{d}_k^2{\sigma }_{rk}^4\le 12({\min }_k{\sigma }_{rk}^2{d}_k{)}^2$); i.e. stratum variances and weights are sufficiently concentrated.

Examples

import numpy as np
def kappa1_plus(tau_r, tau_o, sigma2_r, d=None):
    """Common-shrinkage estimator of Rosenman et al. (2020). Returns (estimate, weight on tau_r)."""
    K = len(tau_r); d = np.ones(K) if d is None else d
    D = d / K
    lam = (sigma2_r * D).sum() / max(((tau_o - tau_r)**2 * D).sum(), 1e-12)   # lambda_1^URE
    lam = min(lam, 1.0)                                                         # positive part
    return tau_r - lam * (tau_r - tau_o), 1 - lam   # weight on tau_o = lam
 
# Toy: 20 strata, RCT sd 0.2, observational sd 0.02, bias sd 0.1
rng = np.random.default_rng(0); K = 20; tau = rng.uniform(0, 1, K)
tr = tau + rng.normal(0, 0.2, K); to = tau + rng.normal(0, 0.1, K) + rng.normal(0, 0.02, K)
est, _ = kappa1_plus(tr, to, np.full(K, 0.04))
print(((tr-tau)**2).mean(), ((est-tau)**2).mean())

Reading the weight

With equal RCT variances ${\sigma }_{rk}^2=0.04$ (so $\text{Tr}({\Sigma }_{r}D)=0.04$), if the observed mean squared discrepancy (with $D=I/K$) ${\Delta }^{T}D\Delta =0.10$ then ${\lambda }_1=0.04/0.10=0.4$: 40% of the way from the RCT toward the observational estimate. If the discrepancy is $0.02$ (below the RCT’s own noise), ${\lambda }_1=2\to$ capped at 1 by the positive part, i.e. full weight on ${\hat{\tau }}_o$.

Connections

• ${\kappa }_1$ has the same form as the James-Stein / empirical Bayes shrinkage: James-Stein Estimator, Empirical Bayes Interpretation of Shrinkage, Stein’s Paradox and Risk Dominance.
• Bayesian equivalent: a normal prior on the stratum bias ${\xi }_k$ with an estimated scale gives the same weight; see Rosenman 2020 - Hierarchical Bias Model Sketch for Few Holdouts.

See Also

Rosenman 2020 - Combining Observational and Experimental Datasets - Overview