Surrogate screening: swapping selected and screened variables

After the LASSO selects an active set $E$, a common heuristic screens the inactive variables whose gradient at the LASSO solution, $|\nabla \ell(\hat{\beta})_j|$, is largest. ScreenedSelection adds this screening step to the selection event, so the polyhedral lemma still gives valid inference. The usual question is whether a screened variable $j$ belongs in the model $E \cup \{j\}$.

Here we screen the same way, keeping the top 10 inactive variables, but ask a different question. A screened variable $j$ may be a surrogate for a selected variable $s$: the two are correlated and the LASSO picked $s$ where it could have picked $j$. To check this, for each selected $s \in E$ and screened $j$ we drop $s$ from the model, put $j$ in its place, and test

\[ H_0: \beta_{j \cdot F_{s,j}} = 0, \qquad F_{s,j} = (E \setminus \{s\}) \cup \{j\}, \]

the coefficient of $x_j$ in the model $F_{s,j}$, or build a confidence interval for it.

Only the target changes. The selection event (LASSO, then screening) stays the same, so we reuse ScreenedSelection as it is and pass it a new contrast.

The swapped coefficient as a contrast

With $Q = X'X$ and $W = Q_{E,E}^{-1}$, the coefficient of $j$ in $F = F_{s,j}$ is

\[ \hat{\theta}_{s,j} = e_j'\, Q_{F,F}^{-1} Z_F = \eta_{s,j}' Z, \qquad (\eta_{s,j})_F = Q_{F,F}^{-1} e_j, \quad (\eta_{s,j})_{-F} = 0, \]

which for least squares is the coefficient of $x_j$ in the regression of $y$ on $X_F$. The estimand is $\eta_{s,j}' X'\mu$, the coefficient of $x_j$ in the regression of $\mu = E[y]$ on $X_F$.

The matrix $Q_{F,F}$ is built from the information matrix of the model $E$. Its block on $E \setminus \{s\}$ is a submatrix of $Q_{E,E} = W^{-1}$, which we already have. The rest is $Q_{j,E \setminus \{s\}}$ and $Q_{jj}$, which come from one matvec $Q e_j$ per screened variable, shared by all $s$. For a GLM, $Q$ is the Q_hat used for selection. We do not refit the model $F$ or recompute its information matrix.

Equivalently, with $G = E \cup \{j\}$ and $M = Q_{G,G}^{-1}$, dropping $s$ from $G$ gives

\[ \hat{\theta}_{s,j} = \hat{\beta}_{j \cdot G} - \frac{M_{js}}{M_{ss}}\, \hat{\beta}_{s \cdot G}. \]

This is the screened coefficient of $j$ minus a multiple of the coefficient of $s$ in the same model. Because $x_j$ and $x_s$ are correlated, part of the effect of $s$ moves onto $j$. The contrast is linear in $Z$, so the polyhedral lemma applies with the screening constraints we already have.

Setup

We use two designs, each with $n = 100$ observations, $p = 300$ features and 15 nonzero coefficients on every 20th feature, and $\sigma^2 = 1$ treated as known:

Independent features with signals of size $3/\sqrt{n}$. Here, ignoring the screening step badly breaks inference.
AR(1) features, with $\text{corr}(x_i, x_l) = \rho^{|i - l|}$, $\rho = 0.8$, and signals of size $5/\sqrt{n}$. Each signal’s neighbors are correlated null features that can stand in for it, so this design produces surrogates.

Selection uses a randomized response $y_{noisy} = y + \epsilon$ with $\epsilon \sim N(0, \frac{1-\pi}{\pi}\sigma^2 I)$ and $\pi = 0.8$.

The package reports variables (li$E, screen$screened) 0-based. In this vignette we add 1 and work with ordinary R indices.

library(lassoinf)

n <- 100; p <- 300; k <- 15
pi_selection <- 0.8
scalar_noise <- (1 - pi_selection) / pi_selection   # Var(omega) = scalar_noise * Sigma
signals <- seq(1, p, by = p / k)

make_data <- function(rho, amp) {
  # AR(1) columns: corr(x_i, x_l) = rho^|i - l|; rho = 0 gives independent columns
  X <- matrix(rnorm(n * p), n, p)
  for (i in 2:p) X[, i] <- rho * X[, i - 1] + sqrt(1 - rho^2) * X[, i]
  beta <- numeric(p)
  beta[signals] <- amp / sqrt(n) * sample(c(-1, 1), k, replace = TRUE)
  mu <- drop(X %*% beta)
  y <- mu + rnorm(n)
  y_noisy <- y + rnorm(n, sd = sqrt(scalar_noise))
  list(X = X, beta = beta, mu = mu, y = y, y_noisy = y_noisy)
}

We fit the LASSO $\frac{1}{2}\|y_{noisy} - Xb\|^2 + \lambda \|b\|_1$ with glmnet. glmnet scales the loss by $1/n$, so we pass it $\lambda / n$. $Q = X'X$ and $\Sigma = \sigma^2 X'X$ are matrix-free operators.

lam <- 0.8 * sqrt(2 * n * log(p))

lasso_inference <- function(d) {
  fit <- glmnet::glmnet(d$X, d$y_noisy, lambda = lam / n, intercept = FALSE,
                        standardize = FALSE, control = list(thresh = 1e-14))
  beta_hat <- as.numeric(coef(fit))[-1]
  beta_hat[abs(beta_hat) < 1e-6] <- 0
  G_hat <- drop(crossprod(d$X, d$X %*% beta_hat - d$y_noisy))
  Q <- xtvx_operator(d$X)
  LassoInference$new(beta_hat, G_hat, Q, rep(lam, p), NULL, NULL,
                     drop(crossprod(d$X, d$y)), Q, NULL, scalar_noise, level = 0.9)
}

Swap estimands

swap_estimands builds the contrasts $\eta_{s,j}$. It stores $Q_{E,E} = W^{-1}$ and caches the column $Q e_j$ of each screened variable. Each contrast then costs one $|E| \times |E|$ solve.

swap_estimands <- function(li) {
  E <- li$E + 1
  Q_EE <- solve(li$W)
  Q_EE <- (Q_EE + t(Q_EE)) / 2
  cache <- new.env()
  column <- function(j) {
    key <- as.character(j)
    if (is.null(cache[[key]])) {
      e_j <- numeric(p); e_j[j] <- 1
      cache[[key]] <- drop(li$Q_hat$multiply(e_j))
    }
    cache[[key]]
  }
  # F = (E minus {s}) plus {j}, with j last
  model <- function(s, j) c(E[E != s], j)
  swap_coef <- function(s, j) {
    if (!(s %in% E) || j %in% E) stop("need s in E and j not in E")
    Q_j <- column(j)
    keep <- E != s
    F <- model(s, j)
    m <- length(F)
    Q_FF <- matrix(0, m, m)
    Q_FF[-m, -m] <- Q_EE[keep, keep]
    Q_FF[-m, m] <- Q_FF[m, -m] <- Q_j[F[-m]]
    Q_FF[m, m] <- Q_j[j]
    eta <- numeric(p)
    eta[F] <- solve(Q_FF, c(numeric(m - 1), 1))
    contrast_estimand(eta)
  }
  list(model = model, swap_coef = swap_coef)
}

analyze runs the whole pipeline on one data set. It fits the LASSO, screens the top 10 inactive variables with the unchanged ScreenedSelection, and then loops over every pair $(s, j)$, giving $|E| \times 10$ estimands. estimand_summary with the ScreenedSelection conditions on both selection steps. The target $\eta_{s,j}' X'\mu$ is shown in the truth column.

swap_table <- function(selection, swaps, selected, screened) {
  pairs <- expand.grid(added = screened, dropped = selected)[, c("dropped", "added")]
  estimands <- Map(swaps$swap_coef, pairs$dropped, pairs$added)
  names(estimands) <- paste(pairs$dropped, pairs$added, sep = ":")
  df <- estimand_summary(selection, estimands)
  list(table = cbind(pairs, df[, -1]), estimands = estimands)
}

analyze <- function(d) {
  li <- lasso_inference(d)
  screen <- ScreenedSelection$new(li, top_k = 10)
  swaps <- swap_estimands(li)
  E <- li$E + 1; screened <- screen$screened + 1
  cat(length(E), "active variables,", sum(d$beta[E] != 0), "of them true signals;",
      sum(d$beta[-E] != 0), "true signals were left out\n")
  cat("screened:", screened, "\n")
  st <- swap_table(screen, swaps, E, screened)
  tab <- st$table
  tab$truth <- sapply(st$estimands, estimand_value, Z = drop(crossprod(d$X, d$mu)))
  tab$dropped_signal <- d$beta[tab$dropped] != 0
  tab$added_signal <- d$beta[tab$added] != 0
  tab$adjacent <- abs(tab$dropped - tab$added) <= 2
  rownames(tab) <- NULL
  list(li = li, screen = screen, swaps = swaps, E = E, screened = screened,
       table = tab, estimands = st$estimands)
}

p_value_grid <- function(tab) round(xtabs(p_value ~ dropped + added, tab), 3)

To compare, we also compute the LASSO-only tests, which use the LASSO constraints but ignore the screening step, and naive tests that ignore selection altogether:

naive_p_value <- function(li, est) {
  sd <- sqrt(li$si$compute_contrast(est$eta)$naive_variance)
  2 * pnorm(-abs(estimand_value(est, li$Z_full)) / sd)
}

compare_p_values <- function(a) {
  lasso_only <- swap_table(a$li, a$swaps, a$E, a$screened)$table
  data.frame(a$table[, c("dropped", "added")],
             screened = a$table$p_value,
             lasso_only = lasso_only$p_value,
             naive = sapply(a$estimands, naive_p_value, li = a$li), row.names = NULL)
}

For repeated experiments, we record for each pair $(s, j)$ whether the 90% interval covers $\eta_{s,j}' X'\mu$, and whether the test rejects at level 0.1. Intervals from the same data set are dependent, so the averages are over both replications and pairs.

one_replication <- function(rho, amp) {
  d <- make_data(rho, amp)
  li <- lasso_inference(d)
  screen <- ScreenedSelection$new(li, top_k = 10)
  swaps <- swap_estimands(li)
  Z_mu <- drop(crossprod(d$X, d$mu))
  q <- qnorm(0.95)
  rows <- list()
  for (s in li$E + 1) for (j in screen$screened + 1) {
    est <- swaps$swap_coef(s, j)
    target <- estimand_value(est, Z_mu)
    sd <- sqrt(li$si$compute_contrast(est$eta)$naive_variance)
    theta <- estimand_value(est, li$Z_full)
    info <- data.frame(dropped_signal = d$beta[s] != 0, added_signal = d$beta[j] != 0,
                       adjacent = abs(s - j) <= 2)
    for (method in c("screened", "lasso only")) {
      res <- contrast_inference(if (method == "screened") screen else li, est)
      rows[[length(rows) + 1]] <- data.frame(method = method,
        cover = res$lower_conf <= target && target <= res$upper_conf,
        length = res$upper_conf - res$lower_conf, reject = res$p_value < 0.1, info)
    }
    rows[[length(rows) + 1]] <- data.frame(method = "naive", cover = abs(theta - target) <= q * sd,
      length = 2 * q * sd, reject = abs(theta) > q * sd, info)
  }
  do.call(rbind, rows)
}

simulate <- function(rho, amp, nsim, seed = 1) {
  set.seed(seed)
  do.call(rbind, replicate(nsim, one_replication(rho, amp), simplify = FALSE))
}

coverage_table <- function(results) {
  data.frame(coverage = tapply(results$cover, results$method, mean),
             median_length = tapply(results$length, results$method, median))
}

Example 1: independent features

set.seed(0)
d <- make_data(rho = 0, amp = 3)
a <- analyze(d)
#> 6 active variables, 1 of them true signals; 14 true signals were left out
#> screened: 21 275 112 128 115 297 253 81 279 125

As a check, the estimate is the least squares coefficient of $x_j$ on $X_F$:

s <- a$E[1]; j <- a$screened[1]
F <- a$swaps$model(s, j)
c(ols = unname(tail(qr.solve(d$X[, F], d$y), 1)),
  swap = estimand_value(a$swaps$swap_coef(s, j), a$li$Z_full))
#>      ols     swap 
#> 0.171792 0.171792

head(a$table, 10)
#>    dropped added   estimate  lower_conf upper_conf    p_value        truth
#> 1        5    21  0.1717920 -0.09606102  0.6231150 0.22803446  0.254789202
#> 2        5   275  0.1624549 -0.28531803  0.3926296 0.79441827  0.010198638
#> 3        5   112 -0.1611645 -0.47233129  0.2600691 0.63349694  0.022210276
#> 4        5   128 -0.1005752 -0.28404646  0.4191999 0.75206212 -0.073187870
#> 5        5   115 -0.1640124 -0.34248901  0.3208632 0.95714797 -0.006686096
#> 6        5   297  0.1543233 -0.32956066  0.4441427 0.80744458  0.139615442
#> 7        5   253 -0.1077188 -0.11600897  0.6941805 0.24067268 -0.062439529
#> 8        5    81  0.2888889  0.04214024  0.8487525 0.06925875  0.295175157
#> 9        5   279 -0.1084395 -0.24839617  0.5555619 0.52974231 -0.040221170
#> 10       5   125  0.2368126 -0.19905394  0.5475778 0.44259914  0.040440723
#>    dropped_signal added_signal adjacent
#> 1           FALSE         TRUE    FALSE
#> 2           FALSE        FALSE    FALSE
#> 3           FALSE        FALSE    FALSE
#> 4           FALSE        FALSE    FALSE
#> 5           FALSE        FALSE    FALSE
#> 6           FALSE        FALSE    FALSE
#> 7           FALSE        FALSE    FALSE
#> 8           FALSE         TRUE    FALSE
#> 9           FALSE        FALSE    FALSE
#> 10          FALSE        FALSE    FALSE

The p-values are easier to read as a grid, with selected variables as rows and screened variables as columns:

p_value_grid(a$table)
#>        added
#> dropped    21    81   112   115   125   128   253   275   279   297
#>     5   0.228 0.069 0.633 0.957 0.443 0.752 0.241 0.794 0.530 0.807
#>     101 0.257 0.084 0.554 0.983 0.408 0.694 0.245 0.854 0.582 0.715
#>     106 0.271 0.075 0.625 0.954 0.402 0.693 0.205 0.799 0.613 0.687
#>     156 0.196 0.077 0.402 0.989 0.239 0.917 0.232 0.996 0.757 0.792
#>     177 0.282 0.086 0.587 0.963 0.424 0.690 0.228 0.852 0.545 0.775
#>     204 0.270 0.094 0.555 0.940 0.459 0.730 0.243 0.829 0.544 0.854

A small p-value in row $s$, column $j$ means that $x_j$ still carries signal once $x_s$ is removed from the model. When $x_j$ is a surrogate for $x_s$ these p-values are small in row $s$ and large in the other rows. A screened variable whose p-value is small in every row is a candidate for a missed signal in its own right. With independent features, most columns of the grid change little from row to row: dropping one selected variable barely changes the coefficient of $x_j$.

Does screening matter?

The three sets of p-values for the same pairs:

head(compare_p_values(a), 10)
#>    dropped added   screened lasso_only       naive
#> 1        5    21 0.22803446 0.03815790 0.080569181
#> 2        5   275 0.79441827 0.17969244 0.081567092
#> 3        5   112 0.63349694 0.67069736 0.109449713
#> 4        5   128 0.75206212 0.27618682 0.297729262
#> 5        5   115 0.95714797 0.92335662 0.070987779
#> 6        5   297 0.80744458 0.48408157 0.146473635
#> 7        5   253 0.24067268 0.31500545 0.335946581
#> 8        5    81 0.06925875 0.01586708 0.008913217
#> 9        5   279 0.52974231 0.25314737 0.323291883
#> 10       5   125 0.44259914 0.01013596 0.019839301

Over 50 replications:

## results_indep <- simulate(rho = 0, amp = 3, nsim = 50)
coverage_table(results_indep)
#>             coverage median_length
#> lasso only 0.8254215     0.7040169
#> naive      0.8430610     0.3528058
#> screened   0.9118029     0.7812204
with(results_indep, tapply(reject, list(method, added_signal = added_signal), mean))
#>             added_signal
#>                  FALSE      TRUE
#>   lasso only 0.2642934 0.4758395
#>   naive      0.3404223 0.5192465
#>   screened   0.1197411 0.2588043

Only the intervals that condition on the screening step come close to the nominal 90% coverage. When the added variable is null, the screened tests reject at about the nominal 10%, and the LASSO-only and naive tests reject far more often. With independent features, $M_{js} \approx 0$ in the decomposition above, so $\hat{\theta}_{s,j}$ is nearly the screened coefficient $\hat{\beta}_{j \cdot G}$. That coefficient is large partly because $x_j$ was chosen for its large gradient, which is the selection effect that screening accounts for.

Example 2: correlated features and surrogates

set.seed(0)
d <- make_data(rho = 0.8, amp = 5)
a <- analyze(d)
#> 23 active variables, 6 of them true signals; 9 true signals were left out
#> screened: 218 140 56 221 234 1 200 247 212 163

The interesting pairs are a selected null variable next to a screened true signal: the LASSO chose $x_s$ in place of the nearby $x_j$. Here are the pairs at distance at most 2:

a$table[a$table$adjacent, ]
#>     dropped added    estimate  lower_conf upper_conf   p_value       truth
#> 6         2     1 -0.37062986 -0.87642351  0.1507497 0.2451725 -0.47395340
#> 53       54    56  0.36505323 -0.30923182  0.5594844 0.6343925  0.21236075
#> 140     162   163  0.31238042 -0.08791716  0.6732315 0.2056827  0.27405821
#> 151     219   218 -0.03023869 -0.30447892  0.6507284 0.5513200  0.03261443
#> 154     219   221 -0.18554004 -0.78502289  0.8737316 0.9299056 -0.52199381
#> 161     220   218  0.06259443 -0.27492338  1.0324281 0.3406162  0.16942370
#> 164     220   221 -0.33518421 -0.95304053  0.2035123 0.2864314 -0.46233925
#>     dropped_signal added_signal adjacent
#> 6            FALSE         TRUE     TRUE
#> 53           FALSE        FALSE     TRUE
#> 140          FALSE        FALSE     TRUE
#> 151          FALSE        FALSE     TRUE
#> 154          FALSE         TRUE     TRUE
#> 161          FALSE        FALSE     TRUE
#> 164          FALSE         TRUE     TRUE

p_value_grid(a$table)
#>        added
#> dropped     1    56   140   163   200   212   218   221   234   247
#>     2   0.245 0.811 0.160 0.129 0.016 0.137 0.515 0.838 0.598 0.840
#>     6   0.573 0.461 0.162 0.168 0.023 0.179 0.634 0.993 0.860 0.545
#>     20  0.866 0.426 0.230 0.178 0.042 0.127 0.486 0.940 0.715 0.647
#>     22  0.626 0.543 0.184 0.155 0.026 0.150 0.519 0.872 0.715 0.635
#>     41  0.492 0.757 0.195 0.158 0.022 0.139 0.424 0.832 0.604 0.788
#>     54  0.601 0.634 0.179 0.159 0.031 0.140 0.542 0.851 0.717 0.611
#>     60  0.575 0.347 0.248 0.163 0.027 0.119 0.521 0.846 0.693 0.585
#>     81  0.643 0.447 0.149 0.129 0.031 0.126 0.542 0.668 0.745 0.595
#>     101 0.680 0.461 0.199 0.155 0.025 0.139 0.498 0.809 0.774 0.613
#>     111 0.611 0.527 0.187 0.164 0.027 0.159 0.524 0.886 0.748 0.662
#>     121 0.611 0.525 0.189 0.163 0.031 0.148 0.523 0.886 0.740 0.648
#>     136 0.725 0.329 0.370 0.127 0.021 0.149 0.640 0.847 0.827 0.548
#>     152 0.620 0.455 0.152 0.137 0.030 0.122 0.576 0.726 0.686 0.695
#>     162 0.836 0.783 0.205 0.206 0.050 0.160 0.417 0.838 0.823 0.610
#>     184 0.592 0.517 0.182 0.168 0.026 0.148 0.530 0.868 0.734 0.634
#>     219 0.630 0.518 0.215 0.160 0.025 0.135 0.551 0.930 0.755 0.575
#>     220 0.530 0.528 0.203 0.163 0.022 0.097 0.341 0.286 0.732 0.449
#>     224 0.619 0.466 0.185 0.183 0.028 0.157 0.537 0.933 0.714 0.622
#>     241 0.732 0.425 0.129 0.130 0.024 0.209 0.762 0.861 0.995 0.830
#>     261 0.905 0.400 0.321 0.126 0.032 0.083 0.529 0.793 0.620 0.704
#>     263 0.663 0.426 0.273 0.100 0.025 0.089 0.561 0.729 0.899 0.675
#>     276 0.540 0.746 0.209 0.119 0.019 0.151 0.497 0.939 0.450 0.804
#>     277 0.619 0.457 0.224 0.186 0.026 0.130 0.543 0.911 0.878 0.618

For each surrogate pair (a selected null $s$ next to a screened signal $j$), we compare the p-value in row $s$ with the median of column $j$ over the other rows:

surrogate_pairs <- a$table[a$table$adjacent & !a$table$dropped_signal & a$table$added_signal, ]
surrogate_pairs$median_other_rows <- sapply(seq_len(nrow(surrogate_pairs)), function(i) {
  col <- a$table[a$table$added == surrogate_pairs$added[i], ]
  median(col$p_value[col$dropped != surrogate_pairs$dropped[i]])
})
surrogate_pairs[, c("dropped", "added", "p_value", "median_other_rows")]
#>     dropped added   p_value median_other_rows
#> 6         2     1 0.2451725         0.6197503
#> 154     219   221 0.9299056         0.8492854
#> 164     220   221 0.2864314         0.8562716

head(compare_p_values(a), 10)
#>    dropped added   screened lasso_only       naive
#> 1        2   218 0.51503535 0.55553556 0.933722950
#> 2        2   140 0.16023117 0.06766424 0.022893000
#> 3        2    56 0.81132276 0.85092786 0.580623487
#> 4        2   221 0.83836968 0.83836968 0.435969343
#> 5        2   234 0.59793153 0.10615870 0.056253583
#> 6        2     1 0.24517247 0.24517247 0.009117723
#> 7        2   200 0.01610315 0.01610315 0.011164592
#> 8        2   247 0.83986640 0.35082434 0.876464877
#> 9        2   212 0.13744425 0.13744425 0.001671224
#> 10       2   163 0.12894095 0.12894095 0.454179662

Over 30 replications:

## results_ar <- simulate(rho = 0.8, amp = 5, nsim = 30)
coverage_table(results_ar)
#>             coverage median_length
#> lasso only 0.8734312     1.0879803
#> naive      0.9221629     0.5403365
#> screened   0.8943925     1.1807819
with(results_ar, tapply(reject, list(method, dropped_signal = dropped_signal,
                                     added_signal = added_signal), mean))
#> , , added_signal = FALSE
#> 
#>             dropped_signal
#>                  FALSE      TRUE
#>   lasso only 0.2162162 0.2241888
#>   naive      0.1592664 0.1723557
#>   screened   0.1452703 0.1508639
#> 
#> , , added_signal = TRUE
#> 
#>             dropped_signal
#>                  FALSE      TRUE
#>   lasso only 0.3459119 0.3175074
#>   naive      0.3411950 0.3175074
#>   screened   0.2877358 0.2700297

When the added variable is null, the target $\beta_{j \cdot F}$ is not exactly zero. If $x_j$ is correlated with the dropped $x_s$ and $x_s$ carries signal, the coefficient of $x_j$ absorbs part of that signal. The table below restricts to the surrogate pairs, where a selected null sits next to a screened signal:

surrogate <- subset(results_ar, adjacent & !dropped_signal & added_signal)
data.frame(pairs = tapply(surrogate$reject, surrogate$method, length),
           coverage = tapply(surrogate$cover, surrogate$method, mean),
           reject = tapply(surrogate$reject, surrogate$method, mean))
#>            pairs  coverage    reject
#> lasso only    29 0.8965517 0.3448276
#> naive         29 0.8275862 0.7241379
#> screened      29 0.8965517 0.2758621

Conditioning on screening can be changed through ScreenedSelection, for example ScreenedSelection$new(li, top_k = 10, conditioning = "exact") or a threshold. The swap contrasts do not depend on that choice.