Investigate maximum feasible QRF subsample size for PUF imputation

[baogorek]

Context

In #537, puf_impute.py fuses PUF tax variables onto CPS records via quantile random forests (QRF). The QRF is trained on PUF donor records and then predicts ~104 tax variables for each CPS recipient, matching on demographic predictors (age, sex, filing status, dependents). To keep this tractable, the CPS side is stratified-subsampled to 20K records (force-including the top 0.5% by AGI to preserve the high-income tail).

The 20K target was chosen pragmatically but hasn't been stress-tested — it's unclear whether 20K is near the ceiling or if we could go significantly higher. Larger subsamples should improve imputation quality (especially in the tails), so it's worth finding the practical upper bound.

Questions to answer

• What is the maximum subsample size that fits in memory during QRF training (given ~104 target variables batched in groups)?
• How does imputation accuracy (e.g., distribution of imputed AGI, top-tail fidelity) change as subsample size increases from 20K toward the full CPS?
• Is there a point of diminishing returns where more data doesn't meaningfully improve the QRF fits?
• Does the current batching strategy (grouping target variables) affect the memory ceiling?

Suggested approach

1. Benchmark memory usage and wall time at 20K, 40K, 60K, etc.
2. Compare imputed distributions (quantiles, max AGI) across subsample sizes.
3. Identify the largest feasible size given CI runner constraints (~7 GB RAM).

Related

• CPS top-coding caps AGI at $6.26M — zero observations above $10M in any state #530 / Add PUF + source impute modules, fix AGI ceiling (issue #530) #537 — PUF imputation module where the 20K default was introduced

🤖 Generated with Claude Code

[juaristi22]

QRF Subsample Size Benchmark Results

Benchmarked PUF subsample sizes for _batch_qrf imputation (67 variables) to find the largest size that fits within the 7 GB GitHub Actions memory limit. Ran on Modal (32 GB RAM, 8 CPU, Python 3.13) using PUF_2024 for training and CPS_2024_Full for test demographics.

Resource Usage

Size	Wall Time	Peak RSS	Headroom	Fits 7 GB?
20,000	16.9 min	6.14 GB	+0.86 GB	✅
40,000	30.1 min	6.65 GB	+0.35 GB	✅
45,000	41.9 min	6.80 GB	+0.20 GB	✅
50,000	44.5 min	6.92 GB	+0.08 GB	✅
55,000	50.6 min	7.05 GB	-0.05 GB	❌
60,000	44.2 min	7.20 GB	-0.20 GB	❌
80,000	58.2 min	7.68 GB	-0.68 GB	❌

The boundary is between 50k (82 MB to spare) and 55k (51 MB over).

Distributional Accuracy (KS Statistic, lower = better)

Variable	20k	40k	45k	50k	55k	60k	80k
employment_income	.058	.062	.060	.060	.061	.059	.055
long_term_capital_gains	.050	.028	.029	.027	.030	.029	.023
social_security	.037	.041	.041	.048	.038	.043	.042
taxable_pension_income	.035	.041	.043	.042	.038	.037	.049
self_employment_income	.031	.036	.035	.035	.033	.035	.035
rental_income	.025	.027	.027	.026	.027	.028	.028
farm_income	.003	.003	.002	.002	.003	.003	.003
trad_ira_contributions	.013	.015	.014	.015	.015	.016	.016
deductible_mort_interest	.022	.020	.023	.022	.024	.026	.029
farm_operations_income	.012	.014	.014	.014	.015	.014	.015
Mean KS	.029	.029	.029	.029	.029	.029	.029

Accuracy is essentially flat across all sizes — mean KS is 0.029 everywhere. The only variable with a clear size-dependent improvement is long_term_capital_gains (heavy-tailed, KS 0.050→0.023), which continues improving well past the memory limit.

Recommendation

Use 40,000 records. Rationale:

• Same accuracy as 50k — mean KS identical (0.029), no variable improves meaningfully between 40k and 50k
• 4x more headroom — 358 MB (5%) vs 82 MB (1.2%), leaving room for OS fluctuations, GH Actions overhead, and other pipeline steps
• 14 min faster — 30 min vs 44 min wall time
• long_term_capital_gains tail fidelity is the only thing larger samples help with, and a targeted fix (oversampling high-AGI strata) would be more effective than increasing overall sample size