DirectRS

Overview

DirectRS is a post-processing method that enhances pre-trained tree ensemble models (XGBoost, LightGBM, CatBoost) by extracting a global stretch matrix S’ from tree geometry, then fitting per-tree Ridge regressions on stretched embeddings. The result is improved accuracy with exact additive interpretability — every prediction can be decomposed as:

\[f(\mathbf{x}) = c_0 + \sum_{j=1}^{p} c_j(\mathbf{x})\]

where each \(c_j\) is the contribution of feature \(j\), and the decomposition holds to machine precision.

Model Architecture

DirectRS has three components:

  1. Tree Geometry → S’ Extraction: The split structure of all trees defines a global Gram matrix G. Its square root S’ is a stretch matrix that captures how the ensemble perceives feature space.

  2. Per-Tree Ridge Embeddings: For each tree \(t\), a Ridge regression is fit on the embedding \(\varphi_t(\mathbf{x}) = [1,\; \mathbf{S'x}]\), mapping stretched features to the tree’s leaf predictions.

  3. Prediction Aggregation: The final prediction averages per-tree Ridge predictions:

    \[f(\mathbf{x}) = \frac{1}{T} \sum_{t=1}^{T} \boldsymbol{\beta}_t^\top \varphi_t(\mathbf{x})\]

Implementation in MoDeVa

from modeva import DataSet, TestSuite
from modeva.models import MoXGBRegressor, MoDirectRSRegressor

# Load data
ds = DataSet()
ds.load(name="CaliforniaHousing")
ds.set_random_split()

# Train base model
model = MoXGBRegressor(
    name="XGB-depth2",
    n_estimators=200, max_depth=2, learning_rate=0.1,
    random_state=42, verbosity=0
)
model.fit(ds.train_x, ds.train_y.ravel())

# Fit DirectRS
drs = MoDirectRSRegressor(base_model=model, ridge_alpha=100.0)
drs.fit(
    ds.train_x, ds.train_y.ravel(),
    X_val=ds.test_x, y_val=ds.test_y.ravel()
)

Interpretability Features

DirectRS provides six interpretability methods:

S’ Stretch Analysis

Analyze the global stretch matrix to understand how the ensemble perceives feature space. The eigenvalue spectrum reveals effective dimensionality.

result = drs.get_global_stretch_analysis(ds.feature_names)
result.plot("eigenvalue_spectrum")
result.plot("feature_activity")

Local Explanation

Exact additive decomposition of each prediction into per-feature contributions. The decomposition holds to machine precision (~1e-14 error).

result = drs.explain_local(ds.test_x, feature_names=ds.feature_names)
result.plot()

Global Feature Importance

Aggregate feature importance across the model. Four modes are available:

  • "slope" — absolute Ridge coefficients (default)

  • "slope2" — squared Ridge coefficients

  • "contrib_abs" — mean absolute contribution

  • "contrib_rms" — root mean squared contribution

result = drs.importance_global(feature_names=ds.feature_names, mode="slope")
result.plot()

Main/Interaction Decomposition

Decompose total model variance into main effects and interaction effects using eta-squared (\(\eta^2\)) diagnostics. Main effects are computed from binned predictions \(f(\mathbf{x})\), not from individual coefficients.

result = drs.importance_main_interaction(ds.test_x, feature_names=ds.feature_names)
result.plot()

Geometric Interaction Tracing

Trace feature interactions through the adjacency matrix A derived from \(\log(\mathbf{S'})\). Higher-order walks \(\mathbf{W}^{(k)} = \mathbf{A}^k\) capture k-step interaction paths. The interaction energy \(E_k\) quantifies the strength at each order.

result = drs.geometric_interaction_traces(
    feature_names=ds.feature_names, K=4, gamma=0.5
)
result.plot("adjacency")
result.plot("spectrum")

Top Feature Interactions

Rank pairwise feature interactions by their off-diagonal entries in the stretch matrix.

result = drs.get_off_diagonal_analysis(ds.feature_names, top_k=15)
result.plot()

FANOVA Comparison

For depth-2 tree ensembles, DirectRS and FANOVA capture the same pairwise interactions. This makes depth-2 models ideal for validating DirectRS importance against the well-established FANOVA decomposition:

ts = TestSuite(ds, model)

# FANOVA feature importance
result_fanova = ts.interpret_fi()
result_fanova.plot()

# DirectRS feature importance
result_drs = drs.importance_global(feature_names=ds.feature_names, mode="slope")
result_drs.plot()

Key Classes

See the API Reference for full method documentation.