Our papers on variance-aware linear UCB for neural contextual bandits [1] and on the learning dynamics of LoRA [2] have been accepted to the International Conference on Artificial Intelligence and Statistics. Congrats Ziqing!
[Bibtex] [Abstract] [Download PDF]
By leveraging the representation power of deepneuralnetworks, neuralupperconfidence bound (UCB) algorithms have shown success in contextual bandits. To further balance the exploration and exploitation, we propose Neural- σ2-LinearUCB, a variance-aware algo- rithm that utilizes σ2 t, i.e., an upper bound of the reward noise variance at round t, to enhance the uncertainty quantification quality of the UCB, resulting in a regret performance improvement. We provide an oracle version for our algorithm characterized by an oracle variance upper bound σ2 tand a practical ver- sion with a novel estimation for this variance bound. Theoretically, we provide rigorous re- gret analysis for both versions and prove that our oracle algorithm achieves a better regret guarantee than other neural-UCB algorithms in the neural contextual bandits setting. Em- pirically, ourpracticalmethodenjoysasimilar computational efficiency, while outperforming state-of-the-art techniques by having a better calibration and lower regret across multiple standard settings, including on the synthetic, UCI, MNIST, and CIFAR-10 datasets.
@inproceedings{bml2025aistats,
abstract = {By leveraging the representation power of deepneuralnetworks, neuralupperconfidence bound (UCB) algorithms have shown success in contextual bandits. To further balance the exploration and exploitation, we propose Neural- σ2-LinearUCB, a variance-aware algo- rithm that utilizes σ2 t, i.e., an upper bound of the reward noise variance at round t, to enhance the uncertainty quantification quality of the UCB, resulting in a regret performance improvement. We provide an oracle version for our algorithm characterized by an oracle variance upper bound σ2 tand a practical ver- sion with a novel estimation for this variance bound. Theoretically, we provide rigorous re- gret analysis for both versions and prove that our oracle algorithm achieves a better regret guarantee than other neural-UCB algorithms in the neural contextual bandits setting. Em- pirically, ourpracticalmethodenjoysasimilar computational efficiency, while outperforming state-of-the-art techniques by having a better calibration and lower regret across multiple standard settings, including on the synthetic, UCI, MNIST, and CIFAR-10 datasets.},
author = {Bui, Ha Manh and Mallada, Enrique and Liu, Anqi},
booktitle = {International Conference on Artificial Intelligence and Statistics (AISTATS)},
grants = {No Grant},
month = {4},
publisher = {PMLR},
record = {accepted Jan 2025, submitted Oct 2024},
series = {Proceedings of Machine Learning Research},
title = {Variance-Aware Linear UCB with Deep Representation for Neural Contextual Bandits},
url = {https://mallada.ece.jhu.edu/pubs/2025-AISTATS-BML.pdf},
year = {2025}
}[Bibtex] [Abstract] [Download PDF]
Despite the empirical success of Low-Rank Adaptation (LoRA) in fine-tuning pretrained models, there is little theoretical understanding of how first-order methods with carefully crafted initialization adapt models to new tasks. In this work, we take the first step towards bridging this gap by theoretically analyzing the learning dynamics of LoRA for matrix factorization (MF) under gradient flow (GF), emphasizing the crucial role of initialization. For small initialization, we theoretically show that GF converges to a neighborhood of the optimal solution, with smaller initialization leading to lower final error. Our analysis shows that the final error is affected by the misalignment between the singular spaces of the pre-trained model and the target matrix, and reducing the initialization scale improves alignment. To address this misalignment, we propose a spectral initialization for LoRA in MF and theoretically prove that GF with small spectral initialization converges to the fine-tuning task with arbitrary precision. Numerical experiments from MF and image classification validate our findings.
@inproceedings{xmmltmv2025aistats,
abstract = {Despite the empirical success of Low-Rank Adaptation (LoRA) in fine-tuning pretrained models, there is little theoretical understanding of how first-order methods with carefully crafted initialization adapt models to new tasks. In this work, we take the first step towards bridging this gap by theoretically analyzing the learning dynamics of LoRA for matrix factorization (MF) under gradient flow (GF), emphasizing the crucial role of initialization. For small initialization, we theoretically show that GF converges to a neighborhood of the optimal solution, with smaller initialization leading to lower final error. Our analysis shows that the final error is affected by the misalignment between the singular spaces of the pre-trained model and the target matrix, and reducing the initialization scale improves alignment. To address this misalignment, we propose a spectral initialization for LoRA in MF and theoretically prove that GF with small spectral initialization converges to the fine-tuning task with arbitrary precision. Numerical experiments from MF and image classification validate our findings.},
author = {Xu, Ziqing and Min, Hancheng and MacDonald, Lachlan Ewen and Luo, Jinqi and Tarmoun, Salma and Mallada, Enrique and Vidal, Rene},
booktitle = {International Conference on Artificial Intelligence and Statistics (AISTATS)},
grants = {Global Centers},
month = {4},
publisher = {PMLR},
record = {accepted Jan 2024, submitted Oct 2024},
series = {Proceedings of Machine Learning Research},
title = {Understanding the Learning Dynamics of LoRA: A Gradient Flow Perspective on Low-Rank Adaptation in Matrix Factorization},
url = {https://mallada.ece.jhu.edu/pubs/2025-AISTATS-XMMLTMV.pdf},
year = {2025}
}