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

Tree-based Search Graph for Approximate Nearest Neighbor Search

Si Adek Keras is software VR dangerous object detection.

ViewFormer: NeRF-free Neural Rendering from Few Images Using Transformers

Cowsay - A rewrite of cowsay in python

[NeurIPS 2021] “Improving Contrastive Learning on Imbalanced Data via Open-World Sampling”,

Editing a classifier by rewriting its prediction rules

Project dự đoán giá cổ phiếu bằng thuật toán LSTM gồm: code train và code demo

Cross-Task Consistency Learning Framework for Multi-Task Learning

Code repository for our paper "Learning to Generate Scene Graph from Natural Language Supervision" in ICCV 2021

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

ICCV2021: Code for 'Spatial Uncertainty-Aware Semi-Supervised Crowd Counting'

SporeAgent: Reinforced Scene-level Plausibility for Object Pose Refinement

MoViNets PyTorch implementation: Mobile Video Networks for Efficient Video Recognition;

PyTorch implementation of Weak-shot Fine-grained Classification via Similarity Transfer

This repository contains an implementation of ConvMixer for the ICLR 2022 submission "Patches Are All You Need?".

QQ Browser 2021 AI Algorithm Competition Track 1 1st Place Program

This repository contains code accompanying the paper "An End-to-End Chinese Text Normalization Model based on Rule-Guided Flat-Lattice Transformer"

[CVPR 2022] Semi-Supervised Semantic Segmentation Using Unreliable Pseudo-Labels

[NeurIPS 2021] Shape from Blur: Recovering Textured 3D Shape and Motion of Fast Moving Objects

DCGAN LSGAN WGAN-GP DRAGAN PyTorch