Crafting Interpretable Embeddings by Asking LLMs Questions (2405.16714v1)

Abstract: LLMs have rapidly improved text embeddings for a growing array of natural-language processing tasks. However, their opaqueness and proliferation into scientific domains such as neuroscience have created a growing need for interpretability. Here, we ask whether we can obtain interpretable embeddings through LLM prompting. We introduce question-answering embeddings (QA-Emb), embeddings where each feature represents an answer to a yes/no question asked to an LLM. Training QA-Emb reduces to selecting a set of underlying questions rather than learning model weights. We use QA-Emb to flexibly generate interpretable models for predicting fMRI voxel responses to language stimuli. QA-Emb significantly outperforms an established interpretable baseline, and does so while requiring very few questions. This paves the way towards building flexible feature spaces that can concretize and evaluate our understanding of semantic brain representations. We additionally find that QA-Emb can be effectively approximated with an efficient model, and we explore broader applications in simple NLP tasks.

Crafting Interpretable Embeddings by Asking LLMs Questions

The paper presents a novel technique to generate interpretable text embeddings, designated as Question-Answering Embeddings (QA-Emb). Authors Vinamra Benara, Chandan Singh, John X. Morris, and Richard Antonello spearhead the research from prominent institutions like UC Berkeley and Microsoft Research, focusing on the intersection of machine learning, NLP, and neuroscience.

Problem Statement

Traditional methods for generating text embeddings, such as bag-of-words or transformer-based embeddings (e.g., BERT, LLaMA), often produce opaque representations, complicating their interpretability. This opaqueness poses significant challenges in domains that demand trustworthy interpretation, such as neuroscience. The authors propose QA-Emb to bridge this gap by rendering each dimension of the embedding human-interpretable through a series of yes/no questions administered to a pre-trained autoregressive LLM.

Methodology

QA-Emb involves querying an LLM with a set of yes/no questions related to the input text. Each question’s binary answer (mapped to 0 or 1) forms a specific dimension of the resulting embedding. Notably, this method does not require fine-tuning the LLM or altering its internal parameters but rather relies on carefully crafted natural language prompts.

Learning the Set of Questions

The selection of yes/no questions is optimized to suit the downstream task. In the case of predicting fMRI responses, the authors formulate the learning problem as an optimization task for ridge regression. The questions are heuristically generated via prompts to capable LLMs like GPT-4 and are fine-tuned by methods such as Elastic Net for redundancy reduction.

Neuroscience Application

Focusing on a neuroscience application, the authors employ QA-Emb to predict human brain responses (measured through fMRI) to natural language stimuli. The paper uses data from narrative podcast stories heard by subjects, with the embedding inputs used in ridge regression models to predict fMRI responses. The results showcase a 26% improvement over the existing interpretable baseline (Eng1000) and competitive performance compared to black-box models like BERT and LLaMA.

Numerical Results

Key findings include:

• QA-Emb outperforms Eng1000 by 26% in terms of average test correlation.
• Even with only 29 questions, QA-Emb achieves superior interpretability and performance compared to Eng1000 which used a larger set of features.
• QA-Emb achieves a 0.116 average test correlation, slightly better than BERT but 7% lower than the best-performing LLaMA model.

Limitations and Optimizations

Two primary limitations cited are the high computational cost and potential inaccuracies in the LLM's answers to the yes/no questions:

1. Computational Cost: QA-Emb requires numerous LLM calls, rendering it computationally intensive. To alleviate this, the authors explore model distillation, whereby a RoBERTa model predicts multiple questions' answers in a single feedforward pass, yielding nearly equivalent performance with significantly reduced computational overhead.
2. LLM Accuracy: The reliability of QA-Emb depends on the LLM’s ability to faithfully answer the yes/no questions. Variability in LLM performance on diverse binary classification tasks underscores the necessity for strong LLMs and optimized prompt engineering.

Broader Applications and Future Work

QA-Emb demonstrates potential applications beyond neuroscience, including information retrieval and text clustering, where it provides modest improvements and a high degree of interpretability. The paper outlines several avenues for future research:

• Enhanced optimization techniques for selecting questions.
• A broader range of applications in domains requiring interpretable text embeddings.
• Improved discrete optimization methods and constraints for more direct optimization of QA-Emb.

Moreover, the authors highlight the societal benefits of interpretable AI systems and the importance of transparency in AI applications, especially in high-stakes fields such as medicine and social sciences.

Conclusion

In summary, QA-Emb introduces a promising method for generating interpretable text embeddings by leveraging the capabilities of LLMs through strategic questioning. This innovation aligns high interpretability with robust performance, addressing a significant challenge in embedding techniques and opening new pathways for applications in various domains. As LLMs evolve, QA-Emb stands to benefit from increased efficiency and capability, further cementing its utility in NLP and beyond.