# Michael A. Alcorn (malcorn@redhat.com)
# A (slightly modified) implementation of LamdaRank as described in [1].
#   [1] https://papers.nips.cc/paper/2971-learning-to-rank-with-nonsmooth-cost-functions.pdf
#   [2] https://www.microsoft.com/en-us/research/wp-content/uploads/2016/02/MSR-TR-2010-82.pdf

import numpy as np
import torch
import torch.nn as nn

from scipy.stats import rankdata

INPUT_DIM = 50
N_DOCS = 20
N_REL = 5


def idcg(n_rel):
    # Assuming binary relevance.
    nums = np.ones(n_rel)
    denoms = np.log2(np.arange(n_rel) + 1 + 1)
    return (nums / denoms).sum()


# Model.
model = torch.nn.Sequential(
    nn.Linear(INPUT_DIM, 128),
    nn.ReLU(),
    nn.Linear(128, 64),
    nn.ReLU(),
    nn.Linear(64, 32),
    nn.ReLU(),
    nn.Linear(32, 1))

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# Document scores.
doc_features = np.random.randn(N_DOCS, INPUT_DIM)
relevant = np.zeros(N_DOCS)
relevant[:N_REL] = 1

docs = torch.from_numpy(np.array(doc_features, dtype = "float32"))
docs = docs.to(device)
doc_scores = model(docs)
doc_scores_np = doc_scores.cpu().detach().numpy().flatten()

# Document ranks.
doc_ranks = rankdata(-doc_scores_np)

# Compute lambdas.
N = 1 / idcg(N_REL)
lambs = np.zeros(N_DOCS, dtype = "float32")
for i in range(N_REL):
    rel_score = doc_scores_np[i]
    diffs = rel_score - doc_scores_np[N_REL:]
    exped = np.exp(diffs)
    # See equation (6) in [2].
    lamb = -1 / (1 + exped) * N * np.abs(1 / np.log2(1 + doc_ranks[i]) - 1 / np.log2(1 + doc_ranks[N_REL:]))
    # See section 6.1 in [1], but lambdas have opposite signs from [2].
    lambs[i] -= lamb.sum()
    lambs[N_REL:] += lamb

# Accumulate lambda scaled gradients.
grads = {}
for i in range(N_DOCS):
    model.zero_grad()
    doc_scores[i].backward(retain_graph = True)
    for param in model.parameters():
        # See section 4.1 in [2].
        if param not in grads:
            grads[param] = lambs[i] * param.grad
        else:
            grads[param] += lambs[i] * param.grad

# Update model weights.
lr = 0.00001
with torch.no_grad():
    for param in model.parameters():
        # See section 4.1 in [2].
        param += lr * grads[param].to(device)

@@ -1734,327 +1734,120 @@
 ts.%0A
-grads = %7B%7D%0Afor i in range(N_DOCS):%0A    model.zero_grad()%0A    doc_scores%5Bi%5D.backward(retain_graph = True)%0A    for param in model.parameters():%0A        # See section 4.1 in %5B2%5D.%0A        if param not in grads:%0A            grads%5Bparam%5D = lambs%5Bi%5D * param.grad%0A        else:%0A            grads%5Bparam%5D += lambs%5Bi%5D * param.grad
+model.zero_grad()%0Alambs = torch.from_numpy(lambs.reshape((len(lambs), 1))).to(device)%0Adoc_scores.backward(lambs)
 %0A%0A# 
@@ -1952,76 +1952,29 @@
     
-# See section 4.1 in %5B2%5D.%0A        param += lr * grads%5Bparam%5D.to(device)
+param += lr * param.grad
 %0A

Use lambdas in backward call.