ICL-MoE

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:

1. Leaf Expert Predictions: Each tree’s leaf-local linear model produces a prediction via the same tree routing as DirectRS.
2. 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:

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:

1. 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}$.

2. 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.

3. 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:

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

• ~modeva.models.directrs.icl_moe_api.MoDirectRSICLRegressor — ICL-MoE for regression tasks
• ~modeva.models.directrs.icl_moe_api.MoDirectRSICLClassifier — ICL-MoE for classification tasks

Parameters

See the API Reference </modules/icl_moe> for full method documentation.