activity
20212025
most citedBlackJAX: Composable Bayesian inference in JAX

8 citations · 19 across the 9 of their papers we have counts for

collaborators

9 papers

cs.CL2025

Scale-invariant Attention

Ben Anson, Xi Wang, Laurence Aitchison

One persistent challenge in LLM research is the development of attention mechanisms that are able to generalise from training on shorter contexts to inference on longer contexts. W…

cs.LG2025

Revisiting Gradient Descent: A Dual-Weight Method for Improved Learning

Xi Wang

We introduce a novel framework for learning in neural networks by decomposing each neuron's weight vector into two distinct parts, and , thereby modeling contrastive inf…

cs.LG2024

Batch size invariant Adam

Xi Wang, Laurence Aitchison

We propose a batch size invariant version of Adam, for use in large-scale, distributed settings, in which the mini-batch is divided into micro-batches which are distributed among w…

cs.MS2024★ 8 cited

BlackJAX: Composable Bayesian inference in JAX

Alberto Cabezas, Adrien Corenflos, Junpeng Lao +19

BlackJAX is a library implementing sampling and variational inference algorithms commonly used in Bayesian computation. It is designed for ease of use, speed, and modularity by tak…

cs.LG2023★ 6 cited

LoRA ensembles for large language model fine-tuning

Xi Wang, Laurence Aitchison, Maja Rudolph

Finetuned LLMs often exhibit poor uncertainty quantification, manifesting as overconfidence, poor calibration, and unreliable prediction results on test data or out-of-distribution…

cs.LG2023★ 1 cited

Bayesian Low-rank Adaptation for Large Language Models

Adam X. Yang, Maxime Robeyns, Xi Wang +1

Low-rank adaptation (LoRA) has emerged as a new paradigm for cost-efficient fine-tuning of large language models (LLMs). However, fine-tuned LLMs often become overconfident especia…

cs.LG2022

Joint control variate for faster black-box variational inference

Xi Wang, Tomas Geffner, Justin Domke

Black-box variational inference performance is sometimes hindered by the use of gradient estimators with high variance. This variance comes from two sources of randomness: Data sub…

cs.LG2022

Robustness to corruption in pre-trained Bayesian neural networks

Xi Wang, Laurence Aitchison

We develop ShiftMatch, a new training-data-dependent likelihood for robustness to corruption in Bayesian neural networks (BNNs). ShiftMatch is inspired by the training-data-depende…

stat.ML2021★ 4 cited

Bayesian OOD detection with aleatoric uncertainty and outlier exposure

Xi Wang, Laurence Aitchison

Typical Bayesian approaches to OOD detection use epistemic uncertainty. Surprisingly from the Bayesian perspective, there are a number of methods that successfully use aleatoric un…