GPU_REML

GPU_REML is a GPU-accelerated statistical framework for SNP heritability estimation, genetic-variance decomposition, and downstream mixed-model inference at biobank scale.

The central statistical problem is restricted maximum likelihood (REML) estimation in linear mixed models where the genetic covariance is defined by one or more genomic relationship matrices (GRMs). These models are the standard language for estimating SNP heritability and asking how heritable signal is distributed across chromosomes, annotations, MAF bins, LD environments, or user-defined genomic regions.

In the standard formulation, GPU_REML fits a linear mixed model

$$ \begin{aligned} y &= X\beta + u_1 + \cdots + u_G + e, \\ u_g &\sim \mathcal{N}(0, \sigma_g^2 K_g), \\ e &\sim \mathcal{N}(0, \sigma_e^2 I), \\ V(\theta) &= \sum_g \theta_g K_g + \theta_e I. \end{aligned} $$

The restricted log likelihood is

$$\ell_R(\theta)=-\frac{1}{2}\left[\log|V(\theta)|+\log|X^TV(\theta)^{-1}X|+y^TP(\theta)y\right]$$

where

$$P(\theta)=V(\theta)^{-1}-V(\theta)^{-1}X\left(X^TV(\theta)^{-1}X\right)^{-1}X^TV(\theta)^{-1}.$$

Each K_g is a genotype-defined covariance component. Let a_g = tr(K_g) / n and b_j = tr(R_j) / n denote average diagonal atoms. SNP heritability is estimated from sample-average variance contributions:

$$ h^2 = \frac{\sum_g \theta_g a_g} {\sum_g \theta_g a_g + \sum_j \eta_j b_j} $$

For unit-trace kernels and one identity residual this reduces to the familiar ratio of raw variance-component sums. The trace-weighted form is required for admixed components and effective-rank SMILE normalization.

The computational obstacle is that the natural GRM representation is dense: constructing, storing, and repeatedly factorizing n x n kernels becomes the bottleneck as cohorts, marker counts, and component counts increase. GPU_REML therefore keeps the statistical REML model but changes how each K_g is applied numerically. Instead of materializing a GRM, each covariance component is represented as a matrix-free genotype operator:

$$ K_g v = \frac{Z_g (Z_g^T v)}{m_{\mathrm{eff},g}} $$

Genotype blocks are decoded on the host, streamed to the GPU, and multiplied in batches. Evaluating and optimizing the REML likelihood without explicit GRMs leads to the main numerical machinery in GPU_REML: block PCG solves, Hutchinson trace estimates, SLQ log-determinant estimates, constrained AI/Fisher updates, and projected-core preconditioning.

The goal is not only to produce one whole-genome heritability number. GPU_REML is designed as a method-development workbench for comparing single-GRM and multi-GRM covariance representations. It also includes a SMILE-inspired weighted-GRM extension, with explicit attribution to the original JianqiaoWang/SMILE project. This path adapts the SMILE idea of introducing a SNP-space weight matrix W into the genetic covariance, while implementing the form that matches GPU_REML's matrix-free REML engine: a block-diagonal W, evaluated without materializing the sample-space kernel:

$$K_g=\frac{X_gW_gX_g^T}{c_g},\quad W_g=\mathrm{blockdiag}(W_{g,1},\ldots,W_{g,B}),\quad c_g=\frac{\mathrm{tr}(X_gW_gX_g^T)}{n}$$

Each W_{g,i} must be a finite symmetric positive-semidefinite dense block. GPU_REML treats this as a trusted-input contract and does not run a cubic-time PSD check. Blocks inside one GRM are summed into one variance component; multiple GRM groups can be supplied when a multi-component REML model is desired.

The sparse fixed-effect path uses the fitted covariance V(theta) to define a penalized GLS likelihood over candidate SNP effects:

$$ (\hat\alpha_\lambda,\hat b_\lambda)=\arg\min_{\alpha,b}\frac{1}{2}(y-C\alpha-Z_Sb)^TV(\theta)^{-1}(y-C\alpha-Z_Sb)+\lambda|b|_1. $$

The sparse command first alternates this weighted-LASSO step with information-corrected REML for the residual y-C alpha-Z b in the (n-\operatorname{rank}(C)) dimensional space orthogonal to the complete nuisance design. Supplying the full C matrix to the REML routine profiles its unpenalized coefficients at every candidate covariance. Since (P_C C=0), using a residual that already subtracts the current nuisance score is algebraically equivalent to applying (P_C) to y-Z b, while retaining the numerically convenient residual scale. For an orthonormal basis B of the selected SNP span, the covariance objective adds -0.5 log|B' P_C B| to the log likelihood. The sparse quadratic subtracts tr{B (B' P_C B)^-1 B'}/n. These are paired local mean-information corrections; the Lasso residual quadratic itself retains P_C. Within each candidate problem, coordinate descent is accepted solely when the active and inactive score-KKT conditions pass at the configured numerical tolerance; coefficient change is only an active-set scheduling heuristic. Every evaluated lambda on the complete path must pass that finite-tolerance certificate. An unsolved path point is reported as a numerical failure rather than being skipped in favor of the lambda-max empty model. At each covariance update, held-out validation predictive R² (1 - SSE/SST) selects lambda before the variance-component update, and the full-marker KKT scan adds any omitted violating variants to the candidate set. The KKT tolerance is matched to the ordinary PCG precision; an independent finite-PCG score on candidate coordinates is diagnostic rather than a second rejection gate. The sparse pipeline standardizes the phenotype once at entry and uses that single analysis scale throughout every LASSO, REML, CHIVE, and prediction step. The outer loop stops when both the complete fitted mean and primary COHERIT heritability stabilize, or after the configured maximum number of updates. It then checks a covariance-aligned LASSO update against the same tolerances (fitted mean 0.05, absolute h² change 0.01). If alignment breaks stability and the update budget remains, that LASSO is reused in the next variance update. Only a stable final aligned pair is labelled converged. Reaching the outer limit is reported as a warning and does not invalidate a finite final pair. If that final update itself is unavailable, the most recent complete covariance-aligned pair is returned with a warning. The selected lambda/lambda-max ratio is then frozen before the combined train+validation refit. A prediction phenotype is never accepted or read. Because that ratio is already fixed, each final-refit LASSO block solves only the lambda-max warm-start point and the exact target point, rather than recomputing the unused validation grid. The run stops with this covariance-aligned LASSO pair and reports the COHERIT estimator.

The production gpu-reml-sparse runner has exactly two modes. fixed uses one whole-genome GRM or an exhaustive user-supplied --component-spec. adaptive starts from K=1, freezes the converged K=1 sparse mean, and tests balanced boundaries along one deterministic LD-score ordering. Independent trace probes estimate nuisance-adjusted covariance scores and their Fisher information. The maximum absolute standardized score is calibrated against one shared Gaussian quadratic-form reference over all available candidates. Its joint p-value must pass --split-alpha before the maximizing boundary is added. A covariance-only REML refit follows each split, keeping both the K=1 sparse mean and its information-corrected sparse quadratic fixed. The process stops when evidence is insufficient or no contrast is identifiable. The common quadratic reference retains dependence and non-Gaussian score tails. Non-finite inputs or unresolved trace precision cause an explicit numerical error. The selected endpoint partition then receives one full validation-lambda alpha/theta refit. Both modes finally freeze the selected lambda ratio and automatically warm-refit on training + validation samples.

Adaptive scores use --score-core-rank (64), --score-reference-samples (16383), --score-trace-probes (512 initial probes per independent group), --score-trace-max-probes (4096), --score-trace-tol (0.05), and --score-trace-seed. Score JSON schema 3 records global_p_value, signed standardized scores, the joint critical value, and trace/PCG diagnostics. Large arrays use temporary files beside the score output, removed on completion or error; GPU reference evaluation uses bounded batches. The fitted null and finite probes give a plug-in p-value, not a finite-sample guarantee for the complete adaptive path. See the score-test derivation for the formulas, precision criterion, and validation procedure.

Adaptive mode accepts an aligned --ld-score table with ID and ld_score columns. If it is omitted, the workflow invokes PLINK2 on the training samples, computes the pilot-compatible 1 + sum(r²) unphased-dosage score within the configured physical window (1 Mb by default), validates it by variant ID, and caches the compact score table. PLINK's large temporary pair table is removed after successful conversion.

Sparse-run output semantics are deliberately explicit:

• estimator_mode is coherit and computed_estimators contains only h2_chive.
• New sparse runs use output schema version 11. Raw/standardized duplicate estimator fields and downstream phenotype-scale conversions no longer exist.
• In each outer_update history record, theta, coherit_h2, and validation predictive R² describe the same selected LASSO state. The subsequent covariance update is recorded separately as theta_after_variance_update and coherit_h2_after_variance_update; its h² convergence fields use the variance_update_ prefix. The validation audit also records the aligned h². q_sparse_after_variance_update includes information and nuisance effects recomputed at that new covariance, rather than reusing the pre-update Q.
• BASIL validation reuses its training dual and marker scores after a true residual check; only inaccurate columns are refined. Fixed-K variance blocks reuse response-independent REML probe products and single-GRM Lanczos work. These caches do not cross sample sets or GRM partitions. Resuming a pipeline also checks a source-code digest; changed code requires a fresh work directory.
• var_components_lasso_ml contains one kernel coefficient per fixed GRM followed by the identity-residual variance. The covariance remains H = sum_g theta_g K_g + theta_e I; the kernels are not rescaled.
• genetic_trace_atoms records a_g = tr(K_g) / n on the samples used for that fit, and grm_variance_scale is trace_weighted. Background genetic variance is sum_g theta_g a_g, including when missing genotypes are mean-imputed. Atoms are recomputed for each sample set and GRM partition.
• h2 is the primary total estimate. For a valid covariance-aligned LASSO pair it is (q_chive + sum_g theta_g a_g) / (q_chive + sum_g theta_g a_g + theta_e). The same trace-weighted background is used by outer-loop h² checks and adaptive covariance-refit summaries. q_chive already measures variance on the analysis scale and is not multiplied by a GRM trace atom.
• h2_chive_guarded is the validated counterpart of h2_chive and equals the top-level h2 field whenever the COHERIT branch is valid.
• q_chive_components retains the squared fitted-mean and residual-correction terms, and term3_mean_uncertainty_subtracted, which is subtracted from their sum to form the COHERIT sparse variance contribution; the squared fitted-mean term is not exposed as a standalone heritability estimator.
• Sparse prediction emits only the lasso_* branch: the Lasso fixed-SNP score plus its matched background BLUP.
• --compute-effects emits nuisance effects and one source-order SNP table containing the sparse effect, matched background-BLUP effect, and their sum.
• Every sparse quadratic and variance component is already on the one analysis scale established by input phenotype standardization. The summary retains only the input normalization metadata needed to transform external outcomes.
• lasso_branch_valid validates the COHERIT output. An invalid branch has JSON null in its guarded fields.
• Ordinary REML is a separate baseline and is never substituted into any sparse estimator.
• Sparse prediction follows the same contract. A valid run emits only lasso_* scores, with no baseline substitution.
• In REML history, accepted refers to the current line-search candidate. A terminal ll_down rejects that candidate and returns the most recent accepted variance vector; an intermediate BCD variance block records this as a no-update and continues. converged is true both when an accepted relative-likelihood increment meets the threshold and when every backtracking candidate decreases the likelihood (ll_down).

Research Use Cases

GPU_REML is most useful when the scientific question requires more than a single whole-genome GRM. It is designed to make multi-GRM REML practical by keeping each covariance component as a streamed genotype operator rather than a stored dense matrix. This is the main advantage of the project: users can expand from one GRM to many GRMs while keeping wall time low through GPU batched products and keeping CPU memory controlled by avoiding explicit n x n GRM storage.

Typical use cases include:

• comparing single-GRM and multi-GRM heritability estimates on the same cohort;
• decomposing SNP heritability across chromosomes, LD environments, MAF bins, annotations, or custom SNP sets;
• fitting many covariance components without constructing and storing one dense GRM per component;
• benchmarking alternative covariance representations under matched phenotype, covariate, and sample filters;
• testing SMILE-style block-diagonal weighted GRMs where dense W_i blocks encode local SNP covariance or effect-correlation structure.

Installation

GPU_REML requires Python 3.10 or newer. For large runs, install a GPU-enabled JAX build before installing GPU_REML.

git clone https://github.com/Asiandier/GPU_REML.git
cd GPU_REML
python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip

Install JAX for the local CUDA driver following the official JAX installation guide. For current NVIDIA CUDA pip wheels, the command is typically:

python -m pip install -U "jax[cuda13]"

Then install GPU_REML:

python -m pip install -e .

Optional PGEN support:

python -m pip install -e ".[pgen]"

Development install:

python -m pip install -e .

Check that JAX can see the GPU:

python - <<'PY'
import jax
print(jax.devices())
PY

CPU-only JAX is sufficient for small examples. Large REML jobs are intended for GPU execution.

Quick Start

Single-GRM REML from PLINK1 BED:

gpu-reml \
  --bed-prefix /path/to/data \
  --pheno-txt pheno.txt \
  --covar-txt covar.txt \
  --out-prefix out/reml

PGEN input:

gpu-reml \
  --pgen-prefix /path/to/data \
  --pheno-txt pheno.txt \
  --covar-txt covar.txt \
  --out-prefix out/reml

Multiple GRMs from multiple BED prefixes:

gpu-reml \
  --bed-prefix /path/to/grm1,/path/to/grm2,/path/to/grm3 \
  --pheno-txt pheno.txt \
  --covar-txt covar.txt \
  --out-prefix out/multi_grm

Arbitrary SNP components from one genotype file:

gpu-reml \
  --bed-prefix /path/to/data \
  --component-spec components.json \
  --pheno-txt pheno.txt \
  --covar-txt covar.txt \
  --out-prefix out/partitioned

Z-score one-shot weak-component merge:

gpu-reml \
  --merge \
  --bed-prefix /path/to/data \
  --component-spec fine_components.json \
  --pheno-txt pheno.txt \
  --covar-txt covar.txt \
  --keep-path keep.txt \
  --out-prefix out/zscore_merge

This mode first fits the fine component model, computes component-level Wald-style z-scores from the fitted variance components and AI matrix, keeps components with z >= 1.6448536269514722, merges all weaker components into one background GRM, and refits once.

SMILE-style block-diagonal weighted GRM:

gpu-reml \
  --smile \
  --bed-prefix /path/to/data \
  --w-files W_block_1.npy,W_block_2.npy,W_block_3.npy \
  --pheno-txt pheno.txt \
  --covar-txt covar.txt \
  --out-prefix out/smile

Multiple weighted GRMs are supplied as semicolon-separated groups. Blocks within a group are summed into one GRM, and each group receives its own variance component:

gpu-reml \
  --smile \
  --bed-prefix /path/to/data \
  --grm-groups 'A_1.npy,A_2.npy;B_1.npy,B_2.npy' \
  --pheno-txt pheno.txt \
  --covar-txt covar.txt \
  --out-prefix out/smile_multi

Fixed single-GRM sparse COHERIT with validation predictive-R² lambda selection and an automatic train+validation final refit:

gpu-reml-sparse \
  --mode fixed \
  --bed-prefix /path/to/data \
  --pheno-txt train.pheno \
  --validation-pheno-txt validation.pheno \
  --covar-txt covar.txt \
  --keep-path train.keep \
  --validation-keep-path validation.keep \
  --prediction-bed-prefix /path/to/prediction_data \
  --prediction-covar-txt prediction.covar \
  --prediction-keep-path prediction.keep \
  --compute-effects \
  --out-prefix out/sparse_single

The combined phenotype and keep files are constructed automatically from the training and validation inputs. They may instead be supplied explicitly with --fit-pheno-txt and --fit-keep-path. A fixed multi-GRM run only adds an exhaustive component spec:

gpu-reml-sparse \
  --mode fixed \
  --bed-prefix /path/to/data \
  --component-spec components_ld4.npz \
  --pheno-txt train.pheno \
  --validation-pheno-txt validation.pheno \
  --covar-txt covar.txt \
  --keep-path train.keep \
  --validation-keep-path validation.keep \
  --out-prefix out/sparse_ld4

Adaptive K with an existing LD-score table:

gpu-reml-sparse \
  --mode adaptive \
  --bed-prefix /path/to/data \
  --ld-score metadata/ldscore.tsv \
  --pheno-txt train.pheno \
  --validation-pheno-txt validation.pheno \
  --covar-txt covar.txt \
  --keep-path train.keep \
  --validation-keep-path validation.keep \
  --out-prefix out/sparse_adaptive

Omit --ld-score to calculate and cache it automatically. Use --plink2 when the executable is not on PATH, and --ld-window-kb to change the default 1-Mb physical window.

Continuous-trait marginal GWAS:

gpu-reml-gwas \
  --bed-prefix /path/to/data \
  --pheno-txt pheno.txt \
  --out-prefix out/gwas

Add --covar-txt covar.txt when covariates should be included.

The repository-local run_gpu.sh launcher remains available for environment-heavy benchmark runs.

Component Specifications

Component specs define how SNPs are grouped into GRM components. A JSON spec can name components and carry metadata:

{
  "components": [
    {
      "name": "maf_0_01",
      "variant_indices": [0, 4, 9],
      "annotation": {"maf_bin": "0-1%"}
    },
    {
      "name": "maf_01_05",
      "variant_indices": [1, 2, 8],
      "annotation": {"maf_bin": "1-5%"}
    }
  ]
}

NPZ specs are also supported for compact programmatic construction. See docs/component_specs.md.

Python API

import jax.numpy as jnp
from GPU_REML import FitConfig, InfinitesimalREMLFitter

cfg = FitConfig(
    bed_prefix="/path/to/data",
    n_rand_vec=100,
    minq_iter=10,
    slq_samples=4,
    slq_m=8,
    precond_rank=500,
    verbose=True,
)

fitter = InfinitesimalREMLFitter(cfg)
result = fitter.fit_infinitesimal(
    y=jnp.asarray(y),
    covar=jnp.asarray(covar),
)
print(result.var_components)

Lower-level users can call fit_reml with custom K @ V operators and diagonal atoms. This makes it possible to prototype new covariance representations without rewriting the REML optimizer. The supplied fixed-effect matrix is used exactly as given, so low-level callers should include an intercept when their model requires one and remove linearly dependent columns.

How It Works

At each REML step, GPU_REML needs repeated applications of:

$$ H(\theta)V = \theta_e V + \sum_g \theta_g K_g V $$

The implementation builds this product from streamed genotype blocks. REML evaluation then combines:

• block PCG solves for H^-1 [X | y | random probes];
• Hutchinson probes for trace terms in the score;
• stochastic Lanczos quadrature for log|H|;
• one-pass affine Lanczos reuse for the single-GRM identity-residual model;
• projected Fisher / AI-style variance-component updates with nonnegative genetic-variance constraints;
• a projected-core preconditioner dI + U C(theta) U.T that captures leading covariance structure. Residual SLQ keeps a fit-wide fixed reference, while PCG independently rebuilds its basis after accepted nonterminal updates.

For SMILE-style weighted kernels, the same REML loop is reused after replacing the standard GRM operator by the block-diagonal weighted operator. This keeps the new covariance representation isolated from the ordinary single-GRM, multi-GRM, partitioned, and sparse paths.

For routine runs, the most important user-facing resource controls are the GPU budget and the genotype-streaming ring depth.

See the mathematical overview for the score, AI, trace, SLQ, preconditioner, and heritability formulas, and the architecture guide for module boundaries and the fit lifecycle.

Key Runtime Parameters

• --gpu-budget-gib: planner-side budget for active GPU allocations. The planner uses it to choose the streamed SNP call width and the size of GPU-resident work arrays, including REML random-probe blocks and projected-core state. If omitted or set to 0, GPU_REML uses 85% of the currently available GPU memory estimate. This is not a hard nvidia-smi process-memory cap: JAX's pooled allocator can retain inactive blocks for reuse, and CUDA also owns context/workspace memory. GPU_REML reports the JAX peak_active value and, when the backend exposes it, peak_reserved at the end of every CLI run so users can distinguish real live use from allocator retention. Lower this budget to leave more VRAM for other processes. For memory diagnostics, XLA_PYTHON_CLIENT_ALLOCATOR=platform releases allocations eagerly, at a potential performance cost.
• --ring-depth: number of CPU-side staging buffers used for genotype streaming. This is the main knob for controlling CPU memory peak during data movement. Larger values can give smoother host-to-GPU streaming but allocate more pinned/staging memory on the CPU. The default 0 lets the planner choose a conservative value.
• GPU_REML_MATMUL_PRECISION: JAX matmul precision policy. The robust default is highest; an alternative should be used only after validating numerical agreement on the target GPU.
• --pcg-tol: ordinary PCG tolerance used by screening, candidate construction, and outer iterations; the default is 5e-3.
• --kkt-tol and --kkt-rel-tol: absolute and lambda-scaled tolerances for the signed score-KKT conditions. Their effective minimum is max(1e-4, 2*pcg_tol), so the certificate does not demand more precision than its PCG inputs provide.
• --outer-max: maximum number of variance updates; the default is 20.
• --sparsity-validation-pheno-txt and --sparsity-validation-out: required together for model selection. The complete lambda path is evaluated on the validation samples inside every alpha/theta outer iteration, and lambda is selected by maximum predictive R² (1 - SSE/SST).
• --effect-rel-tol: relative tolerance for change in the complete fitted fixed mean; the default is 5e-2.

The startup report labels the pinned streaming ring as Host memory plan (CPU RAM, not GPU VRAM). A line such as streaming_ring=35.4GiB therefore describes host RAM. At shutdown, peak_active is the value to compare with the planner's active estimate; peak_reserved is closer to what nvidia-smi observes.

Validation

Build a wheel:

python -m pip wheel --no-deps --no-build-isolation \
  --wheel-dir /tmp/gpu_reml_wheel /path/to/GPU_REML

Limitations

• REML likelihood terms use randomized approximations; results can vary with seed and SLQ/Hutchinson settings.
• A PCG residual tolerance does not itself give an analytic coordinatewise bound on KKT-score error. Reported KKT status is a finite-tolerance numerical certificate, not a proof of exact optimality.
• The package currently focuses on continuous traits.
• GPU performance depends on JAX/CUDA versions, PCIe bandwidth, call width, sample size, SNP count, and component count.
• This is research software. Validate settings against small exact references or matched external software before using it for production scientific conclusions.
• No public license has been selected yet. Do not redistribute until the project owner adds an explicit license.