PaperlayerView exact revision on arXiv

Source-linked AI summary

LaRS: Latent Reasoning Skills for Chain-of-Thought Reasoning

Zifan Xu, Haozhu Wang, Dmitriy Bespalov, Xian Wu, Peter Stone, Yanjun Qi

arXiv:2312.04684v4cs.CLcs.AI

TL;DR

CoT prompts include rationales, so selecting examples by question similarity alone can miss the skills needed for a question; existing skill-description approaches can be costly and difficult to scale. LaRS learns latent rationale skills and predicts which skills a question requires to align examples. It consistently outperforms Skill-KNN and matches oracle performance in almost half of experiments.

  • Problem

    CoT demonstration selection must account for rationale skills, while skill-description approaches can rely on arbitrary, manually crafted prompts and require a description for each example.

    Source: §1 · p. 2

  • Method

    LaRS uses a CVAE to learn latent skill representations from rationales and a question-conditioned policy to predict the skills needed for demonstration selection.

    Source: §1 · p. 2, §4.1 · p. 5

  • Results

    Across four benchmarks and five backbones, LaRS consistently outperforms Skill-KNN and matches oracle performance in almost half the experiments.

    Source: §1 · p. 3

  • Takeaways & Limitations

    LaRS selects CoT examples through unsupervised rationale-skill discovery without human prompt design or extra LLM inference during selection.

    Source: §1 · p. 3, §6 · p. 9

    A single reasoning skill may not represent a rationale containing multiple reasoning steps.

    Source: §7 · p. 9

Abstract

from arXiv · show

Chain-of-thought (CoT) prompting is a popular in-context learning (ICL) approach for large language models (LLMs), especially when tackling complex reasoning tasks. Traditional ICL approaches construct prompts using examples that contain questions similar to the input question. However, CoT prompting, which includes crucial intermediate reasoning steps (rationales) within its examples, necessitates selecting examples based on these rationales rather than the questions themselves. Existing methods require human experts or pre-trained LLMs to describe the skill, a high-level abstraction of rationales, to guide the selection. These methods, however, are often costly and difficult to scale. Instead, this paper introduces a new approach named Latent Reasoning Skills (LaRS) that employs unsupervised learning to create a latent space representation of rationales, with a latent variable called a reasoning skill. Concurrently, LaRS learns a reasoning policy to determine the required reasoning skill for a given question. Then the ICL examples are selected by aligning the reasoning skills between past examples and the question. This approach is theoretically grounded and compute-efficient, eliminating the need for auxiliary LLM inference or manual prompt design. Empirical results demonstrate that LaRS consistently outperforms SOTA skill-based selection methods, processing example banks four times faster, reducing LLM inferences during the selection stage by half, and showing greater robustness to sub-optimal example banks.

1 Introduction

CoT demonstration selection must account for the skills expressed in rationales, but existing skill-description methods depend on manual prompts and scale poorly. LaRS learns latent rationale skills and a question-conditioned reasoning policy, and reports consistent gains over Skill-KNN.

1 Introduction

  • Selecting examples by question similarity can surface rationales with the wrong skill, whereas skill-based selection aligns examples with the operation a question requires.For example, a question needing addition may be paired with a subtraction rationale under question-similarity selection.Source: §1 · p. 2
  • Skill-KNN relies on LLM-generated skill descriptions that can be arbitrary, prompt-dependent, and difficult to scale because each example needs its own description.The passage also notes this reliance constrains applicability across diverse reasoning tasks.Source: §1 · p. 2
  • LaRS discovers latent representations of rationale skills without relying on LLM-generated skill descriptions.Its latent skill variable abstracts rationale features such as formats, equations, or knowledge.Source: §1 · p. 2
  • A CVAE trained on a small example-bank dataset learns both an encoder of demonstrated skills and a policy predicting the skills required by a question.The policy selects a skill from the question, while the rationale encoder maps a question-rationale pair to a skill.Source: §1 · pp. 2–3
  • Across four benchmarks and five backbones, LaRS consistently outperforms Skill-KNN and matches oracle performance in almost half the experiments.It also reduces LLM inference, avoids human prompt design, and is more robust to sub-optimal example banks.Source: §1 · p. 3
  • The paper presents LaRS as an unsupervised CoT demonstration selector, a latent-skill formulation of reasoning, and a theoretically analyzed selection method.The introduction describes its empirical verification as large-scale experiments.Source: §1 · p. 3

2 Related Work

Related work contrasts LaRS’s automatic discovery of rationale styles with manually designed prompting skills. It also situates LaRS among demonstration-selection heuristics and skill-based selection methods.

2.1 CoT Reasoning

  • LaRS automatically discovers diverse rationale styles and dynamically selects skills, unlike SKiC’s manually designed skills and full skill list in the prompt.The related-work discussion also notes PoT’s separation of textual reasoning from computation.Source: §2.1 · p. 3

2.2 Demonstration Selection

  • Demonstration selection chooses examples from an example bank; prior heuristics include similarity, diversity, coverage, and uncertainty.Skill-KNN is the closest related approach described, but depends on pre-trained LLM skill annotations.Source: §2.2 · p. 3
  • LaRS instead learns a lightweight CVAE to discover skills, avoiding extensive LLM inferences and human prompt design while grounding selection in the latent-skill model and theoretical analyses.Source: §2.2 · p. 3

3 Formulation

LaRS’s formulation treats rationale generation as a mixture over latent reasoning skills, with a question-conditioned policy selecting skills. It extends this view to CoT prompting and skill-based selection of examples from a bank.

3.1 Skill Model

  • The skill model represents rationale generation as a mixture over latent skills z, with a question-conditioned policy selecting z and a skill-conditioned distribution generating the rationale.The formulation allows the skill policy to vary with the question.Source: §3.1 · p. 4
  • Diverse rationale formats and reasoning paths motivate modeling human rationales as a stochastic mixture of skills rather than a single deterministic path.The paper notes that different paths and formats can lead to correct answers.Source: §3.1 · p. 4

3.2 CoT prompting

  • The formulation treats LLM rationale generation as an implicit skill model that selects a latent skill conditioned on the question before generating a rationale.This parallels the proposed model of real-world rationale generation.Source: §3.2 · p. 4, §3.1 · p. 4
  • CoT prompts are modeled as shaping the distribution over skills, covering both prefix-based zero-shot prompts and question-rationale few-shot prompts.The formulation simplifies prompts as influencing skill probabilities rather than the rationale-generation distribution directly.Source: §3.2 · p. 4
  • The paper attributes zero-shot generation failures to potentially sub-optimal skill policies and describes CoT prompts as improving reasoning by shaping skill distributions.The proposed explanation connects noisy real-world data and generic LLM training to sub-optimal skill selection.Source: §3.2 · p. 5

3.3 Skill-Based Demonstration Selection

  • The formulation grounds skill-based selection in a specially crafted example bank whose rationales may differ from real-world data and may showcase better reasoning skills.The selection method’s claimed optimality is conditioned on ideal assumptions about the bank and LLMs.Source: §3.3 · p. 5
  • Skill-based selection samples bank examples according to their skill-conditioned distribution and the test question’s skill policy, favoring skills likely for that question.The selection method is defined over question-rationale pairs and is analyzed for optimality under ideal assumptions.Source: §3.3 · p. 5

4 Method

LaRS approximates the example bank’s rationale distribution with a CVAE and selects demonstrations by matching inferred example skills to the skill predicted for a test question. The paper also proves an optimality result under stated assumptions.

4 Method

  • LaRS learns a CVAE approximation of the example-bank distribution to enable skill-based demonstration selection.The method overview presents this learned model as the basis for a practical selection process.Source: §4 · p. 5, §4.1 · p. 5

4.1 Latent Reasoning Skill Discovery

  • The CVAE couples an encoder for question-rationale pairs, a decoder for rationales conditioned on questions and skills, and a policy predicting skills from questions.The models are trained using variational expectation maximization and the reparameterization trick; the decoder is auxiliary and is not deployed downstream.Source: §4.1 · pp. 5–6
  • LaRS uses pretrained embeddings and feed-forward networks to represent the encoder, decoder, and reasoning policy instead of training full language models.The policy and encoder predict Gaussian latent-skill means and variances, while the decoder deterministically predicts an embedding.Source: §4.1 · p. 6

4.2 Demonstration Selection

  • For each test question, the policy predicts a skill, examples receive cosine-similarity scores against their inferred skills, and the top-k examples form the prompt.The test skill is the policy’s maximum-likelihood skill; each example’s skill is inferred from its question-rationale pair.Source: §4.2 · p. 6

4.3 Theoretical Analysis

  • The theorem states that skill-based selection yields the optimal conditional rationale distribution as k approaches infinity, provided Assumptions 1 and 2 hold.The result depends on the stated assumptions about the example bank and rationale generation.Source: §4.3 · pp. 6–7

5 Experiments

LaRS is evaluated on four benchmarks with five backbone models, against demonstration-selection baselines, with results on accuracy, computational overhead, and robustness to sub-optimal example banks.

5.1 Dataset

  • Evaluation uses TabMWP, GSM8K, Spider, and COGS, with 1K randomly sampled training examples as each example bank.Accuracy is measured by answer accuracy for TabMWP and GSM8K, execution-with-values accuracy for Spider, and exact match for COGS.Source: §5.1 · p. 7

5.2 Selection Methods

  • Baselines include Skill-KNN variants, random selection, question-embedding retrieval, and oracle retrieval using ground-truth rationales.Skill-KNN-zero is the primary comparison because it does not rely on human prompt design.Source: §5.2 · p. 7

5.3 Backbones and Hyper-parameters

  • The five backbones are gpt-4o, gpt-3.5-turbo, claude-3-sonnet, claude-3-haiku, and Falcon-40B-Instruct; prompts use dataset-specific shot counts.The shot counts are 2, 4, 4, and 8 for TabMWP, GSM8K, Spider, and COGS, respectively.Source: §5.3 · p. 7

5.4 Performance comparison results

  • LaRS outperforms Skill-KNN-zero in 18 of 20 experiments across four benchmarks and five backbones.The main-results table summarizes performance across backbones and datasets.Source: §5.4 · p. 7, Table 1 · p. 8
  • LaRS matches Skill-KNN-zero accuracy, uses no LLM inferences (about $30 saved per 1k examples), and reduces preprocessing time by 1.5 hours per 1k examples.The authors also report more than 100% less cost per input query.Source: §5.4 · p. 8, Table 3 · p. 8

LaRS is more robust to sub-optimal example

6 Conclusions

LaRS selects CoT demonstrations using reasoning skills learned as latent representations of rationales. Experiments report comparable selection performance while avoiding extra LLM inference and saving example-bank preprocessing time.

6 Conclusions

  • LaRS selects demonstrations for CoT prompting using reasoning skills learned as latent representations from rationales via a CVAE.The conclusion describes experiments across four LLMs and more than four reasoning tasks.Source: §6 · p. 9
  • LaRS achieves comparable performance in selecting effective few-shot examples for CoT reasoning across the evaluated tasks and LLMs.Source: §6 · p. 9
  • The selection process requires no extra LLM inference and saves hours in preprocessing the example bank.Source: §6 · p. 9

7 Limitations

The authors identify limitations in example ordering, the CVAE decoder, and representing multi-step rationales with one skill; the appendix outlines LaRS’s demonstration-selection procedure.

7 Limitations

  • A single reasoning skill may not represent an entire rationale containing multiple reasoning steps.The authors identify learning and selecting skills for individual reasoning steps as a direction to explore.Source: §7 · p. 9
  • Prompt example order is not considered, and the CVAE decoder is an MLP rather than a prompt-tuning module.The authors suggest that sorting heuristics could potentially improve performance and that prompt tuning would better align with the implicit skill model assumption.Source: §7 · p. 9

A LaRS Demonstration Selection

  • For a test question, LaRS predicts a desired skill, scores bank examples by cosine similarity to their inferred skills, and selects the top-k pairs.The appendix describes the reasoning policy as predicting the test question’s skill and the skill encoder as inferring each example’s skill from its question and rationale.Source: A LaRS Demonstration Selection · p. 12

B.1 Dataset

The appendix describes four evaluation datasets spanning mathematical reasoning, text-to-SQL, and semantic parsing, and includes a visualization of learned skills and embeddings.

B.1 Dataset

  • TabMWP provides semi-structured math problems using text and tables, while Spider is a large-scale text-to-SQL dataset.LaRS is trained on the TabMWP train set and evaluated on test1k; Spider uses its train set and dev set, respectively.Source: §B.1 · p. 12
  • COGS tests compositional generalization in semantic parsing, while GSM8K contains linguistically diverse grade-school math word problems.The COGS evaluation combines primitive substitution and primitive structural alternation; GSM8K evaluation uses its test set.Source: §B.1 · p. 12
  • Figure 5 compares skill projections from question-rationale pairs and questions alone with raw question and rationale embeddings, using 12 human-labeled skill colors.Source: Figure 5 · p. 13

B.2 LaRS Implementation Details

LaRS uses a shared model architecture and training setup across four datasets, while the appendix also records dataset-specific skill annotations and zero-shot prompts for Skill-KNN.

B.2 LaRS Implementation Details

  • LaRS uses an encoder, decoder, and reasoning policy with a 128-dimensional continuous reasoning-skill space and shared training hyperparameters across four datasets.The encoder and policy are MLPs with two 256-unit hidden layers; the decoder has two 256-unit hidden layers. Training uses batch size 256, learning rate 0.0001, and 1,000 epochs.Source: §B.2 · p. 13

B.3 Skill-KNN Implementation Details

  • For TabMWP and GSM8K, the authors manually annotated eight questions per dataset; COGS and Spider use the original Skill-KNN annotations.The new TabMWP and GSM8K annotations appear in Tables 4 and 5.Source: §B.3 · p. 13
  • Skill-KNN-zero uses zero-shot skill-description prompts tailored to each of the four datasets.The prompts are listed in Table 6.Source: §B.3 · p. 13

C Analysis and Ablation

The analyses examine whether LaRS's latent reasoning skills align with human skill labels and whether skill-guided selection remains effective across embedding models and example counts.

Why reasoning skill is a better guidance for demonstration selection?

  • Both LaRS's reasoning-skill encoder and reasoning policy separate 12 human-labeled skills, unlike raw question or rationale embeddings.The separation appears in t-SNE projections despite the absence of explicit skill labels during training.Source: Why reasoning skill is a better… · p. 14
  • LaRS consistently outperforms Retrieval-Q with matched embedding models and with 2, 4, or 8 in-context examples.Retrieval-based methods improve with more capable embeddings, and accuracy increases as the number of examples grows.Source: Why reasoning skill is a better… · p. 14, Figure 6 · p. 14

D Case Study

  • K-means clustering of latent reasoning skills from 1,000 TabMWP examples discerns examples showcasing specific skills, including searching minimum/maximum and computing rate change.Source: D Case Study · p. 14

E Theoretical Analysis

The theoretical analysis formalizes rationale generation under skill-based demonstration selection and proves an optimal conditional-rationale result under stated assumptions.

E Theoretical Analysis

  • The proof represents the prompt as selected question–rationale pairs followed by the query when integrating Equation 5.The prompt sequence is specified as (Q1, R1, …, Qk, Rk, Q).Source: E Theoretical Analysis · p. 14
  • The derivation forms a question-conditioned skill distribution from selected question–rationale pairs, then takes the limit as k → ∞.The resulting skill distribution is reintegrated into the rationale-generation expression.Source: E Theoretical Analysis · p. 15
  • Under Assumptions 1 and 2, skill-based CoT selection gives the optimal conditional distribution of rationales given questions.Assumption 1 states that the example bank can approximate expert rationale generation.Source: E Theoretical Analysis · p. 15
Loading 2312.04684v4…