Using foundation models as flexible priors and contextual memory for downstream inference.
Prior knowledge is a powerful form of context in statistical inference. Traditionally, applying such knowledge required expert intervention on each new problem. Today, foundation models encode broad domain knowledge in a reusable, black-box format. Our work focuses on extracting and operationalizing this implicit knowledge by connecting foundation models to structured, parametric statistical models.
We are building a bi-directional bridge between foundation models and structured statistical estimation:
Clinical machine learning models must adapt to new settings such as different hospitals, clinicians, or patient populations. These differing environments present related but subtly distinct tasks, where diseases and medical interventions share common foundations but vary in meaningful ways. In contrast to one-size-fits-all invariant feature learning, we believe representing meaningful differences between domains and adapting to these differences will improve accuracy, utility, and interpretability of machine learning in health. Here, we introduce Retrieval-Augmented Generation of Interpretable Models (RAG-IM), a highly performant method for adapting statistical models to new domains based on their descriptions. By leveraging the strengths of Retrieval-Augmented Generation (RAG), our framework retrieves relevant models from related tasks and combines them with contextual insights from pre-trained language models. RAG-IM generates task-specific, interpretable models that perform reliably, even in few-shot and zero-shot scenarios where data are limited or completely unavailable. Through experiments on 7487 related tasks, we find that RAG-IM is a promising general-purpose platform to enable model-based analysis to data-limited and heterogeneous regimes by connecting statistical analysis with natural language.
@article{mahbub2024one,author={Mahbub, Sazan and Ellington, Caleb and Alinejad, Sina and Wen, Kevin and Lengerich, Ben and Xing, Eric P.},title={From One to Zero: RAG-IM Adapts Language Models for Interpretable Zero-Shot Clinical Predictions},journal={NeurIPS Workshop on Adaptive Foundation Models (NeurIPS AFM)},year={2024},}
Recent years have seen important advances in the building of interpretable models, machine learning models that are designed to be easily understood by humans. In this work, we show that large language models (LLMs) are remarkably good at working with interpretable models, too. In particular, we show that LLMs can describe, interpret, and debug Generalized Additive Models (GAMs). Combining the flexibility of LLMs with the breadth of statistical patterns accurately described by GAMs enables dataset summarization, question answering, and model critique. LLMs can also improve the interaction between domain experts and interpretable models, and generate hypotheses about the underlying phenomenon. We release TalkToEBM as an open-source LLM-GAM interface.
@article{bordt2024data,author={Bordt, Sebastian and Lengerich, Ben and Nori, Harsha and Caruana, Rich},title={Data Science with LLMs and Interpretable Models},journal={AAAI Explainable AI for Science},year={2023},informal_venue={AAAI XAI4Sci},keywords={Interpretable, LLMs},}
LLMs Understand Glass-Box Models, Discover Surprises, and Suggest Repairs
We show that large language models (LLMs) are remarkably good at working with interpretable models that decompose complex outcomes into univariate graph-represented components. By adopting a hierarchical approach to reasoning, LLMs can provide comprehensive model-level summaries without ever requiring the entire model to fit in context. This approach enables LLMs to apply their extensive background knowledge to automate common tasks in data science such as detecting anomalies that contradict prior knowledge, describing potential reasons for the anomalies, and suggesting repairs that would remove the anomalies. We use multiple examples in healthcare to demonstrate the utility of these new capabilities of LLMs, with particular emphasis on Generalized Additive Models (GAMs). Finally, we present the package 𝚃𝚊𝚕𝚔𝚃𝚘𝙴𝙱𝙼 as an open-source LLM-GAM interface.
@article{lengerich2023llms,author={Lengerich, Benjamin J. and Bordt, Sebastian and Nori, Harsha and Nunnally, Mark E. and Aphinyanaphongs, Yin and Kellis, Manolis and Caruana, Rich},title={LLMs Understand Glass-Box Models, Discover Surprises, and Suggest Repairs},year={2023},}