Publications

2024

1. 

   Bayesian Optimization of Antibodies Informed by a Generative Model of Evolving Sequences

   Alan Nawzad Amin, Nate Gruver*, Yucen Lily Li*, Yilun Kuang*, Hunter Elliott, Calvin McCarter, Aniruddh Raghu, Peyton Greenside, and Andrew Gordon Wilson

   NeurIPS Workshop on AI for New Drug Modalities (NeurIPS Workshop), 2024 Spotlight

   Abs Bib HTML

   To build effective therapeutics, biologists iteratively mutate antibody sequences to improve binding and stability. Proposed mutations can be informed by previous measurements or by learning from large antibody databases to predict only typical antibodies. Unfortunately, the space of typical antibodies is enormous to search, and experiments often fail to find suitable antibodies on a budget. Here we introduce Clone-informed Bayesian Optimization (CloneBO), a Bayesian optimization procedure that efficiently optimizes antibodies in the lab by teaching a generative model how our immune system optimizes antibodies. Our immune system makes antibodies by iteratively evolving specific portions of their sequences to bind their target strongly and stably, resulting in a set of related, evolving sequences known as a clonal family. We train a large language model, CloneLM, on hundreds of thousands of clonal families and use it to design sequences with mutations that are most likely to optimize an antibody within the human immune system. We guide our designs to fit previous measurements using a twisted sequential Monte Carlo procedure. We show that CloneBO optimizes antibodies substantially more efficiently than previous methods in realistic in silico experiments and designs stronger and more stable binders in in vitro wet lab experiments.

   @article{amin2024clonebo,
     title = {Bayesian Optimization of Antibodies Informed by a Generative Model of Evolving Sequences},
     author = {Amin, Alan Nawzad and Gruver*, Nate and Li*, Yucen Lily and Kuang*, Yilun and Elliott, Hunter and McCarter, Calvin and Raghu, Aniruddh and Greenside, Peyton and Wilson, Andrew Gordon},
     journal = {NeurIPS Workshop on AI for New Drug Modalities (NeurIPS Workshop), 2024},
     spotlight = {true},
     year = {2024},
   }

2. 

   Unlocking Tokens as Data Points for Generalization Bounds on Larger Language Models

   Sanae Lotfi*, Yilun Kuang*, Brandon Amos, Micah Goldblum, Marc Finzi, and Andrew Gordon Wilson

   Neural Information Processing Systems (NeurIPS), 2024 Spotlight
   ICML Workshop on Theoretical Foundations of Foundation Models (ICML Workshop), 2024 Best Paper Award

   Abs Bib HTML

   Large language models (LLMs) with billions of parameters excel at predicting the next token in a sequence. Recent work computes non-vacuous compression-based generalization bounds for LLMs, but these bounds are vacuous for large models at the billion-parameter scale. Moreover, these bounds are obtained through restrictive compression techniques, bounding compressed models that generate low-quality text. Additionally, the tightness of these existing bounds depends on the number of IID documents in a training set rather than the much larger number of non-IID constituent tokens, leaving untapped potential for tighter bounds. In this work, we instead use properties of martingales to derive generalization bounds that benefit from the vast number of tokens in LLM training sets. Since a dataset contains far more tokens than documents, our generalization bounds not only tolerate but actually benefit from far less restrictive compression schemes. With Monarch matrices, Kronecker factorizations, and post-training quantization, we achieve non-vacuous generalization bounds for LLMs as large as LLaMA2-70B. Unlike previous approaches, our work achieves the first non-vacuous bounds for models that are deployed in practice and generate high-quality text.

   @article{lotfi2024tokenbounds,
     title = {Unlocking Tokens as Data Points for Generalization Bounds on Larger Language Models},
     author = {Lotfi*, Sanae and Kuang*, Yilun and Amos, Brandon and Goldblum, Micah and Finzi, Marc and Wilson, Andrew Gordon},
     journal = {Neural Information Processing Systems (NeurIPS), 2024; ICML Workshop on Theoretical Foundations of Foundation Models (ICML Workshop), 2024},
     spotlight = {true},
     bestpaper = {true},
     year = {2024},
   }

3. 

   Non-Vacuous Generalization Bounds for Large Language Models

   Sanae Lotfi*, Marc Finzi*, Yilun Kuang*, Tim G. J. Rudner, Micah Goldblum, and Andrew Gordon Wilson

   International Conference on Machine Learning (ICML), 2024
   NeurIPS Workshop on Self-Supervised Learning & Mathematics of Modern Machine Learning (NeurIPS Workshop), 2023

   Abs Bib HTML

   Modern language models can contain billions of parameters, raising the question of whether they can generalize beyond the training data or simply regurgitate their training corpora. We provide the first non-vacuous generalization bounds for pretrained large language models (LLMs), indicating that language models are capable of discovering regularities that generalize to unseen data. In particular, we derive a compression bound that is valid for the unbounded log-likelihood loss using prediction smoothing, and we extend the bound to handle subsampling, accelerating bound computation on massive datasets. To achieve the extreme level of compression required for non-vacuous generalization bounds, we devise SubLoRA, a low-dimensional non-linear parameterization. Using this approach, we find that larger models have better generalization bounds and are more compressible than smaller models.

   @article{lotfi2024nonvacuous,
     title = {Non-Vacuous Generalization Bounds for Large Language Models},
     author = {Lotfi*, Sanae and Finzi*, Marc and Kuang*, Yilun and Rudner, Tim G. J. and Goldblum, Micah and Wilson, Andrew Gordon},
     journal = {International Conference on Machine Learning (ICML), 2024; NeurIPS Workshop on Self-Supervised Learning & Mathematics of Modern Machine Learning (NeurIPS Workshop), 2023},
     year = {2024},
   }

2023

1. 

   Unsupervised Learning on Spontaneous Retinal Activity Leads to Efficient Neural Representation Geometry

   Andrew Ligeralde*, Yilun Kuang*, Thomas Yerxa, Miah N. Pitcher, Marla Feller, and SueYeon Chung

   NeurIPS Workshop on Unifying Representations in Neural Models (NeurIPS Workshop), 2023

   Abs Bib HTML

   Prior to the onset of vision, neurons in the developing mammalian retina spontaneously fire in correlated activity patterns known as retinal waves. Experimental evidence suggests that retinal waves strongly influence the emergence of sensory representations before visual experience. We aim to model this early stage of functional development by using movies of neurally active developing retinas as pre-training data for neural networks. Specifically, we pre-train a ResNet-18 with an unsupervised contrastive learning objective (SimCLR) on both simulated and experimentally-obtained movies of retinal waves, then evaluate its performance on image classification tasks. We find that pre-training on retinal waves significantly improves performance on tasks that test object invariance to spatial translation, while slightly improving performance on more complex tasks like image classification. Notably, these performance boosts are realized on held-out natural images even though the pre-training procedure does not include any natural image data. We then propose a geometrical explanation for the increase in network performance, namely that the spatiotemporal characteristics of retinal waves facilitate the formation of separable feature representations. In particular, we demonstrate that networks pre-trained on retinal waves are more effective at separating image manifolds than randomly initialized networks, especially for manifolds defined by sets of spatial translations. These findings indicate that the broad spatiotemporal properties of retinal waves prepare networks for higher order feature extraction.

   @article{ligeralde2023unsupervised,
     title = {Unsupervised Learning on Spontaneous Retinal Activity Leads to Efficient Neural Representation Geometry},
     author = {Ligeralde*, Andrew and Kuang*, Yilun and Yerxa, Thomas and N. Pitcher, Miah and Feller, Marla and Chung, SueYeon},
     journal = {NeurIPS Workshop on Unifying Representations in Neural Models (NeurIPS Workshop), 2023},
     year = {2023},
   }

2. 

   Learning Efficient Coding of Natural Images with Maximum Manifold Capacity Representations

   Thomas Yerxa, Yilun Kuang, Eero Simoncelli, and SueYeon Chung

   Neural Information Processing Systems (NeurIPS), 2023
   Computational and Systems Neuroscience (COSYNE), 2023

   Abs Bib HTML

   The efficient coding hypothesis proposes that the response properties of sensory systems are adapted to the statistics of their inputs such that they capture maximal information about the environment, subject to biological constraints. While elegant, information theoretic properties are notoriously difficult to measure in practical settings or to employ as objective functions in optimization. This difficulty has necessitated that computational models designed to test the hypothesis employ several different information metrics ranging from approximations and lower bounds to proxy measures like reconstruction error. Recent theoretical advances have characterized a novel and ecologically relevant efficiency metric, the manifold capacity, which is the number of object categories that may be represented in a linearly separable fashion. However, calculating manifold capacity is a computationally intensive iterative procedure that until now has precluded its use as an objective. Here we outline the simplifying assumptions that allow manifold capacity to be optimized directly, yielding Maximum Manifold Capacity Representations (MMCR). The resulting method is closely related to and inspired by advances in the field of self supervised learning (SSL), and we demonstrate that MMCRs are competitive with state of the art results on standard SSL benchmarks. Empirical analyses reveal differences between MMCRs and representations learned by other SSL frameworks, and suggest a mechanism by which manifold compression gives rise to class separability. Finally we evaluate a set of SSL methods on a suite of neural predictivity benchmarks, and find MMCRs are higly competitive as models of the ventral stream.

   @article{yerxa2023learning,
     title = {Learning Efficient Coding of Natural Images with Maximum Manifold Capacity Representations},
     author = {Yerxa, Thomas and Kuang, Yilun and Simoncelli, Eero and Chung, SueYeon},
     journal = {Neural Information Processing Systems (NeurIPS), 2023; Computational and Systems Neuroscience (COSYNE), 2023},
     year = {2023},
   }