使用LightGBM.LGBMRanker交叉验证时如何保持组完整性？

问题背景 #

处理搜索排序任务，数据集包含查询（query）、URL及对应相关性得分（目标值），需用lightgbm.LGBMRanker以NDCG为目标做交叉验证。已知需将同一query的URL归为同一group，原数据集所有group大小均≥2，但执行时出现错误：

ValueError: Computing NDCG is only meaningful when there is more than 1 document. Got 1 instead.

调试发现自定义_feval函数中出现size=1的group，说明lightgbm.cv未保持组的完整性。原代码如下：

def eval_model(
    self,
    model: lightgbm.LGBMRanker,
    k_fold: int = 3,
    seed: int = 42,
):
    """Evaluates with NDCG"""

    def _feval(y_pred: np.ndarray, lgb_dataset: lightgbm.basic.Dataset):
        y_true = lgb_dataset.get_label()
        serp_sizes = lgb_dataset.get_group()

        ndcg_values = []
        start = 0
        for size in serp_sizes:
            end = start + size
            y_true_serp, y_pred_serp = y_true[start:end], y_pred[start:end]
            ndcg_serp = sklearn.metrics.ndcg_score(
                [y_true_serp], [y_pred_serp], k=10
            )
            ndcg_values.append(ndcg_serp)
            start = end

        eval_name = "my-ndcg"
        eval_result = np.mean(ndcg_values)
        greater_is_better = True
        return eval_name, eval_result, greater_is_better

    lgb_dataset = lightgbm.Dataset(data=self.X, label=self.y, group=self.serp_sizes)
    cv_results = lightgbm.cv(
        params={**model.get_params(), "group": self.serp_sizes},
        train_set=lgb_dataset,
        num_boost_round=1_000,
        nfold=k_fold,
        stratified=False,
        seed=seed,
        feval=_feval,
    )
    ndcg = np.mean(cv_results["my-ndcg"])

    return ndcg

错误原因分析 #

1. lightgbm.cv默认拆分不考虑group完整性：该函数默认按行随机拆分数据集，不会识别group参数对应的分组逻辑，导致单个query的样本被拆分到训练集和验证集，验证集剩下单个样本的group，触发NDCG计算错误。
2. 参数传递误解：即使将group放入params字典，也仅用于模型训练阶段的group感知排序，不会改变交叉验证的数据集拆分规则。

解决方法：使用Group-aware交叉验证 #

最简便的方式是用sklearn.model_selection.GroupKFold实现按group拆分的交叉验证，确保每个fold中的query（group）完整不拆分。修改后的代码如下：

def eval_model(
    self,
    model: lightgbm.LGBMRanker,
    k_fold: int = 3,
    seed: int = 42,
):
    """Evaluates with NDCG using group-aware cross-validation"""
    from sklearn.model_selection import GroupKFold
    import sklearn.metrics

    # 假设你有每个样本对应的query ID数组（用于分组），比如self.query_ids
    gkf = GroupKFold(n_splits=k_fold)
    ndcg_values = []

    for train_idx, val_idx in gkf.split(self.X, self.y, groups=self.query_ids):
        # 拆分训练/验证集
        X_train, X_val = self.X.iloc[train_idx], self.X.iloc[val_idx]
        y_train, y_val = self.y.iloc[train_idx], self.y.iloc[val_idx]
        
        # 重新生成训练/验证集的group数组（统计每个query的样本数）
        train_groups = self.query_ids.iloc[train_idx].value_counts().sort_index().values
        val_groups = self.query_ids.iloc[val_idx].value_counts().sort_index().values

        # 初始化并训练Ranker模型
        ranker = lightgbm.LGBMRanker(**model.get_params(), random_state=seed)
        ranker.fit(
            X_train, y_train,
            group=train_groups,
            eval_set=[(X_val, y_val)],
            eval_group=[val_groups],
            eval_metric="ndcg@10",
            early_stopping_rounds=50,
            verbose=False
        )

        # 计算当前fold的NDCG
        y_pred = ranker.predict(X_val, num_iteration=ranker.best_iteration_)
        fold_ndcg = []
        start = 0
        for size in val_groups:
            if size <= 1:
                start += size
                continue
            end = start + size
            y_true_serp = y_val.iloc[start:end].values
            y_pred_serp = y_pred[start:end]
            ndcg = sklearn.metrics.ndcg_score([y_true_serp], [y_pred_serp], k=10)
            fold_ndcg.append(ndcg)
            start = end
        
        ndcg_values.append(np.mean(fold_ndcg))

    return np.mean(ndcg_values)

关键说明 #

• GroupKFold拆分逻辑：基于query_ids分组，确保每个query的所有样本要么在训练集，要么在验证集，不会被拆分。
• 重新生成group数组：原serp_sizes是全量数据集的分组，拆分后需重新统计每个fold的query样本数，保证模型训练和评估时的group正确性。
• 内置NDCG指标：可直接用LightGBM内置的ndcg@10评估指标，无需自定义_feval，简化代码。

内容的提问来源于stack exchange，提问作者Durand