排序算法常用损失函数

ListMLE

$S=left{s_{1}, s_{2}, s_{3}, ldots, s_{n}right}$ 表示模型对n个doc的得分，将预测得分按目标顺序排列，

$hat{S}=left{s_{pi_{1}}, s_{pi_{2}}, s_{pi_{3}}, ldots, s_{pi_{n}}right}$

$P(hat{S})=prod_{i=1}^{n} frac{exp left(s_{pi_{i}}right)}{sum_{j=i}^{n} exp left(s_{pi_{j}}right)}$

$L=-log P(hat{S})$

def list_mle(logits, labels, topn):
    """
    将预测得分按真实顺序排列，根据PL模型构造排序的概率
    """
    _, indices = tf.math.top_k(labels, k=topn, sorted=True)
    indices = tf.cast(indices, tf.float32)
    batch_ids = tf.ones_like(indices, dtype=tf.float32) * tf.expand_dims(
          tf.cast(tf.range(tf.shape(input=indices)[0]),dtype=tf.float32), 1)
    
    # batchsize x topk x 2
    nd_indices = tf.cast(tf.stack([batch_ids, indices], axis=-1),tf.int32)
    
    sorted_labels, sorted_logits = [tf.gather_nd(f, nd_indices) for f in [logits, labels]]
    
    raw_max = tf.reduce_max(sorted_logits, axis=1, keepdims=True)
    sorted_logits = sorted_logits - raw_max
    sums = tf.cumsum(tf.exp(sorted_logits), axis=1, reverse=True)
    sums = tf.math.log(sums) - sorted_logits
    
    negative_log_likelihood = tf.reduce_sum(
        input_tensor=sums, axis=1, keepdims=True)
    return negative_log_likelihood
    

ApproxNDCG

y_pred = tf.convert_to_tensor([[0.5, 0.2, 0.1, 0.4, 1.0, -1.0]], dtype=tf.float32)
y_true = tf.convert_to_tensor([[1.0, 2.0, 2.0, 4.0, 1.0, 4.0]], dtype=tf.float32)
loss = approx_ndcg_loss(y_true, y_pred)

def approx_rank(logits, temperature=0.1):
    """计算列表的近似排名
    rank(j) = 1 + sum_{j neq i} I_{s_j > s_i}, I is indicator function
    
    I_{s_j > s_i} = approx 1 / (1 + exp(-(s_j - s_i) / temperature))
    """
    list_size = tf.shape(input=logits)[1]
    x = tf.tile(tf.expand_dims(logits, 2), [1, 1, list_size])
    y = tf.tile(tf.expand_dims(logits, 1), [1, list_size, 1])
    
    pairs = tf.sigmoid((y - x) / temperature)
    rank = tf.reduce_sum(input_tensor=pairs, axis=-1) + 0.5
    
    return rank

def ndcg(labels, ranks):
    """
    labels: [Batch_size, list_size]
    ranks:  [Batch_size, list_size]
    """
    discounts = 1. / tf.math.log1p(tf.cast(ranks, dtype=tf.float32))
    gains = tf.math.pow(2.0, tf.cast(labels, dtype=tf.float32))
    
    dcg = tf.reduce_sum(input_tensor=gains * discounts, axis=-1, keepdims=True)
    
    def inverse_max_dcg(labels):
        """calc inverse of ideal dcg
        """
        ideal_sorted_labels = tf.sort(labels, axis=-1, direction='DESCENDING')
        rank = tf.range(tf.shape(ideal_sorted_labels)[1]) + 1
        
        discounted_gain = tf.math.pow(2.0, ideal_sorted_labels) / tf.math.log1p(tf.cast(rank, dtype=tf.float32))
        discounted_gain = tf.reduce_sum(discounted_gain, axis=-1)
        
        return 1.0 / discounted_gain
    
    normalized_dcg = dcg * inverse_max_dcg(labels)
    return normalized_dcg

def approx_ndcg_loss(labels, logits):
    ranks = approx_rank(logits)
    
    return -ndcg(labels, ranks)