Selection-Inference: Exploiting Large Language Models for Interpretable Logical Reasoning

In plain words

Language models often get logic puzzles wrong when the answer needs several linked deductions. This paper breaks the work into small steps. In each step, the model first picks out the few facts from the text that matter for one deduction, and then, seeing only those facts and not the question, draws one new conclusion from them. That conclusion is added to the pile of known facts, and the loop repeats until the answer comes out.

The point of the split is that the second step cannot cheat. It never sees the question, so it cannot skip the reasoning and guess, and it cannot make up facts that are not in the text. The result is a chain of short, plain-language steps that a person can read, check, and correct. The method needs no training: it works through example prompts given to a pre-trained model.

TL;DR

• The method wraps a frozen LLM in a loop of selection (choose facts from the context) and inference (derive one fact from them), and appends each new fact to the context §4.
• It is tested on 10 logic tasks from bAbI, ProofWriter, and related sources, using a 7B Gopher model in a 5-shot setting §5.
• Author-reported: the 7B model in SI reaches 58.75% average generative accuracy, against 2.94% for the same model used naively and 41.32% with chain-of-thought prompting §5, fig. 4a.
• Fine-tuning the selection and inference steps separately raises accuracy further, to 78.95% §6, fig. 4b.
• Relevant if you want traceable multi-step reasoning from a model without training, and can accept several model calls per step.

Problem

Vanilla prompting of LLMs, where the model gives the answer directly, performs well on single-step entailment but worse on multi-step logic. Performance falls as the number of steps grows and as irrelevant facts are added §3, fig. 3b. Chain-of-thought prompting produces reasoning text, but the answer does not causally depend on it, so the reasoning can be wrong even when the answer is right §1. The paper assumes that each question comes with a context that holds all needed information plus possible distractors, and that questions are well posed and answerable from that context. Retrieval of facts is left to future work §4.

Method

Overview. Each reasoning step has two modules that both use the same pre-trained LLM with a different few-shot prompt §4, alg. 1:

1. Selection picks a subset of the current context, given the context and the question §4.1.
2. Inference sees only the selected facts (no question) and writes one new fact §4.2.

The new fact is appended to the context, and the next step runs on the enlarged context §4. The paper runs a fixed number of steps; a learned or rule-based stop is left for future work §4, §7. The final inference is taken as the answer §4, app. D.

Selection module §4.1.

1. The few-shot selection prompt shows examples of a context, a question, and a selection. Each selection has the form "fact. We know that fact [and fact]*. Therefore," §4.1.
2. The paper first tried free generation from the prompt. The authors found that this often made up facts, which breaks the causal link §4.1.
3. Instead, the LLM scores each candidate fact or rule in the context, and the best-scoring one is added to the selection. The next candidate is scored with that selection in the prompt, until the fixed number of facts is reached §4.1, alg. 2.
4. For ProofWriter, an extra choice ". Therefore," can be scored, and it ends the selection. At most two facts are selected app. D.

Inference module §4.2.

1. The few-shot inference prompt shows examples of the form "fact. We know that fact [and fact]*. Therefore, new fact." §4.2, app. C.
2. The selection is appended after the examples, and the LLM generates text.
3. The new fact is the first sentence of the generated text, extracted with a rule that stops at ". ? ! ;" or a newline §4.2, app. D. The paper follows the BigBench evaluation pipeline for this step §4.2.

Fine-tuned variant §6, app. F. The paper also trains the two modules separately, on single steps only.

• The selection model is trained to output sentence labels such as "sent 2. We know that sent 4 and sent 7." instead of sentence text. The authors say this stops it from inventing facts §6.
• The inference model is trained to compute an entailment from the selection §6.
• Training data for each step is built by appending ground-truth intermediate inferences to the context fig. S2.
• At selection time, labels can be attached to each context element and read back through a dictionary, which is faster than scoring each candidate app. F.

# Selection-Inference loop  [alg. 1] [§4]
# Input:  context (list of facts and rules), question, LLM
# Output: final inference, taken as the answer  [app. D]

def selection_inference(context, question, llm, steps):
    # Repeat for a fixed number of steps; no learned halting  [§4] [§7]
    for step in range(steps):
        # Choose a subset of the current context  [§4.1] [alg. 2]
        selection = selection_module(context, question, llm)
        # Derive one new fact from the selection, without the question  [§4.2]
        inference = inference_module(selection, llm)
        # Add the new fact to the context for the next step  [§4] [alg. 1]
        context = context + [inference]
    return inference

# Scoring selection  [alg. 2] [§4.1]
# Input:  selection prompt, context, question, LLM
# Output: selection string

def selection_module(prompt, context, question, llm):
    # Start with an empty selection  [alg. 2]
    selection = ""
    # Add one fact per pick, up to the fixed number of facts  [§4.1] [app. D]
    for pick in range(max_facts):
        # Score each candidate by its summed token log-likelihood  [alg. 2]
        scores = [
            sum_log_likelihood(llm, candidate, prompt, context, question, selection)
            for candidate in context
        ]
        # Append the best-scoring candidate to the selection  [alg. 2]
        selection = join(selection, context[argmax(scores)])
    return selection

# Inference module  [§4.2] [app. D]
# Input:  selection string, inference prompt, LLM
# Output: one new fact (string)

def inference_module(selection, prompt, llm):
    # Generate after the few-shot examples and the selection  [§4.2]
    generated = llm.generate(prompt + selection + ". Therefore,")
    # Keep only the first sentence (stop at . ? ! ; or newline)  [app. D]
    new_fact = first_sentence(generated)
    return new_fact

Answer extraction. For bAbI, the answer is the last word of the final inference, because the model writes a full sentence app. D. For ProofWriter, the statement is labelled True if one of the derived implications matches it, following the ProofWriter procedure app. D.

Hyperparameters

Results

Author-reported. Setup: 5-shot, Gopher 7B (SI and baselines) and 280B, 10 tasks: bAbI 1, 2, 3, 15, 16 and ProofWriter depths 0, 1, 2, 3, 5 §5. Generative accuracy is exact match of the first sentence; multi-choice scores each option §5. SI is evaluated only in the generative setting §5.

The paper also reports effects of extra steps on ProofWriter (larger gains at depth 1 than depth 2) app. D, fig. S1.

Paper vs code

No official code. Searched: the paper and arXiv metadata for code links (none); GitHub repository search for "selection-inference", "selection-inference" within the google-deepmind organisation (no results), "selection inference llm", and "interpretable logical reasoning selection" (no repository by the authors found). No discrepancies to report.

Pitfalls

• Return value conflict. Algorithm 1 returns the last selection s^t alg. 1, but the text takes the final inference as the answer §4 app. D. Implement the text, and check the returned value explicitly.
• Halting is undefined. Algorithm 1 uses a halt() function, but the paper uses a fixed number of steps and says halting is future work §4 §7. The number of steps per task is given only in app. D.
• Scoring formula is ambiguous. Algorithm 2 writes each token's probability as conditioned on the same prefix (prompt, context, question, selection), without the candidate's earlier tokens alg. 2. The text says "highest log-likelihood" §4.1. Implement the usual sequence log-likelihood of the candidate given the prompt, and say so in your code. The exact prompt layout for scoring is not stated.
• No length normalisation. Summed log-likelihood favours shorter candidates. The paper uses unnormalised scores for multiple-choice tasks app. H and does not say whether this applies to selection too.
• Free generation of selections fabricated facts. The authors found this in their experiments, which is why they score candidates instead §4.1. Restricting the selection to context items is what keeps the trace causal §4.1.
• ProofWriter needs a stop option. Without ". Therefore," as a choice, the number of selected facts is hard to control. The paper caps it at two app. D.
• Answer labelling depends on dataset. ProofWriter uses only the open-world (OWA) data with Unknown answers removed, so True and False are the only labels app. D. Labelling uses the ProofWriter matching procedure app. D.
• Step budget affects accuracy. Extra steps can help a model recover from an early error §5, fig. 2, app. D. Reported results use the minimum number of steps; extra steps may improve accuracy app. D, fig. S1.
• Intro ratios do not match §5. The intro says SI "almost triples" the naive 7B and "almost doubles" the COT 7B §1. The §5 averages give about 20× (58.75 / 2.94) and about 1.4× (58.75 / 41.32). The abstract's "over 100%" improvement is also not derivable from the §5 averages.
• Task count mismatch. The abstract and §5 say 10 tasks. The caption of fig. 4a says 11 datasets fig. 4a. The paper does not explain the difference.
• Two different prompt-only SI numbers. §5 gives 58.75% and §6 gives 57.93% for prompt-engineered SI. §6 also gives different 7B vanilla and COT figures (0.34%, 15.73%) from §5 (2.94%, 41.32%). The task subset for §6 is not stated.
• Cross-setting comparisons. The text compares 7B SI (generative) with 280B baselines in multi-choice §5. The authors note the 280B result is surprising and may reflect priors (footnote in §5). Do not compare settings directly.
• Prompt examples are unspecified. The exact choice of few-shot examples and their order is not stated. Results may depend on them.
• Cost (from the algorithm structure, not stated in the paper). Each selected fact costs one scoring pass per context candidate, and each step adds one generation. Cost grows with context size, facts per selection and steps.
• Fine-tuning covers one step. Both fine-tuned models are trained on single steps only §6. The fine-tuned selection still drops at deeper problems, because planning ahead is hard §6.
• Limitations the authors name. The selection module is the main failure point; prompt engineering works poorly for multi-step selection; the method assumes the context is given; and there is no halting or filtering rule app. G.

Lineage

• Builds on: chain-of-thought prompting (Wei et al. 2022) §2, §5, the stepwise prover ProofWriter (Tafjord et al. 2020) §2, app. D, and the Gopher model family (Rae et al. 2021) §3.
• Compared against: vanilla prompting and chain-of-thought baselines with the 7B and 280B Gopher models §5.
• Superseded by: not stated in the paper.