maifs 0.1.1

Model-Agnostic Ising Feature Selection for PyTorch models

MAIFS: Model-Agnostic Ising Feature Selection

MAIFS is a model-agnostic feature selection wrapper for PyTorch models. It learns a binary input-feature mask and uses an Ising/QUBO solver to decide which features should stay active.

The method does not assume a specific estimator such as linear regression, logistic regression, multi-output regression, or Cox regression. Any PyTorch model can be used as long as:

• the feature axis can be masked;
• the model output is accepted by a scalar PyTorch loss function;
• the loss is differentiable with respect to the masked input.

How It Works

MAIFS alternates between two steps:

1. Train model weights with the current binary feature mask fixed.
2. Differentiate the loss with respect to the mask, build a second-order QUBO approximation, solve it with a named backend, and update the mask.

For high-dimensional inputs, use hessian_mode="diagonal" to avoid constructing a dense Hessian. For smaller problems, hessian_mode="full" keeps pairwise feature interactions.

Feature-Count Guardrails

MAIFS applies two safety bounds after each solver update:

• min_selected_features: the minimum number of features that must remain selected after update_mask().
• max_selected_features: the maximum number of features that may remain selected after update_mask().

By default, max_selected_features=feature_dim, so if a solver returns an all-one mask, MAIFS keeps all features. The default min_selected_features is max(1, ceil(0.2 * feature_dim)); this default is used as a fallback when a solver returns an all-zero mask, so MAIFS does not leave the model with no features at all. If the user explicitly sets max_selected_features, the default minimum is clipped so it does not exceed that maximum.

If the user explicitly passes min_selected_features, MAIFS treats it as a hard lower bound after every update. If the user explicitly passes max_selected_features, MAIFS treats it as a hard upper bound after every update.

Users can override the bounds directly:

selector = MAIFSSelector(
    model,
    feature_dim=n_features,
    min_selected_features=3,
    max_selected_features=20,
)

Alternating Optimization Loop

For production-style use, call fit() instead of manually coordinating weight training and mask updates. The loop trains model weights for a fixed number of epochs, updates the mask, records per-cycle history, and can stop early when the mask becomes stable.

history = selector.fit(
    loader,
    loss_fn,
    optimizer,
    max_cycles=10,
    weight_epochs_per_cycle=50,
    warm_start_epochs=100,
    hessian_mode="diagonal",
    mask_patience=2,
)

print(history[-1])
print(selector.selected_indices())

Each history entry includes cycle, train_loss, mask_updated, mask_changes, num_selected, selected_indices, and stop_reason. Advanced users can still call fit_weights() and update_mask() directly when they need full manual control.

Supported Solvers

The solver name is passed directly through solver=...:

• local_search: bundled greedy local QUBO search, no optional dependency.
• dwave_sa: D-Wave simulated annealing; requires dimod and dwave-neal.
• kaiwu_fast_sa: Kaiwu local fast simulated annealing; requires kaiwu.
• kaiwu_split_fast_sa: Kaiwu FastSA with precision splitting; requires kaiwu.
• kaiwu_cim: Kaiwu cloud CIM backend. See the dedicated Kaiwu CIM section below before using it, because it submits remote CIM tasks.

Missing packages raise MAIFSDependencyError with the install command and detected package versions. Incompatible solver APIs raise MAIFSDependencyConflictError with the installed dependency versions.

Installation

For users, install MAIFS as a normal Python package:

pip install maifs

This installs only the MAIFS package itself. It does not automatically install the packages listed in requirements.txt, and it does not force PyTorch, D-Wave, or Kaiwu into the user's environment.

Before running MAIFS, make sure the active Python environment already has the runtime packages required by the functionality you want to use:

• MAIFSSelector: requires numpy and torch.
• solver="local_search": uses only numpy.
• solver="dwave_sa": requires dimod and dwave-neal.
• Kaiwu solvers: require kaiwu.

requirements.txt, requirements-kaiwu.txt, and requirements-dev.txt are reference files for preparing a local environment from source. They are not installed automatically by pip install maifs.

If you are working from a source checkout, install the package itself in editable mode from the project root:

pip install -e . --no-deps

If the environment is missing runtime packages, install only the packages you actually need. For example:

pip install numpy torch
pip install dimod dwave-neal

Install Kaiwu only if you want to use kaiwu_fast_sa, kaiwu_split_fast_sa, or kaiwu_cim:

pip install kaiwu

For development and tests, install:

pip install -r requirements-dev.txt

Run the local tests from the project root:

python -m pytest tests

Quick Start

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from maifs import MAIFSSelector

torch.manual_seed(7)

n_samples = 128
n_features = 20
x = torch.randn(n_samples, n_features)
y = (2.0 * x[:, :1] - 1.5 * x[:, 3:4] + 0.1 * torch.randn(n_samples, 1))

loader = DataLoader(TensorDataset(x, y), batch_size=128, shuffle=False)
model = nn.Linear(n_features, 1)

selector = MAIFSSelector(
    model,
    feature_dim=n_features,
    cardinality_k=2,
    gamma_penalty=5.0,
    solver="local_search",
    solver_kwargs={"max_iter": 200},
)

loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(selector.model.parameters(), lr=0.05)

history = selector.fit(
    loader,
    loss_fn,
    optimizer,
    warm_start_epochs=100,
    max_cycles=5,
    weight_epochs_per_cycle=50,
    hessian_mode="diagonal",
)

print(selector.selected_indices())
print(history[-1])

Using Another Solver

The only change is the solver name and optional solver_kwargs:

selector = MAIFSSelector(
    model,
    feature_dim=n_features,
    cardinality_k=10,
    solver="kaiwu_fast_sa",
    solver_kwargs={
        "sa_num_reads": 64,
        "sa_num_sweeps": 1000,
        "sa_random_state": 0,
    },
)

Kaiwu CIM Backend

kaiwu_cim submits the QUBO/Ising problem to Kaiwu's cloud CIM service. It is therefore different from local_search, dwave_sa, and kaiwu_fast_sa, which can run locally. MAIFS does not run CIM tests by default because a CIM call can consume cloud quota, requires a valid license, and may create a remote task.

Before using solver="kaiwu_cim", make sure you have:

• installed Kaiwu with pip install kaiwu or, from a source checkout, pip install -r requirements-kaiwu.txt;
• initialized the Kaiwu license in the Python environment;
• obtained a valid project_no;
• chosen whether the call should block until completion with wait=True.

Initialize Kaiwu before constructing the selector:

import kaiwu.license as lic

lic.init(user_id, sdk_code)

Then pass CIM options through solver_kwargs. The cim_optimizer_kwargs dict is forwarded to kaiwu.cim.CIMOptimizer; common fields include project_no, task_mode, sample_number, sample_sort_mode, wait, and interval.

selector = MAIFSSelector(
    model,
    feature_dim=n_features,
    cardinality_k=10,
    solver="kaiwu_cim",
    solver_kwargs={
        "target_precision": 8,
        "max_bits": 1000,
        "max_precision": 32,
        "precision_step": 4,
        "cim_cleanup_records": True,
        "cim_optimizer_kwargs": {
            "project_no": "your-project-no",
            "task_mode": "quota",
            "sample_number": 100,
            "sample_sort_mode": 1,
            "wait": True,
            "interval": 10,
        },
    },
)

Use it exactly like the local solvers once the selector is configured:

selector.fit(
    loader,
    loss_fn,
    optimizer,
    max_cycles=5,
    weight_epochs_per_cycle=50,
    hessian_mode="diagonal",
)
print(selector.selected_indices())

For integration testing, CIM is opt-in. Set MAIFS_RUN_CIM=1 only when the license and project settings are ready:

$env:MAIFS_RUN_CIM="1"
python -m pytest tests\test_maifs_solvers.py::test_kaiwu_cim_backend_when_explicitly_enabled -q

If the license is not initialized, MAIFS raises an error explaining how to call lic.init(user_id, sdk_code). If the Kaiwu package version is incompatible, MAIFS reports the detected Kaiwu version and the missing interface.

Error Handling

from maifs import MAIFSDependencyError, MAIFSDependencyConflictError, MAIFSSolverError

try:
    selector.update_mask(loader, loss_fn, hessian_mode="diagonal")
except MAIFSDependencyError as exc:
    print(exc)
except MAIFSDependencyConflictError as exc:
    print(exc)
except MAIFSSolverError as exc:
    print(exc)

These exceptions include the failing solver name, missing package or conflict information, and suggested installation or configuration steps.