Partition, Prompt, Aggregate: Statistical Self-Consistency in Language Models

Core Methodology§

The paper evaluates whether LLMs satisfy the law of total probability under in-context learning. Given a population with a target variable $Y$ and a discrete partition $\{A_i\}$ of the covariate space, the law requires:

$$ P(Y) = \sum_i P(Y \mid A_i) P(A_i) $$

To test this, the authors construct binary trees that recursively partition a population into increasingly fine-grained subpopulations. Each node corresponds to a subpopulation described verbally (e.g., "people aged 30-40, income > $50k"). They prompt an LLM to estimate the conditional probability $P(Y \mid \text{subpopulation})$ and then aggregate estimates bottom-up: the estimate at a parent node is the weighted average of its children's estimates, where weights are the relative sizes of subpopulations (assumed known or estimated from data).

Binary Tree Protocol§

1. Define a root node representing the full population. 2. Recursively split nodes by a binary attribute (e.g., age > 40, income high/low). 3. For each node, prompt the LLM with the textual description of the subpopulation to obtain an estimate $\hat{p}_\text{node}$. 4. Compute self-consistency discrepancy as the absolute difference between the direct LLM estimate at a node and the aggregated estimate from its children: $\Delta = |\hat{p}_\text{node} - \sum_{\text{child}} w_{\text{child}} \hat{p}_{\text{child}}|$.

Key Findings§

• Widespread violations: Across multiple LLMs (GPT-4, Llama-3, etc.) and tasks (demographic surveys, medical risk predictions), self-consistency discrepancies are large, often exceeding 10%.
• Macro fallacy: When comparing direct population-level estimates against fine-grained aggregated estimates, the latter often align better with ground truth. This suggests models possess accurate subpopulation knowledge but fail to integrate it into aggregate answers.
• Partial recovery via implicit prompting: Using persona-prompting (e.g., "Imagine you are a 30-year-old high-income individual") reduces but does not eliminate the inconsistency.

Implementation Snippet§

Below is a Python-like pseudocode illustrating the aggregation procedure:

def binary_tree_aggregate(tree, llm_predict_fn):
    # tree: each node has textual description, children, and weight
    if node.is_leaf():
        return llm_predict_fn(node.description)
    else:
        child_estimates = [binary_tree_aggregate(child) for child in node.children]
        aggregated = sum(w * est for w, est in zip(node.weights, child_estimates))
        return aggregated

# Compute consistency metric
for root in trees:
    direct = llm_predict_fn(root.description)
    aggregated = binary_tree_aggregate(root)
    consistency_gap = abs(direct - aggregated)

Implications§

The paper introduces statistical self-consistency as a reference-free evaluation metric. It reveals a fundamental gap in how LLMs reason about conditional probabilities, suggesting that current models do not perform proper Bayesian conditioning. This has practical implications for any application requiring consistent probabilistic estimates, such as survey simulation or decision support.