Leave-out variance components

The plug-in AKM decomposition is biased. Worker and firm effects are estimated from a handful of observations each, and that estimation error inflates their measured variance and attenuates the covariance between them, so the plug-in numbers overstate how much of wage inequality is worker and firm heterogeneity, and understate how much is sorting. leave_out_kss applies the bias correction of Kline, Saggio and Sølvsten (2020). This example runs it on a simulated panel whose true effects are known, so you can watch the bias appear and be removed. The leave-out variance components section has the full story; this page follows the script.

It is the slowest of the examples, because it runs the correction seven times.

One call

from hdfe_stream import leave_out_kss
from hdfe_stream.simulate import simulate_akm

simulate_akm(n_workers=8_000, seed=17, keep_effects=True).write_parquet("panel.parquet")
lo = leave_out_kss("log_earn ~ age_squared + age_cubed | worker_id + firm_id",
                   "panel.parquet", workdir="scratch", n_draws=250, seed=0)
print(lo.summary())
leave-out variance components (firm_id and worker_id)
  67,913 observations in 13,404 matches, 9,491 of them movers'
  leaving out a match, 250 draws; max mover leverage 0.937202; mean sigma2 0.018147
  sample: 8,000 -> 7,993 worker_id, 533 -> 531 firm_id after pruning to the leave-one-out connected set

component                plug-in        bias   leave-out
--------------------------------------------------------
var(psi)                0.064284    0.001372    0.062912
var(alpha)              0.174041    0.007164    0.166877
cov(psi, alpha)         0.027284   -0.001314    0.028598

This prunes to the leave-one-out connected set, fits, and decomposes, in that order: pruning changes the estimation sample, so it has to happen before the fit. What is left out is a whole worker–firm match, every year of a spell, as the reference implementations do. keep_effects=True makes the simulator return the effects it drew, as true_worker_effect and true_firm_effect.

Against the truth

The estimand is an observation-level moment of the true effects, taken over the sample that survived pruning, so it is computed on lo.fit.sample()’s rows in the sample:

component              truth    plug-in  leave-out  plug-in err  leave-out err
------------------------------------------------------------------------------
var(psi)             0.05864    0.06428    0.06291      0.00564        0.00427
var(alpha)           0.15976    0.17404    0.16688      0.01428        0.00712
cov(psi, alpha)      0.03051    0.02728    0.02860      0.00323        0.00191

The plug-in errors run in the direction theory predicts: both variances too high, the covariance too low. With 8.5 observations per worker the bias is modest; it grows as that number falls, so a short panel needs the correction more, not less.

What pruning cost

Here pruning removed 7 workers (2 of them articulation points) and 2 firms, 59 rows, since the simulated panel is well connected. lo.pruning has the counts, and lo.fit.sample() labels the rows it dropped "leave_one_out". On administrative data KSS report losing about half the firms.

In steps, keeping the regression

leave_out_kss is a wrapper. If you want the regression in its own right, do the steps yourself: connected_components finds the leave-one-out connected set ahead of time, and the fit’s own leave_out_kss method gives the same components.

from hdfe_stream import connected_components, feols_stream

panel = connected_components("panel.parquet", "worker_id + firm_id",
                             leave_one_out=True).largest()
fit = feols_stream(fml, panel, workdir="scratch", stream="worker_id")
step = fit.leave_out_kss(n_draws=250, seed=0)

How much the approximation moves the answer

The leverages are approximated by random projection, so the answer carries Monte Carlo noise, and the way to know how much is to vary the seed. Across four seeds at 250 draws the spread was 0.2–1.4% of the estimates, small next to their sampling uncertainty. Draws and reproducibility measures it properly and says what to report in a paper.

Standard errors and weak identification

se=True adds standard errors and 95% intervals. The example leaves out an observation here (leave_out="observation"), so that all three components get one; leaving out a match reports them for var(psi) and the covariance. With standard errors comes a check on whether the normal interval is justified at all, which on this well-connected panel it is (q = 0 for every component):

component                plug-in        bias   leave-out         se              95% interval
---------------------------------------------------------------------------------------------
var(psi)                0.064284    0.003421    0.060863   0.001336      [0.058245, 0.063481]
var(alpha)              0.174041    0.013201    0.160841   0.001301      [0.158291, 0.163391]
cov(psi, alpha)         0.027284   -0.002529    0.029813   0.000462      [0.028908, 0.030718]

All three intervals contain the truth. The last part of the example builds a panel with a bottleneck, two blocks of firms joined by six movers, where the normal interval is not justified, and the output adds KSS’s q = 1 interval, which stays valid there. Both are explained in standard errors and weak identification.

The full script: examples/leave_out_kss.py
"""Leave-out (Kline-Saggio-Solvsten) variance components.

The plug-in AKM decomposition is biased. Worker and firm effects are estimated
from a handful of observations each, and that estimation error inflates their
measured variance and attenuates the covariance between them -- so the plug-in
numbers overstate how much of wage inequality is worker and firm heterogeneity,
and understate how much is sorting. With few observations per worker the bias is
large; this is the "limited mobility bias" of Andrews et al. (2008).

Kline, Saggio and Solvsten (2020) remove it exactly, without assuming the errors
are identically distributed, by subtracting a bias term built from each
observation's statistical leverage. Computing those leverages exactly would need
one linear solve per observation, so they are approximated by random projection
-- which is what makes the whole thing affordable here.

What is "left out" is a whole worker-firm match -- every year of a spell --
following LeaveOutTwoWay, VarianceComponentsHDFE.jl and xhdfe. Earnings errors
are correlated within a spell, and leaving out one year of it would treat the
others as independent of it. `leave_out="observation"` leaves out a single year
instead, and is what standard errors need; both appear below.

    python examples/leave_out_kss.py
"""

import numpy as np
import polars as pl

from hdfe_stream import connected_components, feols_stream, leave_out_kss
from hdfe_stream.simulate import simulate_akm, simulate_bottleneck
from simulated_data import OUTPUT, workdir

# The simulator can hand back the worker and firm effects it generated, so for
# once the right answer is known. `log_earn` is built as
#
#     alpha + psi + 0.08 age_squared - 0.012 age_cubed + noise
#
# so the model below is correctly specified, and the true variance components
# are moments of columns we have.
OUTPUT.mkdir(parents=True, exist_ok=True)
data = OUTPUT / "leave_out_panel.parquet"
if not data.exists():
    simulate_akm(n_workers=8_000, seed=17, keep_effects=True).write_parquet(data)

FML = "log_earn ~ age_squared + age_cubed | worker_id + firm_id"

# KSS operate at 250-500 draws on panels with millions of effects, and report
# the error *falling* with sample size at fixed draws -- so this is a budget,
# not something that has to scale with the data. The last section shows how to
# check it is enough for your own.
DRAWS = 250

# ----------------------------------------------------------------- one call
# This prunes, fits and decomposes. The order matters and is not negotiable:
# pruning changes the estimation sample, so it has to happen before the fit.
lo = leave_out_kss(FML, str(data), workdir=workdir("leave_out"),
                   n_draws=DRAWS, seed=0)

print()
print(lo.summary())

# ------------------------------------------------------- against the truth
# The estimand is an observation-level moment of the true effects, taken over
# the sample that survived pruning -- so it must be computed on that sample,
# not on the panel we started with.
kept = (lo.fit.sample().filter("in_sample")
        .select("true_worker_effect", "true_firm_effect").collect())
alpha = kept["true_worker_effect"].to_numpy()
psi = kept["true_firm_effect"].to_numpy()
# Like the reference implementations, the components divide by n - 1.
truth = {
    "var(psi)": float(np.var(psi, ddof=1)),
    "var(alpha)": float(np.var(alpha, ddof=1)),
    "cov(psi, alpha)": float(np.cov(psi, alpha, ddof=1)[0, 1]),
}

print(f"\n{'=' * 72}\nagainst the answer the data was built from\n{'=' * 72}")
print(f"{'component':<18}{'truth':>10}{'plug-in':>11}{'leave-out':>11}"
      f"{'plug-in err':>13}{'leave-out err':>15}")
print("-" * 78)
for key, true_value in truth.items():
    plug_in_error = abs(lo.plug_in[key] - true_value)
    corrected_error = abs(lo.leave_out[key] - true_value)
    print(f"{key:<18}{true_value:>10.5f}{lo.plug_in[key]:>11.5f}"
          f"{lo.leave_out[key]:>11.5f}{plug_in_error:>13.5f}"
          f"{corrected_error:>15.5f}")

print("\nThe plug-in errors run in the direction theory predicts: both variances")
print("too high, the covariance too low. Correcting them changes the story about")
print("how much of wage variance is sorting rather than heterogeneity.")
print(f"\nobservations per worker: "
      f"{lo.n_obs / lo.diagnostics['n_levels']['worker_id']:.1f} -- the bias grows")
print("as this falls, so a short panel needs the correction more, not less.")

# --------------------------------------------------------------- the sample
# Pruning is not a detail. Leave-out estimation needs every effect to stay
# estimable when any one observation is dropped, which fails at workers who are
# the only link between two parts of the worker-firm network. Here the simulated
# panel is well connected and little is lost; on administrative data KSS report
# losing roughly half the firms.
p = lo.pruning
print(f"\n{'=' * 72}\nwhat pruning cost\n{'=' * 72}")
for dim in ("worker_id", "firm_id"):
    print(f"  {dim:<9} {p['levels_before'][dim]:,} -> {p['levels_after'][dim]:,}")
print(f"  articulation-point workers removed: {p['cut']:,}")
print(f"  components before pruning: {p['components_before']:,}")

# The fit says which rows pruning dropped: sample() labels them "leave_one_out".
# connected_components finds the same set ahead of time, to fit other things on.
dropped = lo.fit.sample().filter(pl.col("dropped_because") == "leave_one_out")
print(f"\n  rows dropped by pruning: {dropped.select(pl.len()).collect().item():,}")

# ------------------------------------------------------ the two-step version
# `leave_out_kss` is a wrapper. If you want the regression in its own right --
# coefficients, residuals, fixed effects -- do the steps yourself. Leaving out a
# match reads only the fit's residual file, which every fit keeps, and refits on
# the collapsed matches. (Leaving out an observation reads the sorted rows and
# cell tables a fit otherwise deletes, so that needs keep_intermediates=True.)
print(f"\n{'=' * 72}\nthe same thing in steps, keeping the regression\n{'=' * 72}")
panel = (connected_components(str(data), "worker_id + firm_id", leave_one_out=True,
                              verbose=False)
         .largest())
fit = feols_stream(FML, panel, workdir=workdir("leave_out_steps"),
                   stream="worker_id", verbose=False)
step = fit.leave_out_kss(n_draws=DRAWS, seed=0)

print(f"  regression: {dict(zip(fit.coefnames, fit.beta.round(6)))}")
print("              (the data was built with 0.08 and -0.012)")
print(f"  and the same components: "
      f"{np.allclose([step.leave_out[k] for k in truth], [lo.leave_out[k] for k in truth])}")

# ------------------------------------------------------------ how many draws
# The leverages are approximated, so the answer carries Monte Carlo noise. The
# accuracy does not have to grow with the data -- KSS report the error *falling*
# with sample size at fixed draws -- but you should know how big it is, and the
# only way to find out is to vary the seed.
print(f"\n{'=' * 72}\nhow much the approximation moves the answer\n{'=' * 72}")
runs = [step] + [fit.leave_out_kss(n_draws=DRAWS, seed=s) for s in (1, 2, 3)]
values = np.array([[run.leave_out[k] for k in truth] for run in runs])
print(f"across {len(runs)} seeds at {DRAWS} draws:\n")
print(f"{'component':<18}{'mean':>11}{'spread':>11}{'relative':>11}")
for index, key in enumerate(truth):
    column = values[:, index]
    print(f"{key:<18}{column.mean():>11.5f}{np.ptp(column):>11.5f}"
          f"{np.ptp(column) / abs(column.mean()):>11.2%}")

print("\nThat spread shrinks with more draws; docs/kss/draws.md reports it across five")
print("seeds at 250 and 1,000 draws. What matters is that it is small next to the")
print("sampling uncertainty of the components themselves -- about a tenth of the")
print("standard error when leaving out a match -- so the approximation is not")
print("what limits this answer.")
# --------------------------------------------------------- standard errors
# `se=True` adds standard errors and 95% intervals. This section leaves out an
# observation, so that all three components get one. Leaving out a match
# reports them for var(psi) and the covariance only, following LeaveOutTwoWay's
# leave_out_COMPLETE, whose match-level option is marked beta (see
# docs/kss/inference.md).
# The estimator is a quadratic form y'Cy whose kernel has a zero diagonal by
# construction, which is what makes it unbiased and what makes its variance
#
#     V = 4 sum_i sigma2_i (C mu)_i^2 + 2 tr(C Omega C Omega)
#
# hold with no normality assumption. The first term alone is conservative, so it
# is reported as well; `se_trace=False` computes only that and saves roughly
# three solves per draw.
print(f"\n{'=' * 72}\nstandard errors\n{'=' * 72}")
with_se = leave_out_kss(FML, str(data), workdir=workdir("leave_out_se"),
                        n_draws=DRAWS, seed=0, leave_out="observation", se=True)
print()
print(with_se.summary())

# Do the intervals contain the answer the data was built from? With three
# components and one draw of the data this is an anecdote, not a coverage study
# -- docs/kss/inference.md reports coverage measured over many redraws -- but it
# is the check worth making on your own data's scale.
print(f"\n{'component':<18}{'truth':>10}{'leave-out':>11}{'se':>10}"
      f"{'z':>7}{'covers':>8}")
for key, true_value in truth.items():
    se = with_se.se.se[key]
    z = (with_se.leave_out[key] - true_value) / se
    print(f"{key:<18}{true_value:>10.5f}{with_se.leave_out[key]:>11.5f}"
          f"{se:>10.5f}{z:>7.2f}{'yes' if abs(z) < 1.96 else 'no':>8}")

# Those intervals come with a check on whether they are justified at all. The
# normal approximation needs no single eigen-direction to dominate the
# estimator's variance; `q` counts the ones that do, by KSS's threshold of 1/10
# on lambda^2 / sum lambda^2, and only q = 0 justifies the interval. A bottleneck
# in the mobility network -- two groups of firms joined by a few movers -- is
# what breaks it. This panel is well connected, so q should be 0 throughout.
print("\nis the normal interval justified?  q per component:",
      with_se.weak_id.q)

# `se.theta` recomputes the same estimand as y'Cy -- per-row B_ii by random
# projection rather than the coefficient-space trace.
# It shares almost no code with the point estimate, so agreement is a real check.
print("\ncross-check, two independent routes to the same estimand:")
for key in truth:
    gap = abs(with_se.se.theta[key] - with_se.leave_out[key]) / with_se.se.se[key]
    print(f"  {key:<18}{with_se.se.theta[key]:>11.6f} vs "
          f"{with_se.leave_out[key]:>11.6f}   {gap:.2f} se apart")

# ------------------------------------------------- when the interval fails
# Everything above was on a well-connected panel. A bottleneck in the mobility
# network -- two groups of firms joined by a handful of movers -- is where the
# normal interval stops being justified: one linear combination of the effects
# is estimated far worse than the rest, and it dominates the estimator's
# variance. The output then flags it, and for each flagged component adds KSS's
# q = 1 interval, which stays valid there.
print(f"\n{'=' * 72}\na weakly identified panel\n{'=' * 72}")
weak_data = OUTPUT / "bottleneck_panel.parquet"
simulate_bottleneck(n_bridge=6, keep_effects=True).write_parquet(weak_data)
weak = leave_out_kss("log_earn ~ 1 | worker_id + firm_id", str(weak_data),
                     workdir=workdir("leave_out_weak"), n_draws=DRAWS, seed=0,
                     leave_out="observation", se=True, verbose=False)
print(weak.summary())

# Against the truth, once more. One draw of the data is an anecdote;
# docs/kss/inference.md reports coverage over many. Note the q = 1 interval is
# not symmetric about the estimate: the badly estimated direction enters
# squared, so its uncertainty is skewed, which a normal interval cannot express.
kept = (weak.fit.sample().filter("in_sample")
        .select("true_worker_effect", "true_firm_effect").collect())
a_true = kept["true_worker_effect"].to_numpy()
p_true = kept["true_firm_effect"].to_numpy()
weak_truth = {"var(psi)": float(np.var(p_true, ddof=1)),
              "var(alpha)": float(np.var(a_true, ddof=1)),
              "cov(psi, alpha)": float(np.cov(p_true, a_true, ddof=1)[0, 1])}
print(f"\n{'component':<18}{'truth':>10}{'normal covers':>15}{'q=1 covers':>12}")
for key, true_value in weak_truth.items():
    se = weak.se.se[key]
    normal = abs(weak.leave_out[key] - true_value) <= 1.96 * se
    lo, hi = weak.se.weak[key]["interval"] if key in (weak.se.weak or {}) else (
        float("nan"), float("nan"))
    print(f"{key:<18}{true_value:>10.5f}{'yes' if normal else 'no':>15}"
          f"{'yes' if lo <= true_value <= hi else 'no':>12}")

# ------------------------------------------------------------------ weights
# `weights=` works here too, and then the estimand is the *weighted* variance
# decomposition -- the weighted variances and covariance of the effects, over
# observations. That one choice fixes everything else: the leverages become the
# ones of the square-root-weight hat matrix, which is the only version that is a
# projection, and sigma2 carries a factor of w. Nothing to pass but the weights.

print("\nFor a paper: fix the seed and the number of draws and report both, and")
print("report the estimate across several seeds rather than one. The claim worth")
print("making is that your result is stable to the approximation -- not that one")
print("particular draw reproduces.")