GraphIC: A Graph-Based In-Context Example Retrieval Model for Multi-Step Reasoning

## Abstract

Abstract In-context learning (ICL) enables large language models (LLMs) to generalize to new tasks by incorporating a few in-context examples (ICEs) directly in the input, without updating parameters. However, the effectiveness of ICL heavily relies on the selection of ICEs, and conventional text-based embedding methods are often inadequate for tasks that require multi-step reasoning, such as mathematical and logical problem solving. This is due to the bias introduced by shallow semantic similarities that fail to capture the deeper reasoning structures required for these tasks.
We present GraphIC, a novel approach that leverages graph-based representations of reasoning processes, coupled with Bayesian Networks (BNs) to select ICEs. Graph structures inherently filter out shallow semantics while preserving the core reasoning structure. Importantly, BNs capture the dependency of a node’s attributes on its parent nodes, closely mirroring the hierarchical nature of human cognition—where each thought is shaped by preceding ones. This makes BNs particularly well-suited for multi-step reasoning tasks, aligning the process more closely with human-like reasoning.
Extensive experiments across three types of reasoning tasks (mathematical reasoning, code generation, and logical reasoning) demonstrate that GraphIC outperforms both training-free and training-based models in selecting ICEs, excelling in terms of both effectiveness and efficiency. We show that GraphIC enhances ICL’s performance and interpretability, significantly advancing ICE selection for multi-step reasoning tasks.

## 1 Introduction

In-context learning (ICL) (Brown et al., ) represents a paradigm in how large language models (LLMs) perform inference by using a small number of in-context examples (ICEs) within the input prompt. This technique enables LLMs to generalize to new tasks or enhance their performance on existing tasks without updating parameters. However, previous studies have highlighted the sensitivity of ICL performance to the specific ICEs selected (Zhao et al., ; Liu et al., ), underscoring the importance of strategic ICE selection. Consequently, numerous methods have been proposed to optimize the selection of ICEs, focusing on improving task performance and ensuring greater robustness (Liu et al., ; Rubin et al., ; Ye et al., ; Gupta et al., ). These methods frequently rely on text-based embeddings, where both the query and candidate ICEs are embedded using a language encoder, with similarity scores guiding the selection process.

Figure: Figure 1: ICL with different ICE retrieval mechanisms. The left panel shows examples retrieved via BERT embedding (Devlin et al., ), while the right panel displays examples retrieved via GraphIC. Semantically related terms are highlighted in blue, and quantities or variables needing resolution are highlighted in green.
Refer to caption: x1.png

Current text-based embedding selection methods primarily focus on capturing semantic-level similarity, demonstrating their utility in tasks such as sentiment analysis (Liu et al., ) and machine translation (Agrawal et al., ). However, these approaches encounter significant limitations in multi-step mathematical and logical reasoning tasks, such as GSM8K (Cobbe et al., ) and ProofWriter (Tafjord et al., ). The core issue lies in the fact that textual data often encodes a substantial amount of shallow semantic information, which is largely irrelevant to the underlying reasoning processes required for math and logic tasks. This extraneous information introduces bias in the selection of ICEs (An et al., ), and can even lead the LLM to adopt misleading reasoning strategies, thereby degrading task performance. For example, in a problem involving speed calculation (i.e., determining the rate of change of distance over time), text-based embeddings may prioritize examples that focus on solving for time or distance due to their closer semantic similarity, as shown in Figure 1 (left). This misalignment steers the LLM away from the correct problem-solving approach. Therefore, a novel representation is needed to enhance example retrieval for tasks involving multi-step mathematical and logical reasoning.

Extensive research indicates that graph-based representations, in contrast to text, align more closely with human cognitive processes and are better equipped to model multi-step reasoning (Friston, ; Besta et al., ; Yao et al., ). Graphs enable the filtering of irrelevant shallow semantic content while preserving the core reasoning structure, thus facilitating the development of an unbiased example retrieval mechanism. More critically, this representation offers a novel and interpretable approach for constructing example retrievers. When solving multi-step reasoning problems, the ideal examples are those whose reasoning processes can be directly applied to the query. In other words, after extracting the thought pattern from these examples, the same logic can be reused to solve the query question and arrive at the correct answer. Graph structures provide an explicit means to model these reasoning processes, enabling us to extract implicit thought patterns and assess their transferability to new problems.

More formally, a reasoning process can be modeled as a graph $G$ with attributes $(x_{1},\dots,x_{n})$, where each $x_{i}$ represents the attribute of vertex $v_{i}$, corresponding to the thought at $v_{i}$. We further assume that the probability density function governing the sequence of thoughts throughout the reasoning process is given by $p(x_{1},\dots,x_{n};G,W)$, where $W$ denotes the parameters of the underlying thought pattern. Given a query question $q$ and its associated reasoning graph $G^{q}$, our goal is to retrieve ICEs with parameterized thought patterns $W^{i}$ that maximize the probability density $p(x_{1}^{q},\dots,x_{n}^{q};G^{q},W^{i})$, i.e., maximizing the likelihood of solving $q$ correctly. To achieve this, we employ a Bayesian Network (BN) (Pearl, ), a type of probabilistic graphical model, to represent the multi-step reasoning process and to parameterize the thought patterns. The motivation for using BN lies in their inherent similarity to human cognition; they assume that the value of each node (thought) is influenced by the values of preceding nodes, mirroring the way new ideas are constructed based on prior knowledge. This structure makes BNs particularly well-suited for capturing the dependencies and progression of reasoning steps.

In this paper, we introduce GraphIC, a Graph-based In-Context Example Retrieval Model. The key idea is to leverage graph representations of reasoning processes to enhance the retrieval of ICEs. Specifically, we first prompt the LLM to generate “thought graphs” for both the candidate examples and the query question, where each graph encodes the reasoning steps involved. We then employ a BN to model each thought graph of candidate examples and estimate the associated parameters through maximum likelihood estimation, which we reformulate as a bilinear form maximization problem with a closed-form solution. To better simulate human-like reasoning, we incorporate a personalized PageRank (PPR) mechanism (Page, ), reflecting the cognitive tendency to revisit starting point during reasoning, a characteristic well-aligned with the assumptions behind PPR. Once the thought graphs and their parameters are computed, we select examples whose parameters maximize the probability density of thoughts on the query question’s thought graph. Note that we leverage a graph model to aid ICE retrieval, not for direct problem-solving. The solution is driven by the LLM’s reasoning capabilities. This approach ensures that the retrieved examples, though not necessarily semantically similar to the query, align with the underlying reasoning process. As shown in Figure 1(right), GraphIC retrieves examples that, while not focusing on semantically similar tasks, share a common reasoning structure—for instance, solving for rates such as how the number of envelopes changes over time or how costs vary with the number of toys. This guides the LLM towards the correct solution.

We evaluate GraphIC against 10 baseline models on three multi-step reasoning tasks: mathematical reasoning, code generation, and logical reasoning.
Results show GraphIC excels in selecting relevant ICEs, surpassing existing methods in effectiveness and efficiency. Additionally, GraphIC shows faster performance improvements with more ICEs and exhibits asymmetry, mirroring real-world reasoning. Key contributions are summarized as follows:
1) We introduce a novel graph-based representation, the “thought graph”, to model multi-step reasoning processes. This approach effectively filters out irrelevant shallow semantic information while preserving the essential reasoning steps.
2) By leveraging BNs, we design a retrieval mechanism that maximizes the probability density of solving the query problem, providing a more objective-driven and interpretable retrieval process.
3) Our experimental results indicate that GraphIC, despite being a training-free model, outperforms both training-free and training-based models across various multi-step reasoning tasks.

## 2 Related Work

Existing ICE selection techniques can be classified as either training-free or training-based, depending on whether a retriever needs to be trained.

Training-free approaches are generally divided into two types: (i) those that use heuristic criteria such as similarity (Liu et al., ; Hu et al., ), diversity (Cho et al., ; Zhang et al., ; Levy et al., ; Hongjin et al., ; Zhang et al., ), complexity (Fu et al., ), or combinations of these (Agrawal et al., ; Tonglet et al., ; Gupta et al., ) to select in-context examples (ICEs); (ii) those that leverage feedback from LLMs, such as probability distributions (Wu et al., ; Nguyen & Wong, ; Li & Qiu, ; Yang et al., ), perplexity (Gonen et al., ), or the model’s generated output (An et al., ) to guide the selection process. While training-free approaches avoid the computational and time overhead associated with model training, their relatively simplistic architecture often results in sub-optimal performance compared to training-based methods.

Training-based methods are typically divided into two main categories. The first learns to select individual examples and then extends this to $k$-shot scenarios (Rubin et al., ; Xiong et al., ; Gupta et al., ). The second models the selection of a group of examples as a whole (Ye et al., ; Wang et al., ; Zhang et al., ; Scarlatos & Lan, ; Lu et al., ; Peng et al., ; Xu et al., ). While training-based approaches usually achieve superior performance, their reliance on repeated LLM queries and model training makes them both computationally intensive and time-consuming.

Our proposed GraphIC method is not only training-free and inherently efficient but also incorporates an advanced graph-based example retriever specifically designed for multi-step reasoning tasks. This sophisticated design enables GraphIC to achieve a significant performance advantage, even surpassing training-based methods.

## 3 Preliminaries: Bayesian Network

A Bayesian Network (BN) (Pearl, ) is a probabilistic graphical model that represents conditional dependencies among random variables via a directed acyclic graph (DAG). In a DAG $G=(V,E)$, $V=\{v_{1},\dots,v_{n}\}$ denotes the vertices corresponding to random variables, and $E$ denotes the conditional dependencies. Each vertex $v_{i}$ is associated with a random variable $X_{i}$, and the joint probability distribution is factorized as:

$$ $p(x_{1},x_{2},\dots,x_{n})=\prod_{i=1}^{n}p(x_{i}|\text{pa}(v_{i})),$ (1) $$

where $x_{i}\in\mathbb{R}^{n_{f}}$ denotes the value of the random variable $X_{i}$, $\text{pa}(v_{i})$ refers to the set of parent variables for $v_{i}$, and $p(x_{i}|\text{pa}(v_{i}))$ is typically modeled as:

$$ $p(x_{i}|\text{pa}(v_{i}))=g(\text{dist}(x_{i},\hat{x}_{i})),$ (2) $$

with $\hat{x}_{i}=Wz_{i}$ and $z_{i}=f(\text{pa}(v_{i}))$. Here, $\hat{x}_{i}$ represents the predicted value of $x_{i}$ based on $z_{i}$, where $z_{i}$ aggregates information from the parent nodes $\text{pa}(v_{i})$. The weight matrix $W$ is used to predict $\hat{x}_{i}$, $f(\cdot)$ denotes the aggregation function, $\text{dist}(\cdot,\cdot)$ is a distance metric between $x_{i}$ and $\hat{x}_{i}$, and $g(\cdot)$ is a function that satisfies: 1). monotonicity: $g^{\prime}(u)\leq 0$ for $u\geq 0$; and 2) normalization: $\int_{x}g(\text{dist}(x,\hat{x}_{i}))dx=1$.

Given the aggregated features $Z=(z_{1},\dots,z_{n})^{\top}$ where $\top$ denotes the transpose operation, organizing the individual feature vectors $z_{i}$ into a matrix where each row corresponds to a feature vector, along with the distance function $\text{dist}(\cdot,\cdot)$ and function $g(\cdot)$, the joint probability density of the dataset $X=(x_{1},x_{2},\dots,x_{n})^{\top}$ can be computed.

## 4 The Proposed GraphIC

Figure: Figure 2: The overall pipeline of GraphIC. First, the question and candidate examples are processed through the thought graph generation module, where the LLM generates formalized reasoning representations, which are then parsed into thought graphs. For the question’s thought graph, we extract $X$ (cognitive process per vertex) and compute aggregated feature $Z$. For candidates, parameters $\alpha_{i}$ and $\beta_{i}$ are estimated to capture relevant thought patterns. We then evaluate the applicability of these patterns on the query’s thought graph, enabling ICE selection.
Refer to caption: x2.png

In this work, we propose a novel approach called GraphIC for representing the problem-solving process through a graph-based model, intending to select examples that maximize the probability density of capturing the correct reasoning process. First, we introduce “thought graphs”, a formal structure designed to represent the reasoning process underlying each example (Section 4.1). Building on this, we develop a probabilistic framework for modeling the thought graph, enabling us to compute the probability density associated with a given thought graph (Section 4.2). This probabilistic model serves as the foundation for our example selection strategy, which prioritizes examples that maximize the probability density of thoughts on the thought graph of query question (Section 4.3). The overall framework is illustrated in Figure 2.

### 4.1 Thought Graph and its Construction

We begin by introducing the concept of a thought graph and its construction, inspired by the hierarchical structure of human cognition in problem-solving. The human thought process, especially when addressing complex problems, can be naturally modeled as a graph-like structure (Friston, ; Besta et al., ; Yao et al., ). In this work, we present the “thought graph” as a structure for modeling the cognitive process of LLMs during multi-step reasoning tasks. Formally, a thought graph $G$ is represented as a vertex-attributed graph, where each vertex is associated with a natural language text, corresponding to the description of the operation performed or the intermediate conclusion reached at that step. To facilitate computation, we further represent the vertex attributes as the BERT embedding (Devlin et al., ) of corresponding text, denoted as $x_{i}$.

Since LLMs are not natively equipped to output graph structures, we propose a methodology to generate these graphs from LLM outputs. As illustrated in Figure 2, we prompt the LLM to generate a “formalized reasoning representation” (FRR), which is subsequently parsed to construct the thought graph. The detailed prompt and parser pseudo-code are provided in Appendix D.

After constructing the thought graph, a straightforward way to select in-context examples (ICEs) is to compare the similarities between graph embeddings. To compute these embeddings, we employ a widely adopted method where each graph embedding is generated through iterative aggregation of node information, as outlined by Togninalli et al. (). Specifically, this process is formalized as:

$$ $X^{h+1}=\tilde{A}X^{h},\ \ X^{0}=X=(x_{1},x_{2},\dots,x_{n})^{\top},$ (3) $$

where $\tilde{A}=\tilde{D}_{A}^{-\frac{1}{2}}(A+I)\tilde{D}_{A}^{-\frac{1}{2}}$, $\tilde{D}_{A}=1+\sum_{j}A_{ij}$. $A$ represents the adjacency matrix of the thought graph, where $A_{ij}=1$ indicates a directed edge from node $v_{i}$ to node $v_{j}$, and $A_{ij}=0$ otherwise.

While this approach effectively captures the structural properties of human thought as represented in the graph, it is constrained by its focus on graph similarity alone. Importantly, selecting an example based solely on the similarity does not necessarily optimize the possibility of an LLM generating a correct reasoning trajectory. To overcome this, we further propose a novel example retrieval model that prioritizes the optimization of the probability density of producing a correct reasoning process detailed in the next subsection, moving beyond a mere reliance on graph similarity.

### 4.2 Probabilistic Model on Thought Graph

Building on the method for constructing thought graphs, we now turn to developing a probabilistic model for this structure. BNs model the dependencies of a node’s attributes on its parent nodes, which closely mirror the way human cognition functions—where new thoughts are informed by prior ones (Oaksford & Chater, ; Jacobs & Kruschke, ). This makes them a well-suited framework for modeling the thought graphs. In this section, we outline the construction of BNs for thought graphs. As described in Section 3, calculating the joint probability density on the thought graph requires the aggregate feature $Z$, the distance metric $\text{dist}(\cdot,\cdot)$, and the function $g(\cdot)$. We now provide a detailed discussion of how each of these components is constructed.

Computing the Aggregated Feature $\bm{Z}$.  Traditional BNs, which rely on the Markov assumption that a node’s feature distribution depends solely on its parent nodes, are insufficient for modeling a thought graph where reasoning often requires referencing multiple prior steps. For example, problem solvers may need to iteratively review information from earlier steps or return to the beginning to re-examine the entire reasoning process. To address this limitation, we first employ an iterative aggregation mechanism that better captures the human reasoning processes. This iterative approach is formalized in Equation (3). Next, inspired by the Personalized PageRank (PPR) algorithm (Page, ; Gasteiger et al., ), we refine this method to more accurately simulate the flow of information during problem-solving. The PPR framework models a random walk where a user transitions between web pages with some probability of returning to the start. This closely parallels the cognitive process in complex problem-solving, where solvers often revisit initial hypotheses to reassess their reasoning. Therefore, the iterative feature aggregation is defined as follows:

$$ $\displaystyle X^{(h+1)}=\left[(1-\lambda)\tilde{A}+\lambda\tilde{B}\right]X^{( h)},\ \ X^{(0)}=(x_{1},x_{2},\dots,x_{n})^{\top},$ (4) $$

where $\lambda\in(0,1)$, $\tilde{B}=\tilde{D}_{B}^{-\frac{1}{2}}(B+I)\tilde{D}_{B}^{-\frac{1}{2}}$, and $\tilde{D}_{B}=1+\sum_{j}B_{ij}$. The matrix $B$ models the retracing aspect of the thought process, where $B_{ij}=1$ if $\text{deg}(v_{j})=0$, and $B_{ij}=0$ otherwise, with $\text{deg}(v_{j})$ representing the in-degree of node $v_{j}$.

After $H$ iterations, the aggregated feature matrix $Z$ is given by:

$$ $Z=\left[(1-\lambda)\tilde{A}+\lambda\tilde{B}\right]^{H}X.$ (5) $$

Distance Metric $\bm{\text{dist}(\cdot)}$ and Function $\bm{g(\cdot)}$.  In prior works (Rubin et al., ; Ye et al., ; Xiong et al., ), the inner product has been a standard approach for quantifying vector similarity. Building on this, we define the distance function $\text{dist}(\cdot)$ (see equation 1) in terms of the inner production as follows:

$$ $\text{dist}(x_{1},x_{2})=l-x_{1}^{\top}x_{2},$ (6) $$

where $l$ is a sufficiently large constant chosen to ensure that $\text{dist}(x_{1},x_{2})$ remains positive. A suitable choice for $l$ is the square of the maximum norm of the embeddings produced by the model:

$$ $l=\max_{t}\left[\text{Emb}(t)^{\top}\text{Emb}(t)\right],\ t\in\mathcal{NL},$ (7) $$

with $\text{Emb}(\cdot)$ denoting a text embedding model and $\mathcal{NL}$ denoting the set of all natural languages. Additionally, we define $g_{i}(u)=\frac{1}{C_{i}}\exp(-u)$ (see equation 1), allowing us to represent:

$$ $p(x_{i};G,X)=g_{i}(\text{dist}(x_{i},\hat{x}_{i}))=\frac{1}{C_{i}}\exp\left[-( l-\hat{x}_{i}^{\top}x_{i})\right]=\frac{1}{C_{i}}\exp\left[-(l-z_{i}^{\top}W^{ \top}x_{i})\right],$ (8) $$

where $C_{i}$ is a normalization constant.

Note that this formulation establishes a probabilistic model for the thought graph. Given the parameters $W$, we can compute the probability density of each vertex attribute, which in turn allows us to determine the probability density of the entire graph. In essence,
the matrix $W$ governs the generation of new attributes and is meant to capture or represent the underlying structure of reasoning or connections between different concepts or ideas in the thought graph.

### 4.3 Probabilistic Example Retrieval

As outlined in Section 4.2, the parameter $W$ is meant to capture and represent the underlying structure of reasoning or connections between different concepts or ideas within the thought graph.
The task of computing the probability density of generating a particular thought graph given $W$ can be interpreted as evaluating the likelihood of producing the associated reasoning process based on the thought pattern encoded within $W$. Building on this idea, we design an example retrieval mechanism that estimates the model parameters for each candidate example and prioritizes those that maximize the probability density of the thought graph corresponding to the query. These selected examples serve as ICEs, offering the highest potential for accurately solving the problem at hand.

Estimation of Model Parameters.  We estimate the parameter matrix $W$ by maximizing the likelihood function $\mathcal{L}_{W}$ of the thought graph features, which is computed as

$$ $\mathcal{L}_{W}=\prod_{i=1}^{n}p(x_{i}|G),\ \log\mathcal{L}_{W}=\sum_{i=1}^{n} \log p(x_{i}|G).$ (9) $$

To simplify the computation, this reduces to the following:

$$ $\log\mathcal{L}_{W}=-\sum_{i=1}^{n}\log C_{i}+\sum_{i=1}^{n}\left[-(l-z_{i}^{ \top}W^{\top}x_{i})\right]=-\sum_{i=1}^{n}\log C_{i}-nl+\text{tr}(ZW^{\top}X^{ \top}).$ (10) $$

Hence, maximizing $\mathcal{L}_{W}$ is equivalent to maximizing $\text{tr}(ZW^{\top}X^{\top})$, formally expressed as:

$$ $\max_{W}\ \ \text{tr}(ZW^{\top}X^{\top}),\ W\in\mathbb{R}^{n_{f}\times n_{f}}, \ \text{s.t.}\ ||W||_{F}=(\sum_{i=1}^{m}\sum_{j=1}^{n}w_{ij}^{2})^{\frac{1}{2} }=1.$ (11) $$

This constraint ensures that the magnitude of $W$ does not influence the optimization.

Typically, the number of vertices $n$ in the thought graph is much smaller than the embedding dimensionality $n_{f}$ (i.e., $n\ll n_{f}$). For instance, in the GSM8K dataset, thought graphs often contain fewer than 20 vertices, while the embeddings $n_{f}$ can be as high as 768 if BERT is used. This dimensional disparity makes the solution for $W$ non-unique. Moreover, storing and computing a matrix of size $n_{f}\times n_{f}$ is computationally burdensome. To address both uniqueness and computational efficiency, we constrain $W$ to to be of rank 1, reducing it to the form:

$$ $W=\alpha\beta^{\top},\ ||\alpha||_{2}=||\beta||_{2}=1,\ \alpha,\beta\in\mathbb {R}^{n_{f}}.$ (12) $$

This simplifies the optimization to:

$$ $\text{tr}(ZW^{\top}X^{\top})=\text{tr}(Z\beta\alpha^{\top}X^{\top})=\text{tr}( \alpha^{\top}X^{\top}Z\beta)=\alpha^{\top}X^{\top}Z\beta.$ (13) $$

Thus, we reformulated the problem as a bilinear form maximization problem:

$$ $\max_{\alpha,\beta}\alpha^{\top}X^{\top}Z\beta,\quad\text{s.t.}\quad||\alpha|| _{2}=||\beta||_{2}=1.$ (14) $$

The closed-form solution to this problem (Leon, ) can be obtaind as:

$$ $\displaystyle\alpha=U[0,:],\ \beta=V[0,:],\ \text{where}\quad U,\Sigma,V=\text {SVD}(X^{\top}Z).$ (15) $$

Selection of Examples.  We extract parameters $(\alpha,\beta)$ from the thought graph of each candidate example, denoted as $\{(\alpha^{1},\beta^{1})\}_{i=1}^{N}$, where $N$ represents the size of the candidate set.
For a given query $q$, we construct its thought graph $G^{q}$ and derive $Z^{q}$, then
we select the top $k$ candidate examples that maximize the probability density $p_{i}(X^{q})$ of generating the correct reasoning process:

$$ $p_{i}(X^{q})=(\alpha^{i})^{\top}(X^{q})^{\top}Z^{q}\beta^{i}.$ (16) $$

These selected examples are taken as ICEs in ICL, where the detailed templates used in ICL are given in Appendix B.

## 5 Experiments

### 5.1 Datasets and Implementation Details

We conduct a comprehensive evaluation of GraphIC model across four multi-step reasoning benchmarks: two for mathematical reasoning (GSM8K (Cobbe et al., ) and AQUA (Ling et al., )), one for code generation (MBPP (Austin et al., )), and one for logical reasoning (ProofWriter (Tafjord et al., )). For both GSM8K and MBPP, we utilize the original datasets without further preprocessing. For AQUA and ProofWriter, we refine the original dataset to improve the experimental setup, as detailed in Appendix A.

For GSM8K, AQUA, and ProofWriter, model performance is evaluated based on the accuracy of the LLMs’ final answers. For MBPP, we adopt the pass@1 metric (Chen et al., ) to assess the quality of code generation.

We employ GPT-4o-mini and Llama-3.1-8B-Instruct as LLMs. Unless explicitly mentioned otherwise, all evaluations are conducted under an 8-shot paradigm. We set the temperature to 1e-5. We set iterations $H$ in Equation 5 to 3, with $\lambda$ values from $\{0,0.1,0.2,0.3\}$, based on the LLM and dataset (see Appendix D for details). For GSM8K, AQUA, and ProofWriter, we prompt the LLM to create a formalized reasoning representation (FRR) for thought graph construction, using vertex features from a BERT model. For MBPP, we use the staticfg module to parse Python code and generate the control flow graph, embedding each vertex’s features with CodeBERT (Feng et al., ). Variable names in the code are anonymized with arbitrary symbols like ‘a’, ‘b’, and ‘c’.

### 5.2 Baselines

Our model, GraphIC, is designed as a training-free retriever for ICE selection. We compare GraphIC against six training-free retrieval methods spanning random, similarity-based, diversity-based, and complexity-based approaches, including:
1) Randomrandomly selects $k$ unique ICEs from the candidate set;
2) BM25(Robertson et al., )selects the top $k$ examples based on BM25 scoring;
3) BERT(Devlin et al., )is a dense retriever using cosine similarity with BERT-base-uncased embeddings;
4) Complex-CoT(Fu et al., )selects $k$ examples based on complexity, quantified by newline characters;
5) Auto-CoT(Zhang et al., )clusters candidates and selects the closests to each cluster center; and
6) Skill-kNN(An et al., )prompts LLM to generate task-relevant skills for query and candidates, followed by dense retrieval. Since Skill-kNN does not natively support datasets like GSM8K, we manually craft the instructions and examples, detailed in Appendix C.

We also compare with four training-based retrievers, which encompass both single-example and combination-example retrieval strategies, including:
1) EPR(Rubin et al., )is trained to retrieve the single most relevant ICE, with top $k$ examples being selected during inference;
2) CEIL(Ye et al., )uses determinantal point processes to select ICEs balancing similarity and diversity, where three CEIL models per dataset with scaling factors of 0.01, 0.05, and 0.1 are trained and the best results are reported;
3) DQ-LoRe(Xiong et al., )uses dual queries and low-rank approximation re-ranking to identify ICEs; and
4) GistScore(Gupta et al., )encodes task-specific information into gist tokens for selecting ICEs.
Following Skill-kNN, We use GPT-J-6B (Wang & Komatsuzaki, ) as the scoring LLM for EPR, CEIL, and DR-LoRe.

### 5.3 Main Results

Figure 3 illustrates the thought graphs corresponding to each dataset. Table 1 evaluates our GraphIC model against 10 baselines across two LLMs and four datasets. As a training-free method, GraphIC consistently outperforms both training-free and training-based baselines in most settings. With the GPT-4o-mini model, GraphIC achieves the highest performance, averaging 2.57% above the leading training-free model and 1.18% above the best training-based model. For the Llama-3.1-8B-Instruct model, GraphIC ranks first in three out of four datasets, with an average gain of 4.29% over the top training-free competitor and 2.5% over the strongest training-based method.
Our analysis shows that the GraphIC model significantly enhances performance in mathematical and logical reasoning tasks versus code generation, especially for complex problems. For instance, in the GSM8K dataset, GraphIC outperforms all baselines by an average of 0.65% and 3.57% with two LLMs. In the more challenging AQUA dataset, improvements rise to 3.47% and 7.64%.

Figure: Figure 3: Examples of thought graphs. For GSM8K and AQUA, Vertices indicated by dashed lines represent the numbers entered during calculations, which will be removed in subsequent steps. The purple vertice indicates the final step in the reasoning process.
Refer to caption: x3.png

**Table 1: Main results on two LLMs and four datasets. For random retrieval, we present the mean and standard deviation derived from five independent experiments. Bold numbers indicate the best results, while underlined numbers represent the second-best results.**
| LLM | Model | GSM8K | AQUA | MBPP | ProofWriter | Avg. |
| --- | --- | --- | --- | --- | --- | --- |
| GPT-4o-mini | Random | 92.90 (0.31) | 71.58 (0.72) | 72.76 (0.74) | 64.9 (0.93) | 75.54 (0.36) |
| BM25 | 92.64 | 70.47 | 73.4 | 66.25 | 75.69 |  |
| BERT | 93.02 | 66.93 | 74.2 | 65.25 | 74.85 |  |
| Complex-CoT | 92.49 | 67.32 | 74.2 | 64.25 | 74.57 |  |
| Auto-CoT | 92.72 | 69.69 | 73.8 | 62.25 | 74.62 |  |
| Skill-kNN | 92.34 | 71.65 | 72.0 | 66.00 | 75.50 |  |
| EPR | 93.02 | 72.04 | 73.8 | 68.50 | 76.84 |  |
| CEIL | 92.57 | 72.44 | 73.8 | 69.50 | 77.08 |  |
| DQ-LoRe | 93.32 | 69.69 | 74.6 | 66.50 | 76.03 |  |
| GistScore | 93.25 | 69.69 | 72.8 | 67.0 | 75.69 |  |
| GraphIC | 93.48 | 73.62 | 75.2 | 70.75 | 78.26 |  |
| Llama-3.1<br>-8B-Instruct | Random | 78.86 (0.87) | 53.15 (1.85) | 57.72 (1.06) | 76.1 (2.45) | 66.46 (0.84) |
| BM25 | 77.71 | 46.85 | 62.0 | 77.75 | 66.08 |  |
| BERT | 74.15 | 50.39 | 60.8 | 73.75 | 64.77 |  |
| Complex-CoT | 79.30 | 50.00 | 58.6 | 78.25 | 66.54 |  |
| Auto-CoT | 72.78 | 42.91 | 58.4 | 78.00 | 63.02 |  |
| Skill-kNN | 77.56 | 50.39 | 60.8 | 74.00 | 65.69 |  |
| EPR | 75.66 | 53.94 | 62.0 | 79.25 | 67.71 |  |
| CEIL | 75.51 | 51.97 | 62.4 | 81.00 | 67.72 |  |
| DQ-LoRe | 77.93 | 54.33 | 59.8 | 81.25 | 68.33 |  |
| GistScore | 74.60 | 44.49 | 60.4 | 79.50 | 64.75 |  |
| GraphIC | 79.98 | 57.48 | 61.6 | 84.25 | 70.83 |  |

Additionally, we observe that the GPT-4o-mini model’s performance on the GSM8K dataset is relatively invariant to the selection of ICEs.
So, we further perform an additional experiment on the GSM8K dataset using the GPT-3.5-Turbo model. As presented in Table 2, GraphIC model achieves superior performance across all metrics.

**Table 2: Results obtained using GPT-3.5-Turbo as the LLM on the GSM8K dataset.**
| Random | BM25 | BERT | Complex-CoT | Auto-CoT | Skill-kNN | EPR | CEIL | DQ-LoRe | GistScore | GraphIC |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 80.76(0.55) | 82.10 | 80.89 | 81.65 | 82.03 | 81.50 | 81.65 | 81.72 | 82.10 | 81.72 | 82.79 |

### 5.4 Ablation Study

We perform a series of ablation studies to systematically evaluate the contribution of each component within the GraphIC framework, which is built upon three key pillars: the incorporation of thought graphs, PPR for aggregating features, and BN-based retrieval.

To this end, we develop several variants of the GraphIC model:
1) Textrelies solely on text embeddings, the same as the BERT approach;
2) FRRretrieves examples using BERT embeddings derived from FRRs (or CodeBERT embeddings for the MBPP dataset);
3) Graphutilizes the formula (3) to generate graph embeddings, which are employed for dense retrieval;
4) Graph+PPRuses the formula (5) to obtain graph embeddings for dense retrieval;
5) Graph+BNexcludes the backtracking (or PPR) mechanism from the full GraphIC model during computing $Z$;
and
6) Graph+PPR+BNrepresents the full GraphIC model, integrating all components.

We conduct experiments leveraging the Llama-3.1-8B-Instruct model across four datasets, with the outcomes detailed in Table 3. The findings underscore that each component of GraphIC plays a pivotal role in boosting model performance, with the most significant improvements observed when all three components are utilized in conjunction.

Table: Table 3: Ablation Study.

| Model | GSM8K | AQUA | MBPP | ProofWriter |
| --- | --- | --- | --- | --- |
| Text | 74.15 | 50.39 | 60.8 | 73.75 |
| FRR | 78.31 | 50.78 | 60.4 | 82.50 |
| Graph | 78.46 | 54.72 | 60.4 | 83.50 |
| +PPR | 78.92 | 56.30 | 61.0 | 83.75 |
| +BN | 79.07 | 49.21 | 60.4 | 84.25 |
| +PPR+BN | 79.98 | 57.48 | 61.6 | 84.25 |

### 5.5 Analysis

Figure: (a)
Refer to caption: x4.png

Impact of ICE Examples on Model Performance.  We conduct an in-depth investigation into the influence of the number of ICEs on the performance of our proposed GraphIC model and several competitive baselines across four datasets. For each dataset, we select the top three baselines and varied the number of examples in the set {1, 2, 4, 8}. Llama-3.1-8B-Instruct is employed as the underlying LLM. Results in Figure 4 indicate a general trend of improved model performance with an increase in the number of examples. Notably, The performance of GraphIC steadily improves as the number of ICEs increases, unlike some baseline methods, which may experience performance degradation when the number of ICEs increases. Furthermore, while GraphIC initially lags behind the baselines in the low-shot settings, its performance exhibits a more pronounced improvement as the number of examples grew. One can observe that GraphIC surpasses the baselines, demonstrating superior performance and underscoring its robustness as the number of examples increases.

Figure: (a)
Refer to caption: x8.png

Assumption of Symmetry.  Our findings show that the GraphIC model uses an “asymmetric” approach, unlike the common symmetry assumption in most baseline models. To assess the validity of this assumption, we conduct an experiment examining whether symmetry, in the context of retrieval models, holds. Here, a retrieval model is considered symmetric if $\text{score}(i,j)=\text{score}(j,i)$, where $\text{score}(i,j)$ represents the model’s assessment of example $i$ as an ICE for example $j$. For example, the Skill-kNN is symmetric as it uses inner product embeddings for $\text{score}(j,i)$, while the Complex-CoT model is asymmetric, calculating $\text{score}(j,i)$ based on the complexity of example $i$.

We randomly select 10 examples from the GSM8K candidate set and use the top-7 performing models to compute the score matrix $S$ ($S_{ij}=\text{score}(i,j)$), which we then compare against the ground truth matrix $S^{\text{gt}}$. Here, $S^{\text{gt}}_{ij}$ captures the probability that Llama-3.1-8B-Instruct provides the correct answer when example $i$ is used as an ICE for example $j$.

The experimental results show that the ground truth matrix (Figure 5 (a)) is asymmetric, undermining the symmetry assumption. Using symmetric models for inherently asymmetric data introduces significant errors. For instance, the EPR model, trained on correct answer probabilities, struggles with accuracy due to its symmetry reliance (Figure 5 (f)). In contrast, simple asymmetric methods like Complex-CoT (Figure 5 (b)) and BM25 (Figure 5 (d)) perform well, ranking second and fourth as Table 1 shows, and surpassing many symmetric models. However, their simplistic assumptions limit their ability to capture ground truth nuances. In contrast, GraphIC (Figure 5 (h)), a sophisticated asymmetric model, aligns closely with the ground truth, resulting in superior performance.

## 6 Conclusion

We introduce GraphIC, a graph-based method for in-context example (ICE) retrieval aimed at enhancing LLM performance on multi-step reasoning tasks. By modeling reasoning as “thought graphs” and utilizing Bayesian Networks and personalized PageRank, GraphIC selects ICEs that align with the task’s cognitive structure, overcoming the limitations of text-based embedding methods. Extensive experiments on four benchmarks show that GraphIC consistently outperforms both training-free and training-based baselines, especially in mathematical and logical reasoning. Our analysis of symmetry assumptions highlights the advantage of asymmetric retrieval models. A limitation of the GraphIC model is that, as a training-free framework, it may face difficulties in capturing more intricate thought patterns. Beyond this, GraphIC not only introduces a powerful ICEs retrieval method, but more crucially, it provides a way to represent and understand the reasoning process. This capability can be applied to various domains related to LLM reasoning, such as developing novel graph-based reasoning methods, selecting high-quality and diverse training datasets, and more.

## Appendix A Processing of AQUA and ProofWriter

Given the substantial size of the AQUA dataset, which incurs significant retrieval overhead during testing, we followed the methodology outlined in DQ-LoRe (Xiong et al., ), using a 1,000-sample subset for efficient evaluation.

For the ProofWriter dataset, we refined the subset selected by Logic-LM (Pan et al., ), excluding instances labeled as “Unknown,” as these samples lacked explicit reasoning chains. Furthermore, because the original training set did not provide reasoning in natural language, we leveraged the GPT-4o-mini model to generate reasoning sequences for the training set, discarding any generated outputs deemed incorrect. We evaluate the correctness of the reasoning process by the correctness of the final result, which is a commonly used approach (Lightman et al., ; Xiong et al., ; Khattab et al., ). This process resulted in a refined training set of 1,358 examples with their Chains of Thought and 400 test samples from the original ProofWriter dataset.

## Appendix B Prompt Templates

For the four datasets under consideration, we design the following prompt templates to format the ICEs and the question into a prompt, which is then fed into an LLM to generate answers.

GSM8K & AQUA:

MBPP:

ProofWriter:

## Appendix C Skill-kNN

Since Skill-kNN does not offer prompts for skill generation in these three tasks, we referred to the prompt designed for the semantic parsing task in the original paper to write prompts for the four datasets we used. First, we applied the Complex-CoT method to select 8 examples, then employed the GPT-4o model to generate skills in a zero-shot setting. Finally, we integrated these results to construct the final prompt.

GSM8K:

AQUA:

MBPP:

ProofWriter:

## Appendix D GraphIC

### D.1 Formalized Reasoning Representation

The prompt examples below are used to generate formalized reasoning representations for the four datasets being considered. For the test question, since no answer is provided, we will remove the section of the prompt highlighted in blue. This will allow the LLM to generate both the answer and the formalized reasoning representation simultaneously, from which we can then extract the formalized reasoning representation.

GSM8K:

AQUA:

MBPP:

ProofWriter:

The pseudo-code of a parser that transforms formalized reasoning representations into a thought graph is provided in Algorithm 1. Inputs, Output, and OperationName are extracted following the pattern outlined below.

Figure: Algorithm 1 Parsing Formalized Reasoning Representation

### D.2 Values of λ 𝜆 \lambda italic_λ

We select hyper parameter $\lambda$ values from $\{0,0.1,0.2,0.3\}$, and report the
$\lambda$ values chosen on various datasets and LLMs in Table 4.

**Table 4: $\lambda$ values chosen on various datasets and LLMs.**
| Engine | GSM8K | AQUA | MBPP | ProofWriter |
| --- | --- | --- | --- | --- |
| GPT-4o-mini | 0.2 | 0.2 | 0.1 | 0.1 |
| Llama-3.1-8B-Instruct | 0.3 | 0.2 | 0.2 | 0.0 |
| GPT-3.5-Turbo | 0.3 | / | / | / |

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Can you estimate the extra inference and time cost for GraphIC? For other retrieval baselines, they do not need to create the thought graph.  
2. In Figure 4, the performance keeps increasing with more in-context examples. Do you think it will benefit from more than even more examples? For example, on MBPP, can 16-shot outperform the CEIL baseline?
soundness: 3
strengths: 1. This paper uses a formalized reasoning representation to construct a thought graph for complex reasoning problems. Based on that, it can better model the underlying reasoning process than the semantic representation of natural language.  
2. It enhances the graph embedding by the personalized PageRank and establishes a probabilistic model for the thought graph. GraphIC retrieves in-context examples by selecting top-k candidate examples that can maximize the probability of generating the correct thought graph. The probabilistic method can better capture the examples that can enhance the reasoning of the new query.  
3. The paper verifies the proposed method on diverse reasoning benchmarks with multiple training-free and training-based retrieval baselines, and the results demonstrate its effectiveness.   
4. It further conducts a comprehensive ablation study on each component and investigates the symmetry of different retrieval methods.
summary: This paper focuses on improving the selection of the in-context examples. It proposes GraphIC, which leverages the graph-structure and Bayesian Network to select in-context examples for complex reasoning tasks. Experiments on three types of reasoning tasks (math, code, and logical reasoning) demonstrate GraphIC outperforms the other ICE selection methods.
weaknesses: 1. GraphIC relies on the thought graph, which is generated by formalized reasoning representation from LLMs. How to ensure the correctness of the thought graph of candidate examples? Will it be multiple possible thought graphs for the same query? Will these factors affect the robustness of GraphIC?  
2. For a test query q, GraphIC first creates the thought graph G^q without the ground-truth answer and retrieve in-context examples to maximize the probability density p_i(X^q). This also assumes that the thought graph G^q is correct. What if the thought graph has an incorrect reasoning process? Will it mislead the retrieval and thus affect the performance?

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 1
questions: Can you please explicitly state, e.g., what joint the BN is representing? What exactly is a "graph of thought"? Can you please provide a brief, clear, and self-contained description of the proposed method?
soundness: 2
strengths: The proposed method seems to produce marginally better results than the baselines in most cases.
summary: This paper proposes GraphIC, a graph-based method for in-context example retrieval aimed to multi-step reasoning tasks. GraphIC models CoTs as "tought graphs, and uses Bayesian networks and personalised pagerank for selecting in-context examples with similar underlying graph structures. Empirical results show marginal improvements or competitive results on a wide array of reasoning tasks (GSM8K, AQUA, MBPP, ProofWriter).

Bayesian Networks (BNs) are a common way to represent complex joint distributions over sets of variables; however, this work never formalises which joint distribution the BNs in this work aim to represent. Furthermore, in Eq. 2, it is not really true that "commonly, in BNs, $p(x_i | pa(v_i)) = g(dist(x_i, \hat{x}_i))$" -- BNs aim at representing dependence distributions between variables, like "rainy weather" and "risk of aqua planning", and it's not really clear what a distance function between those variables should represent.

Furthermore, the paper is based on the construction of "thought graphs", but it's not really clear what these are -- are they task dependent? How do they look like? Given a taks, how does someone create a "thought graph"? The paper also says that "to facilitate computation, we further represent the vertex attribures as the BERT embedding of corresponding text" -- what is going on exactly?

Then, my understanding is that the BN defines a distribution over "thoguht graphs", and that the distribution over $\hat{x}_i$; what does $\hat{x}_i$ represent? (E.g. in Eq. 8)

Results seem to be marginally better than the baselines in most cases, but it's really hard to understand what's going on in the proposed method.
weaknesses: It is very hard to understand what's going on in the proposed method -- for example, the method uses Bayesian networks but the paper never explicitly states which join distribution the Bayesian network aims to represent.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: * In the 'Estimation of Model Parameters' paragraph, the rank of W is set to 1, there is no theoretical basis for this hypothesis, and how much precision loss is caused?
* As mentioned in weakness, LLM is probably hard to construct thought graph, do you try more complex datasets, such as widely used mathematical MATH dataset [1]?
* How accurate is the LLM at generating a thought graph? 
* Multiple generations of the LLM may produce different thought graphs, how much will this affect the results?

[1] https://github.com/hendrycks/math/
soundness: 3
strengths: * The motive of the paper is reasonable and the method proposed is novel.
* Writing of this paper is good, with reasonable structure.
* The experiments are relatively abundant, and the experimental results can prove the conclusion of the paper.
summary: This paper studies how to retrieve demonstration examples more effectively in in-context learning to improve the performance of in-context learning. Specifically, the paper proposes a method named GraphIC, which establishes thought graphs and models them as Bayesian networks, then retrieves demonstration examples that make the probability density of queries' bayesian networks maximum. Experiments show the effectiveness of the proposed method.
weaknesses: * Some parts of the method section of the paper lack some details, there are many assumptions but no conditions, refer to questions.
* Method relies on LLM to construct a thought graph, which may be difficult or inaccurate to decompose key steps for complex problems.
* The lack of experiments on the thought graph, in my opinion, is an important part of the method and has a big impact on method performance, refer to questions.

---

confidence: 4
contribution: 3
initial_rating: 3
presentation: 2
questions: See the Weaknesses.
soundness: 3
strengths: 1. The idea of selecting ICE with reasoning structure rather than semantics is novel and inspiring, and the human cognition-inspired methodology of capturing the underlying reasoning process with thought graph and BN is interesting.
2. The inference of the proposed method is solid, which makes the method more reasonable.
3. The authors conducted extensive experiments and deeper analysis to demonstrate the effectiveness of the proposed method and its components.
summary: This paper studied in-context example selection for ICL, and proposed a novel GraphIC method to capture reasoning processes for ICL in multi-step reasoning. GraphIC first generated thought graphs for candidate examples and query with LLM to encode the reasoning steps, and then employed a BN with PPR mechanism to studied the associated parameters reflecting the underlying reasoning structure. Finally, GraphIC selected the candidate examples with the parameter maximizing the probability of thought graph for the query in ICL. The authors conducted extensive experiments, and the results demonstrated that GraphIC outperformed both training-free and training-based baselines across various multi-step reasoning tasks.
weaknesses: 1. There are some concerns on the method design.
- According to the prompts in the Appendix, the proposed method seems to exploit the answer to generate the thought graph for both examples and the query. So where does the answer come from? If the answer is fetched from the annotations, it seems to lead to label leakage. If the answer is generated by the LLM, how to ensure its consistency to real solution?
- As the authors have mentioned in section 3, BN works on the DAG. However, there exist loops in thought graphs for codes which violate the DAG assumption as shown in Figure 3.
2. The presentation of the manuscript could be further improved especially for the confusing and unclear contents.
- The authors should explain the formal definition of the thought graph more clearly. Maybe one example could help. For example, what does one vertex and edge mean in the thought graph (does one vertex mean one step), how can it capture the reasoning processes, and how is the corresponding text attribute got. If the text is that in Figure 3, I do not think it contains much useful information. Besides, it would be better to explain the vertex, edge and random variable of BN in section 3 more clearly.
- Related works on BN should be cited properly.
- The insights behind the parameter W should be explained more clearly. Why the parameter can reflect the reasoning structure, and why the probability computed with parameter of the example and graph of the query can reflect their similarity on reasoning?
- The insights of the matrix B in PPR is confusing. What does the matrix mean and why can the matrix realize the retracing. Perhaps the authors could provide one example.
- The author mentioned that the asymmetry of the proposed method, so I wonder the advantages of the asymmetry compared with symmetry.
3. Some concerns on the experiments.
- The authors could add the LLM without ICL as one baseline to demonstrate the improvements more clearly.
- It would be better to add experiments to compare the computation time of the proposed method and other ICE selection baselines.
- In Figure 4, it seems that the proposed method does not work very well with 1-4 shot. It would be better to make clearer explanation.
- It would be better to add studies on the effects of lambda.

Meta-review: Summary: The paper presents GraphIC, an approach for selecting examples for in-context learning (ICL) that focuses on capturing reasoning structure rather than just semantic similarity. The proposed method leverages graph-based representations called "thought graphs", something inspired by Bayesian Networks (BNs) and personalized PageRank to improve relevance scoring. Empirical studies conducted show better results for GraphIC on 4 tasks (spanning mathematical reasoning, code generation, and logical reasoning) against chosen training-free and training-based baselines.

Strengths:
- Important problem of example selection for ICL
- Good direction of using graph representation (but details are confusing as mentioned in weakness)
- Comparison with both training-free and training-based baselines (but gains are somewhat marginal currently, should show on more complex tasks where thought graphs have potential to shine brighter)
- Provides ablation for different components and robustness to incorrect thought graphs

Weakness:
- The submission received a lukewarm response from the reviewers and several concerns were raised.
- Lack of clarity in explaining key concepts like thought graphs and their construction
- Ambiguity in the proposed probabilistic framework as the joint probability distribution modelled by BNs was unclear
- Applicability to complex tasks where thought graph is motivated to shine most is not clear as there is limited evaluation on more complex datasets
- Concerns about the consistency and correctness of LLM-generated thought graphs
- Presentation issues making technical details hard to follow
- Additional inference time compute needed due to thought graph construction for marginal improvements.

Decision:
While the work shows promise and addresses an important problem, unfortunately, the paper can't be accepted in its current form and addressing all the concerns would warrant another round of reviewing.

Decision: Reject