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

Source code for "FastBERT: a Self-distilling BERT with Adaptive Inference Time".

Read Like Humans: Autonomous, Bidirectional and Iterative Language Modeling for Scene Text Recognition

Automatically replace ONNX's RandomNormal node with Constant node.

Submanifold sparse convolutional networks

Few-Shot Object Detection via Association and DIscrimination

Label Studio is a multi-type data labeling and annotation tool with standardized output format

PyTorch implementation of our ICCV 2019 paper: Liquid Warping GAN: A Unified Framework for Human Motion Imitation, Appearance Transfer and Novel View Synthesis

Python scripts form performing stereo depth estimation using the CoEx model in ONNX.

[ICCV 2021 Oral] Just Ask: Learning to Answer Questions from Millions of Narrated Videos

a basic code repository for basic task in CV(classification,detection,segmentation)

Unsupervised Discovery of Object Radiance Fields

JumpDiff: Non-parametric estimator for Jump-diffusion processes for Python

Grammar Induction using a Template Tree Approach

MAGMA - a GPT-style multimodal model that can understand any combination of images and language

Evaluation Pipeline for our ECCV2020: Journey Towards Tiny Perceptual Super-Resolution.

Implementation of Artificial Neural Network Algorithm

Official implementation of "Implicit Neural Representations with Periodic Activation Functions"

Confident Semantic Ranking Loss for Part Parsing

Self-supervised Augmentation Consistency for Adapting Semantic Segmentation (CVPR 2021)

Pytorch version of SfmLearner from Tinghui Zhou et al.