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

Official implementation of paper "Query2Label: A Simple Transformer Way to Multi-Label Classification".

PSML: A Multi-scale Time-series Dataset for Machine Learning in Decarbonized Energy Grids

DenseCLIP: Language-Guided Dense Prediction with Context-Aware Prompting

Pytorch implementation of set transformer

📖 Deep Attentional Guided Image Filtering

Supplementary code for the experiments described in the 2021 ISMIR submission: Leveraging Hierarchical Structures for Few Shot Musical Instrument Recognition.

Official TensorFlow code for the forthcoming paper

A Collection of Papers and Codes for ICCV2021 Low Level Vision and Image Generation

This package is for running the semantic SLAM algorithm using extracted planar surfaces from the received detection

This script runs neural style transfer against the provided content image.

The official implementation of our CVPR 2021 paper - Hybrid Rotation Averaging: A Fast and Robust Rotation Averaging Approach

Square Root Bundle Adjustment for Large-Scale Reconstruction

This is the code used in the paper "Entity Embeddings of Categorical Variables".

This is an open source python repository for various python tests

MINOS: Multimodal Indoor Simulator

Self-supervised Deep LiDAR Odometry for Robotic Applications

Discriminative Region Suppression for Weakly-Supervised Semantic Segmentation

DeepMind's software stack for physics-based simulation and Reinforcement Learning environments, using MuJoCo.

TigerLily: Finding drug interactions in silico with the Graph.

Codes for our paper "SentiLARE: Sentiment-Aware Language Representation Learning with Linguistic Knowledge" (EMNLP 2020)