Selection when p >> n — genomics-scale shrinkage

Bayesian variable selection, Part 8 — the horseshoe at scale

Every dataset so far had a handful of predictors. Real variable selection often looks the opposite: thousands of candidate predictors, tens of samples. Genomics is the canonical case — a few dozen microarrays, several thousand genes — and it is exactly where the earlier machinery breaks down. Enumeration (Part 2) would need $2^{4088}$ models; discrete spike-and-slab (Part 1) mixes painfully with 4088 indicators. The tool built for this regime is the continuous global-local shrinkage prior, above all the horseshoe (Carvalho-Polson-Scott 2010).

The benchmark is the riboflavin dataset (Bühlmann-Kalisch-Meier 2014): $n=71$ Bacillus subtilis cultures, $p=4088$ gene-expression measurements, and a continuous response — the log riboflavin (vitamin B2) production rate. The question: which handful of genes drives production? We fit a from-scratch horseshoe, cross-check against PyMC (same prior, general-purpose NUTS) and R (glmnet, the lasso — the frequentist workhorse for $p\gg n$).

The model, and why a special sampler is needed

Standardize the genes and regress: $$ y_i = \sum_{j=1}^{p} x_{ij}\,\beta_j + \varepsilon_i,\qquad \varepsilon_i\sim\mathcal N(0,\sigma^2),\qquad p=4088 \gg n=71. $$ The horseshoe prior gives each coefficient its own local scale on top of a global one: $$ \beta_j\mid\lambda_j,\tau,\sigma \sim \mathcal N(0,\ \lambda_j^2\tau^2\sigma^2),\qquad \lambda_j\sim C^+(0,1),\qquad \tau\sim C^+(0,1). $$ The global $\tau$ pulls everything toward zero (essential when $p\gg n$); the heavy-tailed local $\lambda_j$ lets a few coefficients escape. The prior is essentially parameter-free — the half-Cauchys need no tuning — which is much of its appeal. Its action is summarised by the shrinkage weight $\kappa_j = 1/(1+\lambda_j^2\tau^2)\in[0,1]$: $\kappa_j\approx1$ means gene $j$ is shrunk to zero, $\kappa_j\approx0$ means it is kept.

Two ingredients make it tractable.

1. Makalic-Schmidt (2016) auxiliary variables turn each half-Cauchy into a pair of conjugate inverse-gamma draws, so the whole sampler is Gibbs.
2. The Bhattacharya-Chakraborty-Mallick (2016) sampler for the coefficient block. The naive draw of $\beta$ needs the $p\times p$ precision matrix — $O(p^3)=O(4088^3)$, hopeless. Their exact algorithm never forms it; it solves an $n\times n$ system instead, costing $O(n^2p)$. With $n=71$ that turns an impossible step into a millisecond one — the entire run is ~9 seconds.

There are no free hyperparameters to set (all half-Cauchy scales are 1, the standard horseshoe); $\sigma^2$ has the usual reference prior. Engine: hshd.py.

import numpy as np, pandas as pd, matplotlib.pyplot as plt, time
import hshd                                                         # from-scratch high-dim horseshoe (Bhattacharya sampler)

d = pd.read_csv("ribo_data.csv")
genes = np.array([c[2:] if c.startswith("X.") else c for c in d.columns[1:]])
y = d["y"].values; X = d.iloc[:, 1:].values                        # 71 x 4088
t0 = time.perf_counter()
res = hshd.horseshoe_hd(y, X, ndraw=4000, burn=2000, seed=1)
gibbs_sec = time.perf_counter() - t0

b, kap = res["beta"], res["kappa"]; order = np.argsort(kap)         # rank by shrinkage weight (low kappa = signal)
print("from-scratch horseshoe (Bhattacharya sampler): %.1f s for 6000 iterations at p = %d\n" % (gibbs_sec, res["p"]))
print("top genes by signal (low shrinkage weight kappa):")
for j in order[:6]: print("  %-10s beta = %+.2f   kappa = %.2f" % (genes[j], b[j], kap[j]))
print("\ngenes with kappa < 0.5 (clear signal): %d of %d    median kappa = %.3f"
      % ((kap < 0.5).sum(), res["p"], np.median(kap)))

# --- Fig: the 4088 coefficients (nearly all zero), and the shrinkage weights ---
sig = kap < 0.5                                                     # the escaped signals
idx = np.arange(res["p"])
fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))
ax[0].scatter(idx[~sig], b[~sig], s=6, color="navy", alpha=.30, label="shrunk to ~0")
ax[0].scatter(idx[sig],  b[sig],  s=40, color="firebrick", zorder=3, label="signal (kappa < 0.5)")
for j in order[:2]:
    ax[0].annotate(genes[j], (j, b[j]), color="firebrick",
                   xytext=(6, 0), textcoords="offset points", va="center", fontsize=9)
ax[0].axhline(0, color="k", lw=.8)
ax[0].set(xlabel="gene index", ylabel="posterior mean coefficient",
          title="Horseshoe coefficients (p = 4088): almost all at zero")
ax[0].legend(loc="lower right"); ax[0].grid(alpha=.25)

ax[1].hist(kap, bins=40, color="firebrick", alpha=.80, edgecolor="white", linewidth=.3)
ax[1].axvline(0.5, color="navy", ls="--", lw=1.2, label="kappa = 0.5")
ax[1].set(xlabel="shrinkage weight  kappa = 1/(1+lambda^2 tau^2)", ylabel="number of genes",
          title="Shrinkage weights pile up at kappa = 1 (off); a few near 0 (on)")
ax[1].legend(); ax[1].grid(alpha=.25)
plt.tight_layout(); plt.show()

from-scratch horseshoe (Bhattacharya sampler): 9.8 s for 6000 iterations at p = 4088

top genes by signal (low shrinkage weight kappa):
  YOAB_at    beta = -0.43   kappa = 0.25
  YXLE_at    beta = -0.41   kappa = 0.29
  ARGB_at    beta = -0.16   kappa = 0.71
  YISU_at    beta = +0.13   kappa = 0.72
  YXIB_at    beta = -0.07   kappa = 0.84
  CARB_at    beta = -0.07   kappa = 0.87

genes with kappa < 0.5 (clear signal): 2 of 4088    median kappa = 0.998

Reading the genome-wide fit

• A needle in 4088 haystacks. The horseshoe drives essentially every coefficient to zero — the median shrinkage weight is 0.998 — and lets just a couple escape. YOAB_at and YXLE_at stand out sharply ($\kappa\approx0.25,0.29$); a second tier (ARGB, YISU) is weaker. This is what selection must look like when $p\gg n$: aggressive global shrinkage, a few local exceptions.
• The signals are biologically sensible. YXL-region genes and YOAB recur across analyses of this benchmark as the drivers of riboflavin production — the sparse Bayesian fit recovers them without any pre-screening.
• The shrinkage weights tell the story (right panel). Plotted for all 4088 genes, the mass of $\kappa_j$ sits against $\kappa=1$; only a handful lie near 0. That two-groups picture — almost everything off, a few on — is the horseshoe's signature and the continuous-shrinkage answer to the discrete spike-and-slab of Part 1.

import pymc as pm, time                                              # same prior, general-purpose sampler
import matplotlib.pyplot as plt
Xz = (X - X.mean(0)) / X.std(0); yc = y - y.mean(); n, p = Xz.shape
with pm.Model() as mod:
    tau = pm.HalfCauchy("tau", 1.0); lam = pm.HalfCauchy("lam", 1.0, shape=p)
    z   = pm.Normal("z", 0, 1, shape=p); sig = pm.HalfNormal("sig", 1.0)
    beta = pm.Deterministic("beta", z*lam*tau*sig)                  # sigma-scaled horseshoe -- SAME prior as hshd.py
    pm.Normal("y", pm.math.dot(Xz, beta), sig, observed=yc)
    t0 = time.perf_counter()
    idata = pm.sample(800, tune=800, chains=2, target_accept=0.9, random_seed=1, progressbar=False)
    nuts_sec = time.perf_counter() - t0
bpm = idata.posterior["beta"].mean(("chain","draw")).values
r = np.corrcoef(res["beta"], bpm)[0, 1]
print("PyMC horseshoe (NUTS, sigma-scaled to match): %.0f s   top genes include %s"
      % (nuts_sec, ", ".join(genes[np.argsort(-np.abs(bpm))[:2]])))
print("beta agreement, all %d genes: r = %.2f  (the strong signals match; the near-zero"
      % (p, r))
print("cloud wobbles -- those coefficients are ~0 with large posterior uncertainty)")
print("-> ~%.0fx slower than the Gibbs: general-purpose NUTS strains at p=4088 where the"
      % (nuts_sec / gibbs_sec))
print("   Bhattacharya-based Gibbs is trivial (%.1f s)." % gibbs_sec)

# --- Fig: every gene's coefficient, from-scratch vs PyMC (do the engines agree?) ---
lim = max(np.abs(res["beta"]).max(), np.abs(bpm).max()) * 1.05
top = np.argsort(-np.abs(res["beta"]))[:2]                          # label the dominant signals
fig, ax = plt.subplots(figsize=(5.6, 5.6))
ax.axline((0, 0), slope=1, color="navy", ls="--", lw=1.2, label="y = x (perfect agreement)")
ax.axhline(0, color="k", lw=.6); ax.axvline(0, color="k", lw=.6)
ax.scatter(res["beta"], bpm, s=10, color="firebrick", alpha=.40, edgecolor="none")
for j in top:
    ax.annotate(genes[j], (res["beta"][j], bpm[j]), color="firebrick",
                xytext=(6, 4), textcoords="offset points", fontsize=9)
ax.set(xlim=(-lim, lim), ylim=(-lim, lim),
       xlabel="from-scratch Gibbs coefficient", ylabel="PyMC NUTS coefficient",
       title="Do the engines agree?  strong signals coincide, near-zero cloud scatters (r = %.2f)" % r)
ax.legend(loc="upper left"); ax.grid(alpha=.25)
plt.tight_layout(); plt.show()

C:\Users\user\anaconda3\envs\pymc-env\Lib\site-packages\threadpoolctl.py:1226: RuntimeWarning: 
Found Intel OpenMP ('libiomp') and LLVM OpenMP ('libomp') loaded at
the same time. Both libraries are known to be incompatible and this
can cause random crashes or deadlocks on Linux when loaded in the
same Python program.
Using threadpoolctl may cause crashes or deadlocks. For more
information and possible workarounds, please see
    https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md

  warnings.warn(msg, RuntimeWarning)
Initializing NUTS using jitter+adapt_diag...

Multiprocess sampling (2 chains in 2 jobs)

NUTS: [tau, lam, z, sig]

Sampling 2 chains for 800 tune and 800 draw iterations (1_600 + 1_600 draws total) took 151 seconds.

There were 94 divergences after tuning. Increase `target_accept` or reparameterize.

We recommend running at least 4 chains for robust computation of convergence diagnostics

The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details

The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details

PyMC horseshoe (NUTS, sigma-scaled to match): 161 s   top genes include YOAB_at, ARGF_at
beta agreement, all 4088 genes: r = 0.80  (the strong signals match; the near-zero
cloud wobbles -- those coefficients are ~0 with large posterior uncertainty)
-> ~16x slower than the Gibbs: general-purpose NUTS strains at p=4088 where the
   Bhattacharya-based Gibbs is trivial (9.8 s).

Results — the same signal, and a lesson about samplers

• Both Bayesian engines agree on what is real, not on what is noise. With the priors matched, the from-scratch Gibbs and PyMC's NUTS pick out the same dominant gene, YOAB (with YXLE the clear runner-up in the from-scratch Gibbs), and R's lasso (next notebook) puts YOAB first too. But across all 4088 coefficients the two samplers correlate only $r\approx0.80$ — they do not even agree on the second-tier gene (the agreement scatter): the strong signals coincide, while the ~4086 near-zero coefficients scatter. That is the posterior being honest — those genes are $\approx0$ with large uncertainty, so their exact values are not meant to be trusted or reproduced. The reliable outputs of a $p\gg n$ fit are the shrinkage weights and the top-gene ranking, not individual near-zero estimates.
• The sampler matters enormously at scale. The specialised Gibbs finishes in ~10 seconds; general-purpose NUTS takes ~2.7 minutes (~16x longer) and still mixes poorly (low ESS) against the horseshoe's funnel geometry in 4088 dimensions — which is part of why its near-zero estimates are noisier. The Bhattacharya trick — solve $n\times n$, never $p\times p$ — is what makes genome-scale Bayesian selection routine. When $p\gg n$, the algorithm, not just the prior, is the enabling idea.
• Continuous vs discrete selection. The horseshoe answers the same question as SSVS (Part 1) but by shrinking rather than switching: no $2^p$ model space, no binary indicators to mix over — just a heavy-tailed prior that lets a few coefficients through. That scalability is precisely why global-local priors dominate the $p\gg n$ regime.

References

• Carvalho, C. M., Polson, N. G. & Scott, J. G. (2010). The horseshoe estimator for sparse signals. Biometrika 97, 465–480.
• Makalic, E. & Schmidt, D. F. (2016). A simple sampler for the horseshoe estimator. IEEE Signal Proc. Letters 23, 179–182.
• Bhattacharya, A., Chakraborty, A. & Mallick, B. K. (2016). Fast sampling with Gaussian scale-mixture priors in high-dimensional regression. Biometrika 103, 985–991.
• Bühlmann, P., Kalisch, M. & Meier, L. (2014). High-dimensional statistics with a view toward applications in biology. Annual Review of Statistics 1, 255–278.

Data: ribo_data.csv — riboflavin production in B. subtilis: n = 71 samples, p = 4088 gene-expression predictors (hdi::riboflavin). Next: hshd_R.ipynb — the R cross-check (glmnet cross-validated lasso).