[python-package] LGBMRanker does not call custom objective function with group

[francoisWeber]

Description

I tried a custom ranking objective function for LGBMRanker but LGBMRanker does not call the custom loss function with the group argument, which is mandatory for learning to rank.

Reproducible example

import lightgbm as lgbm
import numpy as np

# fake data
N, p, gsize = 1000, 10, 5
X = np.random.randn(N, p)
y = np.random.randint(0, 4, N)
groups = [gsize] * (N//gsize)


# fake loss to trace inputs 
_args = None 
_kwargs = None
def loss(*args, **kwargs):
    global _args, _kwargs
    _args = args
    _kwargs = kwargs
    return 0, 0


lgbm_ranker = lgbm.LGBMRanker(n_estimators=10, learning_rate=0.1, objective=loss)
try:
    lgbm_ranker.fit(X, y, group=groups)
except Exception as e:
    print(e)
print(len(_args))
print(_kwargs)

_args nor _kwargs contain reference to the groups I used as ranking group size.

Environment info

LightGBM version or commit hash: 4.6.0

Command(s) you used to install LightGBM

poetry add lightgbm

[jameslamb]

Thanks for using LightGBM.

For the scikit-learn APIs, LightGBM determines how many arguments to pass to custom objective functions by inspecting the signatures of the functions.

LightGBM/python-package/lightgbm/sklearn.py

Line 221 in de4c88f

argc = len(signature(self.func).parameters)

LightGBM/python-package/lightgbm/sklearn.py

Lines 231 to 233 in de4c88f

	if argc == 4:
	group = _get_group_from_constructed_dataset(dataset)
	return self.func(labels, preds, weight, group) # type: ignore[call-arg]

Maybe something else should be done there to handle *args, **kwargs in signatures, but that'd require some design work.

For now, please follow the exact signature indicated in the docs:

LightGBM/python-package/lightgbm/sklearn.py

Lines 162 to 188 in de4c88f

	func : callable
	Expects a callable with following signatures:
	``func(y_true, y_pred)``,
	``func(y_true, y_pred, weight)``
	or ``func(y_true, y_pred, weight, group)``
	and returns (grad, hess):
	y_true : numpy 1-D array of shape = [n_samples]
	The target values.
	y_pred : numpy 1-D array of shape = [n_samples] or numpy 2-D array of shape = [n_samples, n_classes] (for multi-class task)
	The predicted values.
	Predicted values are returned before any transformation,
	e.g. they are raw margin instead of probability of positive class for binary task.
	weight : numpy 1-D array of shape = [n_samples]
	The weight of samples. Weights should be non-negative.
	group : numpy 1-D array
	Group/query data.
	Only used in the learning-to-rank task.
	sum(group) = n_samples.
	For example, if you have a 100-document dataset with ``group = [10, 20, 40, 10, 10, 10]``, that means that you have 6 groups,
	where the first 10 records are in the first group, records 11-30 are in the second group, records 31-70 are in the third group, etc.
	grad : numpy 1-D array of shape = [n_samples] or numpy 2-D array of shape [n_samples, n_classes] (for multi-class task)
	The value of the first order derivative (gradient) of the loss
	with respect to the elements of y_pred for each sample point.
	hess : numpy 1-D array of shape = [n_samples] or numpy 2-D array of shape = [n_samples, n_classes] (for multi-class task)
	The value of the second order derivative (Hessian) of the loss
	with respect to the elements of y_pred for each sample point.

And also note that the docs there say the gradient and hessians should be 1D arrays (not a tuple of ints as your example is using).

[francoisWeber]

OK thanks a lot, @jameslamb - the code is running now. However I must have done mistakes in it because the training's behavior is wrong.

I give up for now but let me share my attempts with others. Here is what I tried:

• use RAX's smoothed ranking loss to have more LTR losses available in LGBM (e.g: optimizing on recall)
• use JAX (which is used by RAX) to obtain gradients and hessians
• plug them into a custom loss

My WIP is:

# Reimplement NDCG loss for LightGBM through RAX API

import rax
from functools import partial
import jax
import numpy as np
from typing import Tuple
from enum import Enum

class RankingLoss(Enum):
    """Loss functions for ranking"""
    NDCG = rax.ndcg_metric
    MRR = rax.mrr_metric
    RECALL = rax.recall_metric
    PRECISION = rax.precision_metric
    

class RankingLossForLGBM:
    def __init__(self, loss: RankingLoss, topn: int = 10):
        self.topn = topn
        self._loss = partial(rax.approx_t12n(loss), topn=topn)
        self._grad = jax.grad(self._loss)
        self._hess = jax.hessian(self._loss)
        
    def _call_on_one_group(self, labels: np.ndarray, preds: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        return self._grad(preds, labels), np.diag(self._hess(preds, labels))

    def __call__(self, labels: np.ndarray, preds: np.ndarray, weights: np.ndarray = None, groups: list[int] | None = None) -> Tuple[np.ndarray, np.ndarray]:
        # Split labels into chunks of size grp for grp in groups
        groups = groups if groups is not None else [len(labels)]
        weights = weights or np.ones(len(groups))
        start = 0
        grads = []
        hessians = []
        for grp, w in zip(groups, weights):
            end = start + grp
            _labels = labels[start:end]
            _preds = preds[start:end]
            grad, hess = self._call_on_one_group(_preds, _labels)
            grads.append(grad * w)
            hessians.append(hess * w)
            start = end
        return np.hstack(grads), np.hstack(hessians)


# fake data
N, p, gsize = 1000, 10, 5
X = np.random.randn(N, p)
y = np.random.randint(0, 4, N)
groups = [gsize] * (N//gsize)

# training
loss = RankingLossForLGBM(RankingLoss.NDCG, topn=2)
lgbm_ranker = lgbm.LGBMRanker(n_estimators=10, learning_rate=0.1, objective=loss)
lgbm_ranker.fit(X, y, group=groups)

I give up for now but this could give ideas to others.