Performance and Scalability

Overview

This vignette benchmarks the forestBalance pipeline to characterize how speed and memory scale with sample size (\(n\)) and number of trees (\(B\)).

The pipeline has four stages, each optimized:

Forest fitting (grf::multi_regression_forest) – C++ via grf.
Leaf node extraction (get_leaf_node_matrix) – C++ via Rcpp.
Indicator matrix construction (leaf_node_kernel_Z) – C++ via Rcpp, building CSC sparse format directly without sorting.
Weight computation (kernel_balance) – adaptive solver:
- Direct: sparse Cholesky on treated/control sub-blocks. Preferred for smaller problems.
- CG: conjugate gradient using the factored \(Z\) representation. No kernel matrix is formed. Preferred for large \(n\) or dense kernels.

library(grf)
library(Matrix)

Mathematical background

The kernel energy balancing system

Given a kernel matrix \(K \in \mathbb{R}^{n \times n}\) and a binary treatment vector \(A \in \{0,1\}^n\) with \(n_1\) treated and \(n_0\) control units, the kernel energy balancing weights \(w\) are obtained by solving the linear system \[ K_q \, w = z, \] where \(K_q\) is the modified kernel and \(z\) is a right-hand side vector determined by a linear constraint.

The modified kernel is defined element-wise as \[ K_q(i,j) = \frac{A_i \, A_j \, K(i,j)}{n_1^2} + \frac{(1-A_i)(1-A_j) \, K(i,j)}{n_0^2}. \]

A critical structural observation is that \(K_q(i,j) = 0\) whenever \(A_i \neq A_j\): the treated–control cross-blocks are identically zero. Therefore \(K_q\) is block-diagonal: \[ K_q = \begin{pmatrix} K_{tt} / n_1^2 & 0 \\ 0 & K_{cc} / n_0^2 \end{pmatrix}, \] where \(K_{tt} = K[A{=}1,\; A{=}1]\) is the \(n_1 \times n_1\) treated sub-block and \(K_{cc} = K[A{=}0,\; A{=}0]\) is the \(n_0 \times n_0\) control sub-block.

The right-hand side vector \(b\) has a similarly separable structure. Writing \(\mathbf{r} = K \mathbf{1}\) (the row sums of \(K\)), we have \[ b_i = \frac{A_i \, r_i}{n_1 \, n} + \frac{(1-A_i) \, r_i}{n_0 \, n}. \]

The full system decomposes into two independent sub-problems:

Treated block: solve \(K_{tt} \, w_t = n_1^2 \, z_t\) of size \(n_1\),
Control block: solve \(K_{cc} \, w_c = n_0^2 \, z_c\) of size \(n_0\),

where \(z_t\) and \(z_c\) are the constraint-adjusted right-hand sides for each group (each requiring two preliminary solves to determine the Lagrange multiplier).

The kernel factorization

The proximity kernel is \(K = Z Z^\top / B\), where \(Z\) is a sparse \(n \times L\) indicator matrix (\(L = \sum_{b=1}^B L_b\), total leaves across all trees). Each row of \(Z\) has exactly \(B\) nonzero entries (one per tree), so \(Z\) has \(nB\) nonzeros total. The \(Z\) matrix is constructed in C++ directly in compressed sparse column (CSC) format, avoiding the overhead of R’s triplet sorting.

Direct solver (block Cholesky)

For moderate \(n\), we form the sub-block kernels explicitly: \[ K_{tt} = Z_t Z_t^\top / B, \qquad K_{cc} = Z_c Z_c^\top / B, \] where \(Z_t = Z[A{=}1, \,\cdot\,]\) and \(Z_c = Z[A{=}0, \,\cdot\,]\). Each sub-block is computed via a sparse tcrossprod. The linear systems are then solved by sparse Cholesky factorization.

Computational Cost: \(O(nB \cdot \bar{s})\) for the two sub-block cross-products (where \(\bar{s}\) is the average leaf size), plus \(O(n_1^{3/2} + n_0^{3/2})\) for sparse Cholesky (in the best case).

CG solver (matrix-free)

For large \(n\), forming \(K_{tt}\) and \(K_{cc}\) becomes expensive. The conjugate gradient (CG) solver avoids forming any kernel matrix by operating on the factored representation. To solve \[ K_{tt} \, x = r \quad \Longleftrightarrow \quad Z_t Z_t^\top x = B \, r, \] CG only needs matrix–vector products of the form \[ v \;\mapsto\; Z_t \bigl(Z_t^\top v\bigr), \] each costing \(O(n_1 B)\) via two sparse matrix–vector multiplies. The same applies to the control block with \(Z_c\).

Here, the memory use is \(O(nB)\) for \(Z\) alone, versus \(O(n^2)\) for the kernel. Each CG iteration costs \(O(n_g B)\) (where \(n_g\) is the group size). Convergence requires a few hundred iterations in the examples below, so the total cost is \(O(n_g B \cdot T_{\text{iter}})\). Only four solves are needed (two per block), since the third right-hand side is a linear combination of the first two.

Computational Cost: \(O\bigl((n_1 + n_0) \cdot B \cdot T_{\text{iter}}\bigr)\) where \(T_{\text{iter}}\) is the number of iterations. For a fixed number of iterations this scales linearly in both \(n\) and \(B\).

Block Jacobi preconditioned CG

Plain CG can require many iterations at large \(n\). The Block Jacobi solver uses the first tree’s leaf partition to build a block-diagonal preconditioner. Each leaf block is a small dense system (~min.node.size \(\times\) min.node.size) that is cheap to factor. Preconditioned CG then needs fewer iterations than plain CG (about 30% fewer at \(n = 10{,}000\) below), but each iteration also applies the preconditioner, so it is not always faster overall.

Note: only 2 linear solves per treatment-group block are needed (not 3), because the third right-hand side is a linear combination of the first two.

Solver comparison

We compare all three solvers at \(n = 10{,}000\) and \(n = 25{,}000\):

# Fit one forest per n, then time each solver on the SAME kernel system.
solver_bench <- do.call(rbind, lapply(c(10000, 25000), function(nn) {
  set.seed(123)
  dat <- simulate_data(n = nn, p = 10)
  B_val <- 1000
  mns <- max(20L, min(floor(nn / 200) + 10, floor(nn / 50)))

  forest <- grf::multi_regression_forest(
    dat$X, scale(cbind(dat$A, dat$Y)),
    num.trees = B_val, min.node.size = mns
  )
  leaf_mat <- get_leaf_node_matrix(forest, dat$X)
  Z <- leaf_node_kernel_Z(leaf_mat)

  solvers <- if (nn <= 10000) c("bj", "cg", "direct") else c("bj", "cg")
  do.call(rbind, lapply(solvers, function(s) {
    if (s == "direct") {
      K <- leaf_node_kernel(leaf_mat)
      t <- system.time(
        w <- kernel_balance(dat$A, kern = K, solver = "direct")
      )["elapsed"]
    } else {
      t <- system.time(
        w <- kernel_balance(dat$A, Z = Z, leaf_matrix = leaf_mat,
                             num.trees = B_val, solver = s)
      )["elapsed"]
    }
    ate <- weighted.mean(dat$Y[dat$A == 1], w$weights[dat$A == 1]) -
           weighted.mean(dat$Y[dat$A == 0], w$weights[dat$A == 0])
    data.frame(n = nn, solver = s, time = t, ate = ate)
  }))
}))

Solver comparison on the same kernel (B = 1,000). ATE agreement confirms solvers find the same solution.
n	Solver	Time (s)	ATE
10000	bj	8.97	0.016298
10000	cg	6.39	0.016298
10000	direct	23.06	0.016298
25000	bj	45.65	0.006126
25000	cg	32.81	0.006126

All solvers produce the same ATE to high precision, confirming they solve the same system. CG (Rcpp) is the robust default; Block Jacobi can be faster in specific regimes (moderate leaf size, many trees) and is available via solver = "bj".

End-to-end timing

We benchmark the full forest_balance() pipeline (forest fitting, leaf extraction, Z construction, and weight computation) across a range of sample sizes up to \(n = 50{,}000\) with \(B = 1{,}000\) trees:

n_vals <- c(500, 1000, 2500, 5000, 10000, 25000, 50000)
B <- 1000
p <- 10

bench <- do.call(rbind, lapply(n_vals, function(nn) {
  set.seed(123)
  dat <- simulate_data(n = nn, p = p)

  # Auto (default)
  t_auto <- system.time(
    fit_auto <- forest_balance(dat$X, dat$A, dat$Y, num.trees = B)
  )["elapsed"]

  data.frame(n = nn, trees = B, auto = t_auto,
             auto_solver = fit_auto$solver)
}))

Full pipeline time (auto solver).
n	Trees	Time (s)	Solver
500	1000	0.16	cg
1000	1000	0.29	direct
2500	1000	1.03	direct
5000	1000	4.26	direct
10000	1000	16.18	direct
25000	1000	27.79	cg
50000	1000	88.21	cg

plot of chunk timing-plot

The solver is selected automatically based on the per-fold sample size. The Block Jacobi preconditioned CG (green) activates for large \(n\) and provides the best scaling.

Scaling with number of trees

tree_vals <- c(200, 500, 1000, 2000)
n_test <- c(1000, 5000, 25000)

tree_bench <- do.call(rbind, lapply(n_test, function(nn) {
  do.call(rbind, lapply(tree_vals, function(B) {
    set.seed(123)
    dat <- simulate_data(n = nn, p = 10)
    t <- system.time(
      fit <- forest_balance(dat$X, dat$A, dat$Y, num.trees = B)
    )["elapsed"]
    data.frame(n = nn, trees = B, time = t)
  }))
}))

Pipeline time across tree counts.
n	Trees	Time (s)
1000	200	0.08
1000	500	0.15
1000	1000	0.25
1000	2000	0.57
5000	200	1.19
5000	500	2.03
5000	1000	3.97
5000	2000	8.06
25000	200	5.33
25000	500	13.80
25000	1000	28.14
25000	2000	55.97

plot of chunk tree-plot

The pipeline scales approximately linearly in both \(n\) and \(B\).

Pipeline stage breakdown

The following experiment profiles each stage of the pipeline individually to show where time is spent at different scales:

breakdown <- do.call(rbind, lapply(
  list(c(1000, 500), c(5000, 500), c(5000, 2000),
       c(25000, 500), c(25000, 1000)),
  function(cfg) {
    nn <- cfg[1]; B_val <- cfg[2]
    set.seed(123)
    dat <- simulate_data(n = nn, p = 10)
    mns <- max(20L, min(floor(nn / 200) + 10, floor(nn / 50)))

    # 1. Forest fitting
    t_fit <- system.time({
      resp <- scale(cbind(dat$A, dat$Y))
      forest <- grf::multi_regression_forest(dat$X, Y = resp,
                        num.trees = B_val, min.node.size = mns)
    })["elapsed"]

    # 2. Leaf extraction (Rcpp)
    t_leaf <- system.time(
      leaf_mat <- get_leaf_node_matrix(forest, dat$X)
    )["elapsed"]

    # 3. Z construction (Rcpp)
    t_z <- system.time(
      Z <- leaf_node_kernel_Z(leaf_mat)
    )["elapsed"]

    # 4. Weight computation (kernel_balance)
    t_bal <- system.time(
      w <- kernel_balance(dat$A, Z = Z, num.trees = B_val, solver = "cg")
    )["elapsed"]

    data.frame(n = nn, trees = B_val, mns = mns,
               fit = t_fit, leaf = t_leaf, Z = t_z, balance = t_bal,
               total = t_fit + t_leaf + t_z + t_bal)
  }
))

Wall-clock time for each pipeline stage across configurations (CG solver, single fold).
n	Trees	mns	Forest fit (s)	Leaf extract (s)	Z build (s)	Balance (s)	Total (s)
1000	500	20	0.05	0.02	0.007	0.11	0.19
5000	500	35	0.27	0.13	0.027	1.04	1.47
5000	2000	35	1.00	0.50	0.098	4.47	6.07
25000	500	135	1.40	0.62	0.120	16.04	18.19
25000	1000	135	2.76	1.27	0.243	32.40	36.68

plot of chunk breakdown-plot

At small \(n\), forest fitting and the balance solver take similar time. At large \(n\), the CG solver dominates — its cost grows with the number of iterations and the sparse matrix–vector product size. The C++ stages (leaf extraction and \(Z\) construction) remain a small fraction throughout.

Memory usage

When the direct solver is used, the kernel is formed as a sparse matrix. For the CG solver, only the sparse indicator matrix \(Z\) is stored. The table below shows the actual memory usage compared to the theoretical dense kernel:

mem_data <- do.call(rbind, lapply(
  c(500, 1000, 2500, 5000, 10000, 25000, 50000),
  function(nn) {
    set.seed(123)
    dat <- simulate_data(n = nn, p = 10)
    B_val <- 1000
    mns <- max(20L, min(floor(nn / 200) + ncol(dat$X), floor(nn / 50)))

    fit_forest <- multi_regression_forest(
      dat$X, scale(cbind(dat$A, dat$Y)),
      num.trees = B_val, min.node.size = mns
    )
    leaf_mat <- get_leaf_node_matrix(fit_forest, dat$X)

    Z <- leaf_node_kernel_Z(leaf_mat)
    z_mb <- as.numeric(object.size(Z)) / 1e6

    # Kernel sparsity (skip forming K for very large n)
    if (nn <= 10000) {
      K <- leaf_node_kernel(leaf_mat)
      pct_nz <- round(100 * length(K@x) / as.numeric(nn)^2, 1)
      k_mb <- round(as.numeric(object.size(K)) / 1e6, 1)
    } else {
      pct_nz <- NA
      k_mb <- NA
    }

    n_fold <- nn %/% 2
    auto_solver <- if (n_fold > 5000 || mns > n_fold / 20) "cg" else "direct"
    dense_mb <- round(8 * as.numeric(nn)^2 / 1e6, 0)

    data.frame(n = nn, mns = mns, pct_nz = pct_nz,
               sparse_MB = k_mb, Z_MB = round(z_mb, 1),
               dense_MB = dense_mb, solver = auto_solver)
  }
))

Memory usage with adaptive leaf size (p = 10).
n	mns	Solver	K % nonzero	Stored	Actual (MB)	Dense K (MB)	Actual / Dense
500	20	cg	39.4%	Z (factor)	6.0	2	300%
1000	20	direct	28.8%	K (sparse)	3.5	8	43.8%
2500	22	direct	20.1%	K (sparse)	15.1	50	30.2%
5000	35	direct	16%	K (sparse)	48.1	200	24%
10000	60	direct	12.6%	K (sparse)	151.7	800	19%
25000	135	cg	–	Z (factor)	300.4	5000	6%
50000	260	cg	–	Z (factor)	600.3	20000	3%

plot of chunk memory-plot

At \(n = 50{,}000\), a dense kernel would require 19 GB. The CG solver stores only the \(Z\) matrix, using a small fraction of this.

Summary

The full forest_balance() pipeline scales to \(n = 50{,}000\) with 1,000 trees.
All computationally intensive steps are in compiled code: C++ for leaf extraction and \(Z\) construction (Rcpp), BLAS for sparse matrix operations, and CHOLMOD for sparse Cholesky.
The adaptive min.node.size heuristic adjusts leaf size to both \(n\) and \(p\), creating denser kernels than the old default of 10.
The solver switchover between direct and CG depends on both the per-fold sample size and the kernel density. The "auto" setting handles this transparently.
The pipeline scales approximately linearly in both \(n\) and \(B\).