ICL-MoE

Overview

ICL-MoE (In-Context Learning Mixture of Experts) is a post-processing layer that builds on a fitted DirectRS model. It replaces hard tree-based leaf routing with soft kNN-based attention gating in the stretched embedding space, enabling local adaptation of per-leaf linear experts to each query point.

Given a fitted DirectRS model with stretch matrix \(\mathbf{S'}\) and per-leaf Ridge coefficients, ICL-MoE computes predictions in two stages:

Leaf Expert Predictions: Each tree’s leaf-local linear model produces a prediction via the same tree routing as DirectRS.
Residual Correction: A kNN-weighted local Ridge regression corrects the residuals of the leaf experts, improving fit in local neighbourhoods of the embedding space \(\mathbf{z} = \mathbf{S'x}\).

The combined prediction is:

\[f(\mathbf{x}) = \underbrace{b + \sum_{t=1}^{T} (\alpha_{t,l(t)} + \boldsymbol{\gamma}_{t,l(t)}^\top \mathbf{z})}_{\text{leaf experts (= DirectRS)}} \;+\; \underbrace{\boldsymbol{\theta}(\mathbf{x})^\top \tilde{\mathbf{z}}}_{\text{residual correction}}\]

where \(l(t)\) is the active leaf for tree \(t\), and \(\boldsymbol{\theta}(\mathbf{x})\) is a locally weighted Ridge solution computed from the \(k\) nearest training neighbours in \(\mathbf{z}\)-space.

Model Architecture

ICL-MoE has three components:

Stretched Embedding: Reuses the global stretch matrix \(\mathbf{S'}\) from the fitted DirectRS model. All distance computations and local regressions operate in the stretched space \(\mathbf{z} = \mathbf{S'x}\).
Leaf Expert Layer: Per-tree, per-leaf linear functions \(\alpha_{t,l} + \boldsymbol{\gamma}_{t,l}^\top \mathbf{z}\) extracted directly from the DirectRS Ridge coefficients. This layer exactly reproduces DirectRS predictions.
kNN Residual Correction: For each query \(\mathbf{x}\), the \(k\) nearest training points in \(\mathbf{z}\)-space are found via KDTree. Softmax attention weights \(\pi_i = \mathrm{softmax}(-d_i^2 / \tau)\) are computed, and a weighted Ridge regression on the residuals \(r_i = y_i - \hat{y}_{\mathrm{leaf}}(\mathbf{x}_i)\) provides the local correction:

\[\boldsymbol{\theta}(\mathbf{x}) = (\tilde{\mathbf{Z}}^\top \mathbf{W} \tilde{\mathbf{Z}} + \lambda \mathbf{I})^{-1} \tilde{\mathbf{Z}}^\top \mathbf{W} \mathbf{r}\]

where \(\mathbf{W} = \mathrm{diag}(\pi_1, \ldots, \pi_k)\) and \(\tilde{\mathbf{Z}} = [1,\; \mathbf{Z}_{\mathrm{neighbours}}]\).

Implementation in MoDeVa

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

# 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())

# Fit ICL-MoE on top of DirectRS
icl = MoDirectRSICLRegressor(
    directrs_model=drs, k=50, tau=1.0, ridge_lambda=1.0
)
icl.fit(ds.train_x, ds.train_y.ravel())

For classification, the same pipeline applies using the classifier variants. ICL-MoE operates in logit space internally; predict_proba applies sigmoid:

from modeva.models import MoXGBClassifier, MoDirectRSClassifier
from modeva.models import MoDirectRSICLClassifier

model = MoXGBClassifier(
    name="XGB-cls-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())

drs = MoDirectRSClassifier(base_model=model, ridge_alpha=100.0)
drs.fit(ds.train_x, ds.train_y.ravel())

icl = MoDirectRSICLClassifier(
    directrs_model=drs, k=50, tau=1.0, ridge_lambda=1.0
)
icl.fit(ds.train_x, ds.train_y.ravel())

proba = icl.predict_proba(ds.test_x)  # (N, 2) probabilities

Interpretability Features

ICL-MoE provides four interpretability methods. All decompositions are exact (hold to machine precision).

Local Explanation

Exact additive decomposition of each prediction into per-feature contributions. The decomposition combines leaf expert contributions and residual correction:

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

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

Global Feature Importance

Aggregate feature importance across the dataset. Two modes are available:

"contrib_abs" — mean absolute contribution (default)
"contrib_rms" — root mean squared contribution

result = icl.importance_global(
    ds.test_x, feature_names=ds.feature_names, mode="contrib_abs"
)
result.plot()

Neighbour Analysis

Inspect the kNN neighbourhood used for residual correction. Returns neighbour weights, targets, and a PCA scatter plot of neighbours in the stretched embedding space.

result = icl.get_neighbor_analysis(
    ds.test_x, sample_index=0, feature_names=ds.feature_names
)
result.plot("weight_bar")
result.plot("neighbor_scatter")

Leaf Gating Analysis

Examine the per-leaf routing for a specific tree. Shows which leaf is active for the query point and the leaf-level prediction contributions.

result = icl.get_leaf_gating_analysis(
    ds.test_x, sample_index=0, tree_index=0,
    feature_names=ds.feature_names
)
result.plot()

FANOVA Comparison

For depth-2 tree ensembles, ICL-MoE, DirectRS, and FANOVA feature importance can be compared side-by-side. All three methods are available through MoDeVa:

ts = TestSuite(ds, model)

# ICL-MoE importance
icl_result = icl.importance_global(
    ds.test_x, feature_names=ds.feature_names, mode="contrib_abs"
)

# DirectRS importance
drs_result = drs.importance_global(
    feature_names=ds.feature_names, mode="contrib_abs"
)

# FANOVA importance
fanova_result = ts.interpret_fi()
fanova_result.plot()

Key Classes

MoDirectRSICLRegressor — ICL-MoE for regression tasks
MoDirectRSICLClassifier — ICL-MoE for classification tasks

Parameters

Parameter	Default	Description
`directrs_model`	(required)	Fitted `MoDirectRSRegressor` or `MoDirectRSClassifier`
`name`	`None`	Model identifier string
`k`	`50`	Number of nearest neighbours for kNN gating
`tau`	`1.0`	Temperature for softmax attention weights; lower values sharpen attention
`ridge_lambda`	`1.0`	Regularization for the residual Ridge regression
`top_m`	`5`	Number of top leaf experts per tree used in the gating analysis

See the API Reference for full method documentation.