InstaDeep open-sources CATX library, enabling contextual bandits in JAX



As part of its commitment to giving back to the international AI and technology community, InstaDeep is pleased to announce it has made its new CATX library available on GitHub. CATX translates the existing CATS algorithm to JAX – the high-performance numerical computing and machine learning research library from Google.

CATX is based on the original CATS algorithm from the paper “Efficient Contextual Bandits with Continuous Actions” by Majzoubi et al. The original algorithm is implemented in the Vowpal Wabbit library from Microsoft and, initially, Yahoo. Although a highly useful tool, this library suffered from a shortcoming: an inability to use more powerful, advanced neural networks within the structure of the CATS algorithm. CATX rectifies this key deficiency and enables the use of custom neural networks, alongside potential computation speed improvements, thanks to JAX’s just-in-time compilation.

CATX is aimed at users facing continuous action contextual bandits problems – any problem where you need to take continuous actions while maximising the desired reward (and, consequently, minimising cost, time or effort expenditures). Contextual bandits settings, where the exploration-exploitation trade-off needs to be dealt with, can be found in many industries and use cases. CATX offers a valuable boost to this type of problem, by implementing contextual bandits with continuous actions in JAX, and allowing custom neural networks in the tree structure of the CATS algorithm.

Speaking on the release of CATX, InstaDeep’s Cyprien Courtot, Research Engineer and co-creator of CATX, said, “We’re delighted to share CATX with the open source community. By utilising JAX, we’ve been able to apply contextual bandits to complex and noisy real world use-cases, and we look forward to many users applying it for other practical problems and applications.”

CATX has been published as an open source library and the full repo and documentation is available now from the company’s GitHub page.

>>Want to work on projects like this? Check out our exciting job opportunities: