mcts-pipeline-search

MCTS Pipeline Search

Context

Grid/random search wastes budget exploring dead branches. MCTS (Monte Carlo Tree Search) uses UCB1 to balance exploitation (best so far) and exploration (uncertain branches). agy recommends MCTS as P2 for Top-1% scenarios: used by AIDE for ML pipeline exploration, AlphaGo for game trees.

The core insight: UCB1 prevents both "stuck on suboptimal" and "random walk", getting best of both.

Guidance

Basic Usage

from framework.src.mcts import MCTSSearch, grid_expansion

# 网格定义
grid = {
    "num_leaves": [15, 31, 63, 127],
    "learning_rate": [0.01, 0.05, 0.1],
    "n_estimators": [100, 500, 1000],
}

# 评估函数: 真实 CV 或 proxy
def evaluator(config):
    return train_and_score(config)  # 返回 0-1 score

# 扩展: 网格生成子配置
expansion = grid_expansion(grid)

# 搜索
search = MCTSSearch(evaluator, expansion, max_depth=3)
result = search.search(initial_config={"num_leaves": 31}, iterations=100)

print(f"Best: {result.best_node.config}")
print(f"Score: {result.best_score:.3f}")
print(f"Tree size: {result.tree_size}")

Custom Expansion (Beyond Grid)

def my_expansion(config):
    """对当前 config 生成邻居(参数 ±10%)"""
    neighbors = []
    for delta in [-0.1, 0.1]:
        new = dict(config)
        new["learning_rate"] = max(0.001, config["learning_rate"] * (1 + delta))
        neighbors.append(new)
    return neighbors

Pipeline Search DAG

# 探索不同模型 pipeline 组合
def pipeline_evaluator(config):
    # config 包含: model_type, features, ensemble_method
    if config["model_type"] == "lgbm":
        return train_lgbm(config)
    elif config["model_type"] == "xgb":
        return train_xgb(config)
    elif config["model_type"] == "ensemble":
        return train_ensemble(config)

def pipeline_expansion(config):
    """从当前 pipeline 探索下一个组件"""
    next_steps = []
    for model in ["lgbm", "xgb", "catboost", "ensemble"]:
        new = dict(config)
        new["model_type"] = model
        next_steps.append(new)
    return next_steps

search = MCTSSearch(pipeline_evaluator, pipeline_expansion, max_depth=5)
result = search.search({"features": "basic"}, iterations=200)