OrionMoE-Small-v1.0

Model Overview

A mixture-of-experts foundation model for tabular classification.

OrionMoE-Small predicts labels for a new dataset using labeled rows as context. The model learns across synthetic classification tasks during pretraining, then adapts to a downstream dataset through in-context learning without updating its weights.

The model combines distribution-aware column embeddings with row and feature attention and dataset-conditioned expert routing. A scikit-learn interface provides fit, predict, and predict_proba, with support for numerical and categorical features, missing values, and many-class inference.

Developed by	Lexsi-Labs
Model	OrionMoE-Small-v1.0
Task	Binary and multiclass tabular classification
Framework	PyTorch, with a scikit-learn-compatible interface
Checkpoint	orionmoe-small-classifier-v1.0.ckpt
Model repository	Lexsi/OrionMoE
Documentation	Lexsi-Labs/OrionMoE

Quick start

Install the inference package from PyPI:

python -m pip install orionmoe

The package requires Python 3.9-3.12 and PyTorch 2.4 or newer. Both CPU and CUDA inference are supported. For CUDA, use a PyTorch build compatible with your system.

import torch
from sklearn.datasets import load_breast_cancer
from sklearn.metrics import accuracy_score, roc_auc_score
from sklearn.model_selection import train_test_split

from orion_moe.sklearn import OrionMoEClassifier

X, y = load_breast_cancer(return_X_y=True, as_frame=True)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, stratify=y, random_state=42
)

classifier = OrionMoEClassifier(
    device="cuda" if torch.cuda.is_available() else "cpu",
    ensemble_contexts=8,
    random_state=42,
)
classifier.fit(X_train, y_train)

probabilities = classifier.predict_proba(X_test)
predictions = classifier.classes_[probabilities.argmax(axis=1)]

print("Accuracy:", accuracy_score(y_test, predictions))
print("ROC AUC:", roc_auc_score(y_test, probabilities[:, 1]))

Input data

• Pass a two-dimensional NumPy array or pandas DataFrame and a one-dimensional target containing class labels.
• Keep feature columns in the same order for fitting and prediction. DataFrame column names must also match.
• Preprocessing is fitted on the labeled context. Numerical and categorical columns and missing feature values are handled internally.
• For categorical features represented by numeric codes, specify their column names or indices using categorical_features.
• Many-class inference is selected automatically when the number of classes exceeds the checkpoint's native classification head.

Architecture

OrionMoE-Small processes a table as a grid of rows and features:

1. Distribution-aware column embeddings represent feature values using information from the labeled support rows.
2. Induced row attention followed by feature attention models relationships across examples and columns while separating support and query rows.
3. Dataset-conditioned mixture-of-experts layers route cell representations to experts using dataset context.
4. Feature registers and an in-context classification head combine the table representation with support labels to predict query labels.

Inference controls

Setting	Purpose
ensemble_contexts	Number of fitted context/preprocessing members; the wrapper defaults to 8.
query_chunk_size	Splits query rows into inference chunks. The example uses 128; the wrapper default is 0, meaning no explicit query chunk size.
max_context_size	Limits the labeled rows used as context; None uses all available rows.
use_kv_cache, persistent_kv_cache	Enable attention caching and reuse across prediction calls. Both default to False.
dtype	Defaults to "float32". CUDA mixed precision supports "bfloat16" and "float16" on compatible hardware.

Memory use depends on the number of context rows, query rows, features, classes, and ensemble members. Automatic OOM recovery can reduce query chunk sizes and, if needed, context or features. Reducing context or features can change predictions. For evaluations that must preserve those inputs, set oom_context_retry=False and oom_feature_retry=False.

Evaluation

TabArena Elo performance

Figure 1. Classification Elo performance on TabArena (higher is better).

OpenML-CC18 performance

Figure 2. Classification Elo performance on OpenML-CC18 (higher is better).

License

The OrionMoE source code and model weights are licensed under the Lexsi Labs Source Available License (LSAL) v1.2, OrionMoE Edition. This is a source-available license, not an OSI-approved open-source license.

• Personal noncommercial research and academic research and teaching are permitted subject to the license.
• Organizational use is subject to the written acknowledgement or permission requirements in Section 1A.
• Commercial use requires a separate commercial license.
• Production or user-facing deployment requires express written authorization.

See LICENSE for the complete terms. Copyright (c) 2025-2026 Lithasa Technologies Pvt. Ltd. (Lexsi Labs). For permissions and commercial licensing, contact support@lexsi.ai.