💬 NLP

What Language Models Know But Don't Say: Non-Generative Prior Extraction for Generalization

The paper introduces LoID, a deterministic method that extracts informative prior distributions for Bayesian logistic regression by analyzing token-level token-level confidence in opposing semantic directions within Large Language Models, thereby significantly improving generalization on small, out-of-distribution tabular datasets compared to existing text-generation-based approaches.

Original authors: Sara Rezaeimanesh, Mohammad M. Ghassemi

Published 2026-04-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sara Rezaeimanesh, Mohammad M. Ghassemi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Small Class" Dilemma

Imagine you are trying to teach a student how to predict the weather.

The Ideal Scenario: You give the student 10 years of weather data from every continent. They learn the patterns perfectly.
The Real-World Problem: In fields like medicine or finance, you often only have data from one small town or one specific hospital. Maybe your data only includes patients who are young and healthy.

If you train a model on this "small class" data, it becomes a genius at predicting weather for that one town but fails miserably when it tries to predict for a different city (like a cold, snowy city). This is called the Out-of-Distribution (OOD) problem. The model is "brittle" because it hasn't seen enough variety.

The Solution: The "Library of Wisdom"

Enter Large Language Models (LLMs) like the one you are talking to right now. These models have read almost everything on the internet, including millions of medical textbooks, finance reports, and scientific papers. They possess a massive "Library of Wisdom" about how the world works.

The Old Way: Researchers tried to ask the LLM, "Hey, what do you think about this patient?" and hoped the LLM would write a paragraph giving a prediction. But LLMs are chatty and inconsistent; sometimes they say one thing, sometimes another, just because of how the question was phrased.
The New Way (LoID): The authors of this paper realized they didn't need the LLM to write an answer. They just needed to peek at what the LLM was thinking before it spoke.

The Method: LoID (Logit-Informed Distributions)

Think of the LLM as a very smart, but shy, librarian. You don't want to ask the librarian to write a whole essay about "Does smoking cause cancer?" because they might ramble. Instead, you want to know: "How confident are you that smoking causes cancer?"

The authors built a tool called LoID (Logit-Informed Distributions) to do exactly this. Here is how it works, step-by-step:

1. The "Yes vs. No" Test

For every feature in a dataset (e.g., "Age," "Smoking Status," "Income"), the researchers ask the LLM a simple question in two different ways:

"Does Age have a positive impact on Stroke Risk?"
"Does Age have a negative impact on Stroke Risk?"

2. Reading the "Thoughts" (Logits)

Instead of waiting for the LLM to type out "Yes, age increases risk," LoID looks at the internal math the LLM uses to decide which word to pick next.

It checks the probability the LLM assigns to the word "Positive" vs. the word "Negative."
If the LLM is 99% sure the word should be "Positive," it gives a high score.
If the LLM is confused (50/50), it gives a low score.

3. Creating a "Mental Map" (The Prior)

The researchers take these scores and turn them into a Prior.

Analogy: Imagine you are baking a cake, but you've never baked one before. You have a recipe (your data), but it's missing some ingredients.
The Prior: This is like a "Grandma's Secret Note" that says, "Usually, you need 2 cups of flour."
LoID's Prior: It says, "Based on everything I've read in the world, Age usually increases Stroke Risk."

This "Grandma's Note" helps the model when the data is scarce or weird. It acts as a safety net, guiding the model toward what is generally true in the real world, even if the specific data it's looking at is biased.

Why is this better than just asking the LLM?

Consistency: If you ask an LLM to write a story, it might change its mind every time. But if you check its internal math (logits), the answer is deterministic. It's like checking a thermometer vs. asking a person "Is it hot?" The thermometer gives a consistent number every time.
No Hallucinations: The LLM isn't making up facts; it's revealing the statistical relationships it learned during its training.

The Results: When it Works and When it Fails

The authors tested this on 15 different datasets (medical, financial, engineering).

✅ When LoID Shines:

Scenario: Predicting Concrete Strength.
Why: Physics is universal. Cement + Water + Time = Strong Concrete. This is true everywhere.
Result: LoID used the LLM's knowledge of physics to fix the model, recovering 94% of the lost performance. The LLM knew the rules of physics even if the data was from a weird factory.

❌ When LoID Stumbles:

Scenario: Predicting Bank Marketing (will a customer buy a term deposit?).
Why: This depends on specific human behavior, local campaigns, and timing. A bank in Portugal acts differently than a bank in the US.
Result: The LLM's "general knowledge" was too vague or wrong for this specific context. It tried to apply a generic rule to a specific situation, making things worse.

The Takeaway

This paper introduces a way to harvest the "common sense" of AI without relying on its "creative writing."

Old Approach: Ask the AI to guess the answer (unreliable).
New Approach (LoID): Ask the AI to show its confidence in the relationship between variables, then use that confidence to guide a standard statistical model.

It's like having a world-class expert (the LLM) stand next to a junior data scientist (the model). Instead of letting the expert take over the whole job, you just let the expert whisper, "Hey, usually, older people are at higher risk," and the junior scientist uses that whisper to correct their calculations when the data is missing or biased.

In short: LoID lets us use the vast knowledge of AI to make small, imperfect datasets behave like big, perfect ones, but only when the rules of the world are consistent enough for the AI to know them.

1. Problem Statement

In critical domains like medicine and finance, machine learning models often suffer from Out-of-Distribution (OOD) generalization failures. This occurs because training data is frequently small, non-representative, or biased (e.g., data from a single hospital or specific demographic), leading to models that fail when deployed on real-world populations with different distributions.

While Bayesian methods offer a principled way to encode prior knowledge to mitigate this brittleness, manually eliciting high-quality priors at scale is difficult. Recent attempts to use Large Language Models (LLMs) as prior sources (e.g., AutoElicit) rely on generative sampling (asking the LLM to write down priors), which introduces:

Run-to-run variability due to stochastic text generation.
Inefficiency in extracting precise numerical distributions.
Inconsistency in capturing the model's internal confidence.

The paper addresses the gap of how to extract deterministic, informative priors directly from the internal structure of LLMs without relying on generated text.

2. Methodology: LoID (Logit-Informed Distributions)

The authors propose LoID, a deterministic framework that extracts Bayesian priors for logistic and linear regression directly from an LLM's token-level logits (probabilities) rather than generated completions.

Core Mechanism

Semantic Probing: For each feature $f_j$ in a dataset, the method constructs paired prompts expressing opposing semantic claims about the feature's impact on the target variable (e.g., "The impact of $f_j$ on $t$ is positive" vs. "...is negative").
Logit Extraction: Instead of asking the LLM to generate a number, LoID queries the model's next-token probabilities for specific sentiment tokens ("positive" vs. "negative") given the prompt context.
Logit Preference Score: The method calculates a preference score based on the log-odds of the positive token versus the negative token:
$\text{logit}_j(p_i) = \log\left(\frac{p_i}{1 - p_i}\right), \quad \text{where } p_i = \frac{P(\text{positive})}{P(\text{positive}) + P(\text{negative})}$
This maps the bounded probability $(0, 1)$ to the unbounded real line, suitable for regression coefficients.
Prior Construction:
- Mean ( $\mu_j$ ): The average of logit scores across multiple paraphrased prompts ( $T$ ) represents the LLM's belief about the direction and strength of the relationship.
- Variance ( $\sigma_j^2$ ): The standard deviation of the logit scores across paraphrases serves as an uncertainty estimate. High variance indicates the LLM is inconsistent about the feature's effect, leading to a wider prior.
Bayesian Integration: These parameters form a Gaussian prior $\beta_j \sim \mathcal{N}(\mu_j, \sigma_j^2)$ for the coefficients in Bayesian logistic/linear regression. Posterior inference is performed using a Laplace approximation (Gaussian approximation around the MAP estimate) for efficiency.

3. Key Contributions

Deterministic Prior Elicitation: LoID is the first method to extract Bayesian priors directly from token-level logits, eliminating the stochasticity associated with text generation and prompt sampling found in methods like AutoElicit.
No Fine-Tuning Required: The method operates in a zero-shot setting, requiring no fine-tuning or in-context learning examples.
Comprehensive Evaluation: The framework is evaluated on 15 real-world tabular datasets (14 public, 1 private medical dataset) under synthetic OOD settings characterized by covariate shifts.
Baseline Comparisons: The study compares LoID against:
- Standard uninformative priors ( $\mathcal{N}(0,1)$ ).
- AutoElicit (LLM-generated expert priors).
- LLMProcesses (LLM as a direct predictor).
- An Oracle upper bound (training on full IID data).

4. Experimental Results

The experiments demonstrate that LoID significantly improves OOD generalization in specific contexts:

Performance Gains: Across 15 datasets, LoID improved performance in 11/15 cases, recovering on average 50% of the performance gap between the OOD-trained model and the Oracle (IID-trained) model.
Domain Success:
- Medical & Demographic: LoID excelled in tasks with stable, generalizable causal mechanisms (e.g., Adult Income, Chronic Kidney Disease, Head & Neck Cancer), closing up to 83.8% of the performance gap.
- Physical Systems: It performed well on regression tasks with known physical laws (e.g., Concrete Strength, closing 93.9% of the gap).
Failure Modes:
- Context-Dependent Tasks: LoID struggled or degraded performance on tasks driven by local, behavioral, or campaign-specific patterns (e.g., Bank Marketing, Bike Sharing). In these cases, the LLM's generic world knowledge conflicted with dataset-specific realities, introducing harmful bias.
- Comparison: LoID outperformed LLMProcesses (direct prediction) on 12/14 datasets and AutoElicit on 12/14 datasets.
Efficiency: LoID is approximately 10x faster than AutoElicit because it avoids iterative multi-round sampling and relies on a single pass of logit extraction.

5. Significance and Conclusion

The paper establishes that LLMs possess valuable, structured domain knowledge that can be extracted deterministically to improve statistical models under distribution shift.

Shift in Paradigm: It moves beyond using LLMs as "black box" predictors or text generators, treating them instead as knowledge bases whose internal confidence (logits) can be quantified to regularize classical statistical models.
Practical Utility: LoID offers a robust solution for low-data regimes in critical fields (medicine, finance) where data is scarce or biased, provided the underlying feature-target relationships are semantically stable and generalizable.
Limitations: The method is less effective for datasets with de-identified features (where semantic meaning is lost) or highly context-dependent behavioral data where the LLM's generic knowledge does not align with specific local dynamics.

In summary, LoID demonstrates that extracting distributional priors from LLM logits is a more robust and reliable strategy for OOD generalization than direct prediction or free-form generative elicitation.