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

An Implementation of Fully Convolutional Networks in Tensorflow.

This is the offical website for paper ''Category-consistent deep network learning for accurate vehicle logo recognition''

Adversarial Adaptation with Distillation for BERT Unsupervised Domain Adaptation

Few-Shot Object Detection via Association and DIscrimination

Unofficial PyTorch implementation of Attention Free Transformer (AFT) layers by Apple Inc.

[ICCV 2021] Learning A Single Network for Scale-Arbitrary Super-Resolution

Aws-machine-learning-university-accelerated-tab - Machine Learning University: Accelerated Tabular Data Class

Official repository for the ICLR 2021 paper Evaluating the Disentanglement of Deep Generative Models with Manifold Topology

Learning hidden low dimensional dyanmics using a Generalized Onsager Principle and neural networks

GUPNet - Geometry Uncertainty Projection Network for Monocular 3D Object Detection

Image Lowpoly based on Centroid Voronoi Diagram via python-opencv and taichi

Res2Net for Instance segmentation and Object detection using MaskRCNN

Code for the published paper : Learning to recognize rare traffic sign

Collection of generative models in Pytorch version.

Differentiable Annealed Importance Sampling (DAIS)

Python scripts form performing stereo depth estimation using the HITNET model in Tensorflow Lite.

Attention-based Transformation from Latent Features to Point Clouds (AAAI 2022)

A PyTorch implementation of EfficientNet and EfficientNetV2 (coming soon!)

SMCA replication There are no extra compiled components in SMCA DETR and package dependencies are minimal

Node Editor Plug for Blender