Rax is a Learning-to-Rank library written in JAX

Last update: Dec 27, 2022

Related tags

Overview

🦖 Rax: Composable Learning to Rank using JAX

Rax is a Learning-to-Rank library written in JAX. Rax provides off-the-shelf implementations of ranking losses and metrics to be used with JAX. It provides the following functionality:

Ranking losses (rax.*_loss): rax.softmax_loss, rax.pairwise_logistic_loss, ...
Ranking metrics (rax.*_metric): rax.mrr_metric, rax.ndcg_metric, ...
Transformations (rax.*_t12n): rax.approx_t12n, rax.gumbel_t12n, ...

Ranking

A ranking problem is different from traditional classification/regression problems in that its objective is to optimize for the correctness of the relative order of a list of examples (e.g., documents) for a given context (e.g., a query). Rax provides support for ranking problems within the JAX ecosystem. It can be used in, but is not limited to, the following applications:

Search: ranking a list of documents with respect to a query.
Recommendation: ranking a list of items given a user as context.
Question Answering: finding the best answer from a list of candidates.
Dialogue System: finding the best response from a list of responses.

Synopsis

In a nutshell, given the scores and labels for a list of items, Rax can compute various ranking losses and metrics:

import jax.numpy as jnp
import rax

scores = jnp.asarray([2.2, -1.3, 5.4])  # output of a model.
labels = jnp.asarray([1., 0., 0.])      # indicates doc 1 is relevant.

rax.ndcg_metric(scores, labels)         # computes a ranking metric.
rax.pairwise_hinge_loss(scores, labels) # computes a ranking loss.

All of the Rax losses and metrics are purely functional and compose well with standard JAX transformations. Additionally, Rax provides ranking-specific transformations so you can build new ranking losses. An example is rax.approx_t12n, which can be used to transform any (non-differentiable) ranking metric into a differentiable loss. For example:

loss_fn = rax.approx_t12n(rax.ndcg_metric)
loss_fn(scores, labels)            # differentiable approx ndcg loss.
jax.grad(loss_fn)(scores, labels)  # computes gradients w.r.t. scores.

Examples

See the examples/ directory for complete examples on how to use Rax.

You might also like...

[ICLR 2021] Rank the Episodes: A Simple Approach for Exploration in Procedurally-Generated Environments.

[ICLR 2021] RAPID: A Simple Approach for Exploration in Reinforcement Learning This is the Tensorflow implementation of ICLR 2021 paper Rank the Episo

48 Nov 21, 2022

Rank 1st in the public leaderboard of ScanRefer (2021-03-18)

InstanceRefer InstanceRefer: Cooperative Holistic Understanding for Visual Grounding on Point Clouds through Instance Multi-level Contextual Referring

63 Dec 7, 2022

Code for "LoRA: Low-Rank Adaptation of Large Language Models"

LoRA: Low-Rank Adaptation of Large Language Models This repo contains the implementation of LoRA in GPT-2 and steps to replicate the results in our re

394 Jan 8, 2023

Official PyTorch Implementation of Rank & Sort Loss [ICCV2021]

Rank & Sort Loss for Object Detection and Instance Segmentation The official implementation of Rank & Sort Loss. Our implementation is based on mmdete

229 Dec 20, 2022

This is the pytorch implementation for the paper: Generalizable Mixed-Precision Quantization via Attribution Rank Preservation, which is accepted to ICCV2021.

GMPQ: Generalizable Mixed-Precision Quantization via Attribution Rank Preservation This is the pytorch implementation for the paper: Generalizable Mix

18 Sep 2, 2022

COD-Rank-Localize-and-Segment (CVPR2021)

Comments

Add Kendall's tau
Kendall's tau (https://en.wikipedia.org/wiki/Kendall_rank_correlation_coefficient) is a ranking metric which we could add to Rax. We can reference the SciPy version of Kendall's tau (https://github.com/scipy/scipy/blob/v1.9.0/scipy/stats/_stats_py.py#L5015-L5222) as inspiration or comparison for the Rax implementation.

Expected usage:

scores = jnp.array([0., 1., 2.]) labels = jnp.array([1., 2., 0.]) rax.kendalltau_metric(scores, labels) # should produce Kendall's tau-b.
enhancement
opened by rjagerman 0

Releases(v0.2.0)

v0.2.0(Aug 11, 2022)
What's Changed

New losses: rax.poly1_softmax_loss and rax.unique_softmax_loss.

New example: T5X integration.

Removed dependencies' version upper-bounds.

Minor documentation fixes.

Full Changelog: https://github.com/google/rax/compare/v0.1.0...v0.2.0
Source code(tar.gz)
Source code(zip)
v0.1.0(Feb 24, 2022)
This is the initial release of Rax: a Learning-to-Rank (LTR) library built on top of JAX. It includes the following functionality:

Ranking losses (rax.*_loss).

Ranking metrics (rax.*_metric).

Function transformations (rax.*_t12n).

Source code(tar.gz)
Source code(zip)

Rax is a Learning-to-Rank library written in JAX

Related tags

Overview

🦖 Rax: Composable Learning to Rank using JAX

Ranking

Synopsis

Examples

You might also like...

[ICLR 2021] Rank the Episodes: A Simple Approach for Exploration in Procedurally-Generated Environments.

Rank 1st in the public leaderboard of ScanRefer (2021-03-18)

Code for "LoRA: Low-Rank Adaptation of Large Language Models"

Official PyTorch Implementation of Rank & Sort Loss [ICCV2021]

This is the pytorch implementation for the paper: Generalizable Mixed-Precision Quantization via Attribution Rank Preservation, which is accepted to ICCV2021.

COD-Rank-Localize-and-Segment (CVPR2021)

Gradient-free global optimization algorithm for multidimensional functions based on the low rank tensor train format

ViViT: Curvature access through the generalized Gauss-Newton's low-rank structure

This is the solution for 2nd rank in Kaggle competition: Feedback Prize - Evaluating Student Writing.

Comments

Add Kendall's tau

Releases(v0.2.0)

v0.2.0(Aug 11, 2022)

What's Changed

v0.1.0(Feb 24, 2022)

Owner

Google

The source code for 'Noisy-Labeled NER with Confidence Estimation' accepted by NAACL 2021

Adversarial-autoencoders - Tensorflow implementation of Adversarial Autoencoders

A bunch of random PyTorch models using PyTorch's C++ frontend

The official PyTorch code for NeurIPS 2021 ML4AD Paper, "Does Thermal data make the detection systems more reliable?"

NLMpy - A Python package to create neutral landscape models

Reinforcement learning models in ViZDoom environment

Code for the paper "Generative design of breakwaters usign deep convolutional neural network as a surrogate model"

nfelo: a power ranking, prediction, and betting model for the NFL

Simple node deletion tool for onnx.

(ICCV 2021 Oral) Re-distributing Biased Pseudo Labels for Semi-supervised Semantic Segmentation: A Baseline Investigation.

Implementation of fast algorithms for Maximum Spanning Tree (MST) parsing that includes fast ArcMax+Reweighting+Tarjan algorithm for single-root dependency parsing.

Repository features UNet inspired architecture used for segmenting lungs on chest X-Ray images

Neural Scene Flow Prior (NeurIPS 2021 spotlight)

you can add any codes in any language by creating its respective folder (if already not available).

Implementation of Restricted Boltzmann Machine (RBM) and its variants in Tensorflow

A collection of loss functions for medical image segmentation

Out-of-Domain Human Mesh Reconstruction via Dynamic Bilevel Online Adaptation

Companion repository to the paper accepted at the 4th ACM SIGSPATIAL International Workshop on Advances in Resilient and Intelligent Cities

[ECCVW2020] Robust Long-Term Object Tracking via Improved Discriminative Model Prediction (RLT-DiMP)

A Neural Net Training Interface on TensorFlow, with focus on speed + flexibility