OCEAN: Offline Chain-of-thought Evaluation and Alignment in Large Language Models

## Abstract

Abstract Offline evaluation of LLMs is crucial in understanding their capacities, though current methods remain underexplored in existing research.
In this work, we focus on the offline evaluation of the chain-of-thought capabilities and show how to optimize LLMs based on the proposed evaluation method.
To enable offline feedback with rich knowledge and reasoning paths,
we use knowledge graphs ( e.g. , Wikidata5M)
to provide feedback on the generated chain of thoughts.
Due to the heterogeneity between LLM reasoning and knowledge graph structures, direct interaction and feedback from knowledge graphs on LLM behavior are challenging, as they require accurate entity linking and grounding of LLM-generated chains of thought in the knowledge graph.
To address the above challenge,
we propose an offline chain-of-thought evaluation framework, OCEAN ,
which models chain-of-thought reasoning in LLMs as a Markov Decision Process (MDP),
and evaluate the policy’s alignment with knowledge graph preference modeling.
To overcome the reasoning heterogeneity
and grounding problems, we leverage on-policy knowledge graph exploration and reinforcement learning to model a knowledge graph policy that generates token-level likelihood distributions for LLM-generated chain-of-thought reasoning paths, simulating knowledge graph reasoning preference.
Then we incorporate the knowledge-graph feedback on the validity and alignment of the generated reasoning paths into inverse propensity scores and propose KG-IPS estimator.
Theoretically, we prove the unbiasedness of the proposed KG-IPS estimator and provide a lower bound on its variance.
With the off-policy evaluated value function, we can directly enable off-policy optimization to further enhance chain-of-thought alignment.
Our empirical study shows that OCEAN can be efficiently optimized for generating chain-of-thought reasoning paths with higher estimated values
without affecting LLMs’ general abilities in downstream tasks or their internal knowledge.

## 1 Introduction

Offline policy evaluation aims to estimate a target policy model’s performance with only collected data,
without requiring direct interactions between the target policy and realistic environments.
Previous offline evaluation methods focus on decision-making policies in recommender systems (Li et al., ), healthcare (Bang & Robins, ),
and other scenarios where online experimentation is costly (Thomas et al., ; Bhargava et al., ), risky, and impractical (Yu et al., ).
Recent studies in LLMs leverage human feedback to align models’ behaviors with human preferences in single-turn generation (Ouyang et al., ; Rafailov et al., )
and multi-step reasoning tasks (Joshi et al., ).
Due to the high cost of deploying LLMs online and interacting with human feedback,
Bhargava et al. () further enables offline evaluation of LLMs from logged human feedback to align LLMs’ response generation.

However, considering annotators may not be comprehensive in various types of knowledge backgrounds and the associated reasoning,
human feedback on chain-of-thought reasoning (Joshi et al., ) can be more challenging to collect.
In addition, since chain-of-thought reasoning involves a sequential decision-making process,
the volume of collected human feedback can be exponentially increased.
Due to such challenges, conventional reinforcement learning from human feedback (RLHF) methods (Ouyang et al., ; Bai et al., )
can suffer from training inefficient and scalability issues.

Motivated by recent works in using knowledge graphs as side information to enable prompt engineering (Wang et al., ),
self-correction (Zhao et al., ; Wang et al., ; Li et al., ; Wu et al., ), evaluating chain-of-thought (Nguyen et al., ), and model fine-tuning (Wang et al., ; Tang et al., ),
we propose leveraging knowledge graphs as weak yet controllable knowledge reasoners to
effectively measure the alignment between LLMs’ multi-step chain-of-thought reasoning and multi-hop knowledge graph trajectories
by inverse propensity scores (IPS) (Joachims et al., ).
In contrast to the existing chain-of-thought evaluation (Nguyen et al., ) method,
which relies on accurate chain-of-thought grounding on specific knowledge graphs,
we propose verbalizing the knowledge graph trajectories and developing a knowledge graph policy that serves as a verbal reasoning mechanism over the graphs.
Therefore, we can bridge the heterogeneity between knowledge graph and LLM reasoning forms,
and the verbalized knowledge graph policy can be generalized to be compatible with various LLMs.

To enable controllable chain-of-thought alignment in LLMs,
we principally track LLMs’ decision-making process in generating chain-of-thought reasoning steps,
by formulating the process as a Markov Decision Process (MDP)
whose goal is to reach the correct final answer with minimal knowledge exploration and exploitation.
Then, we propose offline chain-of-thought evaluation and alignment, OCEAN,
which evaluates the generated chain of thoughts from off-policy LLMs through collected offline data samples with feedback from a knowledge graph.
The improved KG-IPS approach considers the effects of feedback from the knowledge graph policy
that aligns the model’s chain-of-thought generation and the behavior policy, which prevents model degeneration. We prove that the KG-IPS estimator provides an unbiased estimate of the target policy, with a lower bound for the variance, and establish confidence intervals using sub-Gaussian concentration inequalities.
To enable direct optimization of LLM policies, we leverage the proposed KG-IPS policy evaluation approach for LLM fine-tuning by
directly maximize estimated policy values through gradient descent.
Then we empirically evaluate the optimized LLM policy on three types of chain-of-thought reasoning tasks,
and demonstrate the effectiveness of the proposed policy optimization method.
We also observe relative performance improvements across evaluation tasks,
without affecting LLMs’ generalizability or generation quality.
We summarize our contributions as follows:

- •
We propose an offline evaluation framework, OCEAN, which bridges the heterogeneity between LLM and knowledge graph reasoning, for
effective evaluations of chain-of-thought.
- •
With the evaluation framework, we further develop a direct policy optimization method which enables efficient alignment with automatic feedback from the knowledge graph.
- •
To facilitate the evaluation and optimization, we model the knowledge-graph preference and derive feedback by developing a policy which verbalizes knowledge-graph trajectories.
- •
We provide a theoretical analysis of the unbiasedness and establish a lower bound for the variance of our KG-IPS estimator.
- •
Through comprehensive experiments, we demonstrate OCEAN’s effectiveness
in aligning LLMs’ chain-of-thought reasoning through direct optimization of the estimated policy value.
OCEAN also achieves better performance on various downstream tasks without affecting LLMs’ generalizability.

## 2 Related Work

Offline Policy Evaluation
Offline policy evaluation (OPE) is essential when online deploying learned policies is risky and impractical (Levine et al., ).
OPE has been applied to various practical applications,
including evaluating the recommender system’s behavior with offline collected user feedback (Gilotte et al., ; Jeunen, ).
Recent work (Gao et al., ) also develops an OPE estimator for LLM evaluation based on human feedback.
Different from previous works, we study and formulate chain-of-thought generation in LLM as an MDP and use knowledge graph reasoning as automatic feedback to develop a KG-IPS policy value estimator.

LLM Alignment
Reinforcement Learning from Human Feedback (RLHF) has been the dominant approach, optimizing LLMs using human-annotated data to align model behavior with user preferences (Ouyang et al., ; Bai et al., ).
DPO (Rafailov et al., ) and RRHF (Yuan et al., ) are proposed to reduce the training instability of RLHF.
Wu et al. () utilizes varying densities of human feedback to offer fine-grained rewards for RL finetuning, and Sun et al. () focuses on aligning LLMs with reward models driven by human-defined principles.
To address RLHF’s limitations such as heavy reliance on human input, alternative approaches like Reinforcement Learning from AI Feedback (RLAIF) (Bai et al., ; Lee et al., ; Liu et al., ) and self-alignment methods (Sun et al., ) have been proposed, using AI-generated feedback to scale and automate alignment.
Despite advancements, a key challenge remains in aligning LLMs’ internal knowledge with their reasoning, resulting in flawed reasoning even after factual errors are corrected.
Our approach focuses on improving chain-of-thought alignment by modeling reasoning paths as an MDP and using KGs to ensure both factual accuracy and human-like reasoning.

Chain-of-thought Reasoning
Chain-of-thought prompting has been widely applied to elicit the strong reasoning abilities of LLMs (Wei et al., ; Chu et al., ; Xia et al., ).
By decomposing a complex problem into a sequence of intermediate sub-tasks, LLMs are able to focus on important details and solve the step-by-step (Huang & Chang, ; Yu et al., ).
Despite the remarkable performance improvements, recent studies have found that LLMs often generate unfaithful chain-of-thought reasoning paths that contain factually incorrect rationales (Turpin et al., ; Lanham et al., ).
To address this, a number of works leverage LLMs’ self-evaluation abilities to verify and refine each reasoning step (Ling et al., ; Madaan et al., ).
As the factual errors in the generated chain-of-thought may also be caused by the limited or outdated parametric knowledge of LLMs, recent methods incorporate external knowledge sources to further edit unfaithful content in the reasoning path (Zhao et al., ; Wang et al., ; Li et al., ; Wang et al., ; ).
While these methods focus on knowledge augmentation and editing through prompts, our method, in comparison, directly aligns LLM internal knowledge with faithful and factual chain-of-thought, which avoids potential knowledge conflicts between parametric and non-parametric knowledge when generating reasoning paths.

## 3 Preliminary

We first provide the formulation of chain-of-thought reasoning in LLMs as an MDP.
Then we discuss conventional knowledge graph reasoning,
as an alternative to free-form generation by verbalizing structured knowledge graph reasoning paths into natural language,
which is more statistically controllable and generates faithful reasoning paths to the knowledge graph.

### 3.1 Problem Formulation: CHAIN-OF-THOUGHT AS AN MDP

Given the prompt instruction $q$, chain-of-thought reasoning process in a causal language model $\pi_{\theta}$
includes the generation of a trajectory of reasoning steps $\bm{c}=(c_{1},c_{2},\dots,c_{T})$, before the final answer prediction $y$,

$$ $c_{t}\sim\pi_{\theta}(\cdot|q)=\prod_{i=1}^{t-1}\pi_{\theta}(c_{i}|q,c_{<i}), \quad y\sim\pi_{\theta}(\cdot|q)=\pi_{\theta}(y|q,\bm{c})\prod_{i=1}^{T}\pi_{ \theta}(c_{i}|q,c_{<i}),$ $$

where each reasoning step $c_{t}$ comprises a sequence of tokens and the number of reasoning step $T$ is determined by the model’s generation.
Controllable chain-of-thought generation can be challenging due to its nature in autoregressive sequential sampling (Lin et al., ),
which produces a high-dimensional action space in sampling a reasoning step $\pi_{\theta}(c_{t}|q)$ containing multiple tokens.

Chain-of-thought reasoning can be viewed as a Markov Decision Process (MDP) (Sutton, ):
starting with the instruction prompt $q$, the LLM sequentially decides and generates the next-step reasoning path $c_{t}$
that navigates until it arrives at a target final answer $y$.
Given the LLM policy $\pi_{\theta}$, at time step $t$, each state $s_{t}\in\mathcal{S}$
comprises of the instruction prompt $q$ and previously generated reasoning paths $(c_{i})_{i=0}^{t-1}$.
The action space $\left\{1,\dots,|\mathcal{V}|\right\}^{N_{t}}$ in LLMs
is a sequence of $N_{t}$ tokens sampled from an identical and finite vocabulary set $\mathcal{V}$.
The LLM policy $\pi_{\theta}$ samples next-step thought based on current state as $a_{t}\sim\pi_{\theta}(a_{t}|s_{t})$,
which is a sub-sequence in the reasoning path $a_{t}\in c_{t}$ and the surrounding context $c_{t}\backslash a_{t}$ is deterministically generated by LLMs.
The transition in chain-of-thought is concatenating each reasoning path to the current state as $s_{t+1}=[s_{t},c_{t}]$.
Then the reward function is to evaluate each thought given the state as $r_{t}=r(s_{t},c_{t})$.
Although such formulation of chain-of-thought enables direct LLM on-policy optimization via reinforcement learning,
direct interaction with knowledge graphs to collect per-step reward in LLMs can be
practically challenging and require a large effort of engineering
due to the discrepancy between the unstructured generation of LLMs and structured knowledge graphs (Pan et al., ).
Therefore, we propose to offline evaluate and optimize the target policy aligning with knowledge graph preference.

### 3.2 Verbalized Knowledge Graph Reasoning

In contrast to chain-of-thought reasoning, conventional knowledge graph reasoning methods (Lin et al., ; Saxena et al., )
sample a entity-relation pair $(r_{t},e_{t})$ at step $t$ from a subset of the graph $\mathcal{G}=(\mathcal{E},\mathcal{V})$ consisting of the outgoing edges of current edge $e_{t-1}$,

$$ $(r_{t},e_{t})\in\left\{(r^{\prime},e^{\prime})|(e_{t-1},r^{\prime},e^{\prime}) \in\mathcal{G}\right\},$ (1) $$

where the transition feasibility of the entity $e_{t-1}$ to all the outgoing edges is entirely determined by $\mathcal{G}$.
The process of knowledge graph reasoning starts with a decomposed triple $(e_{0},r_{1},e_{1})$ of the instruction $q$,
and produces a chain of triplets $\bm{h}=(e_{0},r_{1},e_{1},\dots,r_{T},e_{T})$ by sampling from a policy $\mu$,

$$ $(r_{t},e_{t})\sim\mu\left((r_{t},e_{t})|e_{0},r_{1},e_{1},\dots,r_{t-1},e_{t-1 }\right),$ (2) $$

where the goal of such knowledge graph exploration is to arrive at the correct answer entity at the end of the search step $T$.
By knowledge graph exploration, we can collect a set of trajectories $\mathbb{H}=\left\{\bm{h}_{k}\right\}_{k=1}^{K}$,
which are used to estimate a parametric probabilistic policy $\mu_{\phi}$ as a proxy to model the preference of the knowledge graph.

To align the action space between the knowledge graph preference policy $\mu_{\phi}$ and the target policy $\pi_{\theta}$,
we leverage a small language model as the backbone of $\mu_{\phi}$ and fine-tune the model on verbalized trajectories as natural language contexts.
Inspired by existing efforts in verbalizing structured knowledge graphs into natural language query (Seyler et al., ) and context (Agarwal et al., ; Wang et al., ),
we leverage the GPT-4 (Achiam et al., ) model $f$ to verbalize each chain of triplets $\bm{h}$ into a chain-of-thoughts $\bm{c}=f(\bm{h})$. The verbalized knowledge-graph trajectories are used to model knowledge graph preference in Section 4.2.

## 4 OCEAN : O ffline C hain-of-thought E valuation and A lignme N t

We propose an offline evaluation of the chain-of-thought generation process aligned with knowledge graph preference.
The off-policy estimator can be used for policy optimization that aligns LLMs with more faithful reasoning paths from knowledge graphs (Lin et al., ).
We develop a small language model as a behavior policy that models the knowledge graph preference.

### 4.1 Offline Evaluation and Optimization

One of the most broadly used offline evaluation approaches is inverse propensity scores (Ionides, ; Dudík et al., ),
which has been used for LLM-based offline policy evaluation for various purposes (Bhargava et al., ; Dhawan et al., ).
Given the offline logged chain-of-thought trajectories $\mathcal{D}=\{\tau_{i}\}_{i=1}^{N}$, where $\tau_{i}=(s_{t}^{(i)},c_{t}^{(i)},r_{t}^{(i)},s_{t+1}^{(i)})_{t=0}^{T_{i}}$,
we propose a KG-IPS estimator considering two-folded weights of entity tokens in the knowledge graph preference policy $\mu_{\phi}$ and of non-entity tokens in the base LLM policy $\pi_{0}$ ,

$$ $\displaystyle\hat{V}_{KG-IPS}(\theta)=\frac{1}{N}\sum_{i=1}^{N}\frac{1}{T_{i}| c_{t}^{(i)}|}\sum_{t=1}^{T_{i}}\sum_{e\in c_{t}}\frac{\pi_{\theta}(e|s_{t}^{(i )})}{\lambda(e|s_{t}^{(i)})}\log\pi_{0}(e|s_{t}^{(i)}),$ (3) $$

where
$\lambda(e|s_{t}^{(i)})=\textbf{1}\{e\in a_{t}^{(i)}\}\cdot\mu_{\phi}(e|s_{t}^{
(i)})+\textbf{1}\{e\in c_{t}^{(i)}\backslash a_{t}^{(i)}\}\cdot\pi_{0}(e|s^{(i
)})$.
We follow (Zhang et al., ) to use the log-likelihood score of each token in the base policy $\pi_{0}$ as the reward function.

Establishing the unbiasedness of the KG-IPS estimator is essential for reliable policy evaluation (Jiang & Li, ; Bhargava et al., ). We formalize this in the following lemma and provide the proof in Appendix 8:

###### Lemma 1 .

The KG-IPS estimator provides an unbiased estimate of the target policy $\pi_{\theta}$.

The standard IPS estimator is known to have high variances considering large behavior discrepancies ($\pi_{\theta}(e|s_{t}^{(i)})/\mu_{\phi}(e|s_{t}^{(i)})$) between the behavior policy $\mu_{\phi}$ and the target policy $\pi_{\theta}$.
In addition, by separately weighting the entity and non-entity tokens with $\mu_{\phi}$ and $\pi_{0}$ respectively,
we avoid the increasing variance accumulated from the long chain-of-thought reasoning process and maintain the LLM’s behaviors on non-entity tokens without model degeneration.
To further formalize our approach and illustrate the variance inherent in the KG-IPS estimator, we present the following Lemma, which provides a lower bound on the variance,

###### Lemma 2 .

The variance of the KG-IPS estimator is lower bounded by $\frac{M^{2}}{4n}$, where $M$ denotes the maximum value of the weighted terms, and $n$ is the number of samples. Given $V(\theta)$ is the true value function of the target policy $\pi_{\theta}$, applying the concentration inequality for sub-Gaussian variables, the KG-IPS estimator satisfies the following confidence interval with probability at least $1-\delta$:

$$ $\left|\hat{V}_{\text{KG-IPS}}(\theta)-V(\theta)\right|\leq O\left(M\sqrt{\log( 1/\delta)/n}\right).$ $$

The detailed proof and variance analysis are provided in Appendix 9.

To further support our findings, we demonstrate that the optimal policy for the final reward is consistent with the optimal policy for the entity-based knowledge graph reward, which means the non-entity-based LLM reward can be considered as a regularization term that does not affect the optimal policy. See Appendix 10 for a complete analysis.

In the end, we could directly optimize the target policy by maximizing the estimated value function through policy gradient,

$$ $\theta\leftarrow\theta+\nabla_{\theta}\hat{V}_{KG-IPS}(\theta).$ (4) $$

###### Lemma 1 .

###### Lemma 2 .

### 4.2 Knowledge Graph Preference Modeling

Figure: (a) Distribution of QA Accuracy
Refer to caption: x1.png

To facilitate the evaluation and optimization,
we model knowledge graph preference and derive feedback by developing the behavior policy $\mu_{\phi}$ which verbalizes knowledge-graph trajectories.
Randomly sampled trajectories $\mathbb{H}$ from $\mathcal{G}$ in Section 3.2 contain samples that may not be transformed into a chain of thoughts leading to a reasonable question-answering.
Following conventional self-consistent measurement (Wang et al., ; Manakul et al., ), given a sampled trajectory $\bm{h}$ and its verbalized chain of thoughts $\bm{c}$,
we prompt the GPT-4 model to propose a question $q$ related to the first entity $e_{0}\in\bm{h}$
whose answer should be exactly the last entity $e_{T}\in\bm{h}$, and query the GPT-4 model with the proposed question,

$$ $\hat{q}\sim f(q|e_{0},e_{T},\bm{c}),\quad\hat{y}\sim f(y|\hat{q},\bm{c}),\quad R (\bm{h}|\bm{c})=\mathrm{E}\left[\mathbf{1}\left\{e_{T}=\hat{y}\right\}\right],$ $$

where the reward of the trajectory is determined by the answer accuracy.
In Figure 1(a), we present the probability distribution of sampled trajectories,
with respect to the number of correct answers generated per trajectory from ten differently sampled questions associated with each trajectory.
Based on such self-consistency measurement, we estimate the reward function $R(\bm{h}|\bm{c})$ as the normalized question-answering accuracy.
Then we fine-tune the preference policy $\mu_{\phi}$ directly via policy gradient optimization,

$$ $\nabla_{\phi}J(\phi)=\nabla_{\phi}\sum_{k=1}^{K}\sum_{t=0}^{|\bm{c}_{k}|-1}R( \bm{h}_{k}|\bm{c}_{k})\log\mu_{\phi}(y_{t}|q,y_{<t}).$ $$

Based on the distribution of relations (Figure 1(b)) and entities (Figure 1(c)) in the sampled knowledge graph trajectories,
we observe that the relation distribution is relatively more skewed toward the most frequent relations.
This suggests that the verbalized knowledge graph reasoning policy is likely to focus on more frequent reasoning transitions,
potentially enhancing its ability to learn meaningful patterns.
In contrast, the entity distribution shows a relatively short tail,
which may help mitigate the risk of overfitting to specific entities or knowledge biases.

## 5 Experiments

In this section, we evaluate our proposed method, OCEAN, by conducting chain-of-thought alignment on four LLM backbone models and evaluating several downstream tasks.
We show our method’s effectiveness in chain-of-thought alignment and its generalizability in various tasks
to understand (i) whether the proposed optimization approach sufficiently aligns LLMs’ chain-of-thought behaviors with higher estimated values on multi-hop question-answering tasks, (ii) how the proposed method performs on knowledge-intensive question-answering tasks and (iii) whether the post-alignment LLM generalizes on commonsense reasoning tasks.

### 5.1 Implementation Details

Datasets.
Following prior work (), we evaluate our approach across three key aspects of question answering. For knowledge-intensive reasoning, we use datasets requiring deeper domain understanding. The AI2 Reasoning Challenge (ARC) (Clark et al., ) tests models’ advanced reasoning abilities with grade-school science questions. PubMedQA (Jin et al., ) assesses biomedical reasoning by requiring yes/no/maybe answers from research abstracts. Finally, SciQA (Auer et al., ) challenges models to reason over scientific knowledge using the Open Research Knowledge Graph (ORKG).

The second focus is multi-hop reasoning, where models must combine information from multiple sources. HotpotQA (Yang et al., ) requires reasoning across multiple Wikipedia documents. MuSiQue (Trivedi et al., ) includes complex questions needing 2-4 inference hops. StrategyQA (Geva et al., ) tests implicit reasoning, where steps must be inferred. These datasets assess the models’ ability to perform connected reasoning tasks for complex questions.

Lastly, we focus on commonsense reasoning, crucial for understanding everyday knowledge. We evaluate three different commonsenseQA benchmarks, CSQA (Talmor et al., ), CSQA2 (Saha et al., ), and CSQA-COT1000 (Li et al., ). In addition, OpenBookQA (Mihaylov et al., ) tests models’ ability to combine common knowledge with science facts for elementary-level questions. WinoGrande (Sakaguchi et al., ) challenges models to avoid biases in commonsense reasoning.
These tasks evaluate a model’s general abilities in commonsense question-answering.

Baselines.
We experiment with four backbone LLMs:
Gemma-2 (Team, ) with 2B model parameters,
Llama-3 (AI@Meta, ) with 8B model parameters,
Phi-3.5-mini (Abdin et al., ) with 3.8B model parameters,
and Mistral-0.2 (Jiang et al., ) with 7B model parameters.
We use the instruction fine-tuned version of backbone LLMs for better instruction following abilities in question-answering.
For chain-of-thought alignment in OCEAN, we use the CWQ question-answering dataset (Talmor & Berant, ) as the source data,
in which the question-answering pairs are developed from knowledge graphs.
OCEAN only uses CWQ questions for the LLM to generate chain-of-thought reasoning paths, which are further aligned using the knowledge graph preference model,
without directly supervised learning on the ground-truth answers.
To compare with direct supervised learning, we also enable instruction-tuning as a baseline (SFT),
which is fine-tuned with the question as instruction and the answer as the response.

Knowledge Graph Preference Model.
The knowledge graph preference model is developed based on the pre-trained GPT2-Medium model (Radford et al., ).
We collected $6$K question-answering pairs from the Wikidata5M (Wang et al., ) knowledge graph based on the sampling strategy in Section 4.2.
The sampled knowledge graph trajectories are composed into natural language prefixed by the corresponding questions by the GPT-4 model,
which verbalizes the knowledge graph reasoning trajectories and aligns with generative language models’ behaviors.
The model is then fine-tuned with a base learning rate of $1e-4$ for $10$ epochs with a linear learning scheduler.

### 5.2 Multi-hop Question Answering

We evaluate the chain-of-thought reasoning performance of OCEAN compared with base LLMs and supervised fine-tuning (SFT),
in three multi-hop question-answering tasks in Table 1.
Comparing SFT and Base LLMs, we observe similar knowledge inconsistency as in knowledge-intensive tasks.
Although SFT improves on MuSiQue with Gemma-2 and Mistral-0.2 backbones whose base models’ performance is relatively inferior on this task,
such knowledge-inconsistent problems result in worse performance on other downstream tasks.

**Table 1: Comparison results of OCEAN, base LLMs (Base), and supervised fine-tuning (SFT), on three Multi-hop Question-answering tasks. We report with context (w/ ctx) and without context (w/o ctx) answer results with the Exact Match (EM) metric on HotpotQA and the Accuracy metric on StrategyQA. Performance on MuSiQue dataset is EM with context. We also use each test/validation split for each dataset and report policy evaluation $\hat{V}(\theta)$ results. We highlight the best-performed metric in bold font and the second-best underline for each task.**
| Model | Method | HotpotQA | MuSiQue | StrategyQA |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| w/ ctx (%) | w/o ctx (%) | $\hat{V}(\theta)$ | w/ ctx (%) | $\hat{V}(\theta)$ | w/ ctx (%) | w/o ctx (%) | $\hat{V}(\theta)$ |  |  |
| Llama-3 | Base | 32.78 | 33.54 | -10.35 | 11.59 | -9.90 | 77.73 | 59.53 | -9.25 |
| SFT | 8.22 (-24.56) | 16.49 (-17.05) | -22.28 | 1.80 (-9.79) | -17.09 | 66.52 (-11.21) | 51.82 (-7.71) | -15.17 |  |
| OCEAN | 33.38 (+0.6) | 33.75 (+0.21) | -8.10 | 11.67 (+0.08) | -9.77 | 75.40 (-2.33) | 59.83 (+0.3) | -5.53 |  |
| Gemma-2 | Base | 26.33 | 18.58 | -31.88 | 5.84 | -26.41 | 76.71 | 60.99 | -14.06 |
| SFT | 29.75 (+3.42) | 15.91 (-2.67) | -46.92 | 12.53 (+6.69) | -40.25 | 64.77 (-11.94) | 51.97 (-9.02) | -23.27 |  |
| OCEAN | 26.20 (-0.13) | 19.70 (+1.12) | -26.43 | 6.87 (+1.03) | -22.15 | 74.24 (-2.47) | 66.23 (+5.24) | -13.52 |  |
| Phi-3.5 | Base | 32.13 | 26.14 | -19.49 | 11.85 | -15.30 | 73.51 | 58.37 | -13.87 |
| SFT | 21.99 (-10.14) | 7.87 (-18.27) | -44.57 | 6.01 (-5.84) | -42.10 | 63.03 (-10.48) | 50.95 (-7.42) | -21.66 |  |
| OCEAN | 35.13 (+3.0) | 26.23 (+0.09) | -14.84 | 10.82 (-1.03) | -13.47 | 72.20 (-1.31) | 57.64 (-0.73) | -12.25 |  |
| Mistral-0.2 | Base | 26.82 | 28.13 | -19.08 | 5.67 | -6.40 | 79.33 | 58.22 | -11.36 |
| SFT | 20.88 (-5.94) | 14.49 (-13.64) | -18.53 | 7.73 (+2.06) | -12.24 | 52.40 (-26.93) | 51.53 (-6.69) | -15.39 |  |
| OCEAN | 27.24 (+0.42) | 27.54 (-0.59) | -3.12 | 5.15 (-0.52) | -5.94 | 77.29 (-2.04) | 56.62 (-1.6) | -11.21 |  |

Since OCEAN is aligned to incorporate more knowledge-faithful chain-of-thought reasoning patterns
learned from knowledge graph reasoning policy without directly editing its internal knowledge,
OCEAN maintains its generalizability in adapting to downstream tasks.
We observe that OCEAN consistently improves on the policy estimated value $\hat{V}(\theta)$
through direct policy optimization proposed in equation 4, which demonstrates the effectiveness of the developed optimization method.
Regarding the question-answering accuracy, OCEAN improves base LLMs,
which achieves the best performance on HotpotQA and StrategyQA without context.

### 5.3 Knowledge-intensive Question Answering

To understand the effectiveness of OCEAN in knowledge-intensive question-answering tasks,
we show performance comparison with base LLMs (Base) and supervised fine-tuning (SFT) in Table 2.
Comparing SFT and Base LLMs, we observe that directly aligning knowledge graphs with LLMs may suffer from
domain change and knowledge inconsistency when downstream tasks require specific domain knowledge that could conflict with the knowledge graph in the fine-tuning stage.
We also observe that SFT achieves $4.85\%$ and $0.55\%$ average improvements of base models on the PubMedQA dataset,
with and without context respectively, whereas it suffers from $29.60\%$, $8.35\%$, $13.6\%$ average performance decreases on the remaining tasks.
Such significant discrepancies in SFT’s effects on different downstream tasks further show the risk in direct knowledge editing in LLMs.

**Table 2: Comparison results of OCEAN, base LLMs (Base), and supervised fine-tuning (SFT), on three Knowledge-intensive Question-answering tasks. We report answers with context (w/ ctx) and without context (w/o ctx) on Exact Match (EM) metric on PubMedQA and SciQA. The EM performance on ARC dataset is without context. We also use the test/validation split for each dataset to report estimated policy values $\hat{V}(\theta)$. We highlight the best metric in bold font for each task.**
| Model | Method | ARC | PubMedQA | SciQA |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| w/o ctx (%) | $\hat{V}(\theta)$ | w/ ctx (%) | w/o ctx (%) | $\hat{V}(\theta)$ | w/ ctx (%) | w/o ctx (%) | $\hat{V}(\theta)$ |  |  |
| Llama-3 | Base | 79.93 | -10.38 | 63.60 | 58.60 | -25.40 | 83.10 | 57.10 | -22.91 |
| SFT | 61.87 (-18.06) | -18.42 | 75.80 (+12.2) | 58.00 (-0.6) | -26.03 | 67.10 (-16.0) | 35.80 (-21.3) | -23.84 |  |
| OCEAN | 80.60 (+0.67) | -12.45 | 66.00 (+2.4) | 59.80 (+1.2) | -9.37 | 83.20 (+0.1) | 57.70 (+0.6) | -16.63 |  |
| Gemma-2 | Base | 65.89 | -15.36 | 34.40 | 40.60 | -24.61 | 76.50 | 47.10 | -26.60 |
| SFT | 18.06 (-47.83) | -25.22 | 35.60 (+1.2) | 21.00 (-19.6) | -26.55 | 79.80 (+3.3) | 51.50 (+4.4) | -36.61 |  |
| OCEAN | 63.21 (-2.68) | -16.20 | 44.60 (+10.2) | 41.60 (+1.0) | -18.72 | 72.20 (-4.3) | 47.50 (+0.4) | -26.77 |  |
| Phi-3.5 | Base | 87.29 | -7.86 | 70.40 | 41.80 | -28.48 | 83.50 | 58.90 | -14.46 |
| SFT | 65.22 (-22.07) | -9.02 | 62.40 (-8.0) | 50.20 (+8.4) | -28.40 | 76.90 (-6.6) | 43.80 (-15.1) | -14.62 |  |
| OCEAN | 87.63 (+0.34) | -7.94 | 68.40 (-2.0) | 47.60 (+5.8) | -11.45 | 84.70 (+1.2) | 63.50 (+4.6) | -13.40 |  |
| Mistral-0.2 | Base | 73.91 | -9.99 | 51.60 | 36.20 | -13.01 | 78.50 | 58.00 | -11.77 |
| SFT | 43.48 (-30.43) | -13.99 | 65.60 (+14.0) | 50.20 (+14.0) | -21.87 | 64.40 (-14.1) | 35.50 (-22.5) | -21.86 |  |
| OCEAN | 68.90 (-5.01) | -10.89 | 52.60 (+1.0) | 33.20 (-3.0) | -12.42 | 79.10 (+0.6) | 58.40 (+0.4) | -12.00 |  |

With the enhancement of OCEAN, question-answering accuracy of knowledge-intensive tasks generally improved,
while OCEAN fine-tuned LLMs achieving the best performance on all three datasets,
except for PubMedQA without context where SFT achieves better performance due to knowledge transfer from knowledge graph dataset.
We also observe consistent policy value improvement on PubMedQA and SciQA, where the original policy values of base LLMs are relatively lower.
For tasks like ARC, which does not require additional reference knowledge from context and reasoning in an easier chain of thought,
OCEAN still maintains comparable policy value to the base LLM,
which demonstrates the robustness and generalizability of the proposed method.

### 5.4 Commonsense Reasoning

**Table 3: Comparison results of OCEAN, base LLMs (Base), and supervised fine-tuning (SFT), on five Commonsense Reasoning tasks. We report the Exact Match (EM) metric on these tasks and the average performance. We highlight the best method in bold font for each task and LLM.**
| Model | Method | CSQA | CSQA-2 | CSQA-COT1000 | OpenBookQA | Winogrande | Average |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Llama-3 | Base | 65.03 | 71.39 | 69.50 | 58.80 | 43.09 | 61.56 |
| SFT | 51.19 (-13.84) | 57.06 (-14.33) | 49.00 (-20.5) | 63.20 (+4.4) | 34.73 (-8.36) | 51.04 |  |
| OCEAN | 65.03 (0.0) | 68.60 (-2.79) | 72.00 (+2.5) | 60.40 (+1.6) | 41.36 (-1.73) | 61.48 |  |
| Gemma-2 | Base | 57.99 | 62.57 | 63.50 | 51.80 | 49.64 | 57.10 |
| SFT | 14.66 (-43.33) | 65.80 (+3.23) | 15.00 (-48.5) | 7.40 (-44.4) | 50.04 (+0.4) | 30.58 |  |
| OCEAN | 67.73 (+9.74) | 63.56 (+0.99) | 72.50 (+9.0) | 57.20 (+5.4) | 50.12 (+0.48) | 62.22 |  |
| Phi-3.5 | Base | 68.55 | 64.70 | 72.50 | 72.40 | 50.51 | 65.73 |
| SFT | 69.94 (+1.39) | 61.47 (-3.23) | 72.00 (-0.5) | 69.40 (-3.0) | 50.51 (0.0) | 64.66 |  |
| OCEAN | 69.62 (+1.07) | 62.77 (-1.93) | 73.50 (+1.0) | 71.20 (-1.2) | 50.12 (-0.39) | 65.44 |  |
| Mistral-0.2 | Base | 61.18 | 68.48 | 65.00 | 64.00 | 46.25 | 60.98 |
| SFT | 35.87 (-25.31) | 22.47 (-46.01) | 33.00 (-32.0) | 34.80 (-29.2) | 29.12 (-17.13) | 31.05 |  |
| OCEAN | 63.80 (+2.62) | 69.19 (+0.71) | 67.00 (+2.0) | 62.60 (-1.4) | 46.49 (+0.24) | 61.82 |  |

Finally, to demonstrate OCEAN’s generalizability in preserving commonsense knowledge and preventing knowledge catastrophic forgetting (Luo et al., ),
we evaluate OCEAN with base LLMs (Base) and supervised fine-tuning (SFT) on five commonsense reasoning tasks in Table 3.
Since such tasks do not require chain-of-thought generation or external domain knowledge, we only evaluate the accuracy of the model’s generated answers.
We observe the high impact of direct SFT on knowledge graph on LLMs in potential catastrophic forgetting of commonsense knowledge,
especially for the backbone LLMs of Gemma-2 and Mistral-0.2.
In contrast, we show that OCEAN achieves robust performance on commonsense reasoning leverage off-policy evaluation and optimization from knowlege graph’s feedback.
OCEAN manages to maintain comparable performance of base LLMs (*e.g.*, Llama-3 and Phi-3.5), which have strong zero-shot commmonsense reasoning abilities.
In addition, we observe that for base LLM with relatively lower performance (*e.g.*, Gemma-2 and Mistral-0.2), OCEAN enables consistent improvements.
Therefore, OCEAN serves as a robust off-policy alignment paradigm to incorporating knowledge graph reasoning without affecting the generalizability and behaviors of pretrained LLMs.

## 6 Analysis

### 6.1 In-context Learning & Instruction tuning

We conduct further analysis to compare the performance of both the base model and our proposed model in the scenarios of In-Context Learning and instruction fine-tuning. Specifically, we conduct this analysis using the Gemma-2 and Phi-3.5 models across three benchmark datasets: SST2 (Socher et al., ) for sentiment classification, AgNews (Zhang et al., ) for topic classification, and BoolQ (Clark et al., ) for reading comprehension. In the In-context Learning setup, we provide the model with a single example for each task in the prompt. For the Instruction tuning experiments, we apply LoRA (Hu et al., ) to the pre-trained model and fine-tune it on each dataset for 10 epochs. Throughout these experiments, the rank parameter in LoRA is fixed at 16, and we set $\alpha$ in LoRA to 32 across all tasks. The results of the In-context Learning and Instruction Tuning are presented in Table 4. Overall, we observe that the performance of the base model and our proposed model is largely comparable across most scenarios, except in the AG News task with Gemma-2, where OCEAN demonstrates greater performance after instruction tuning.

**Table 4: Performance Comparison of In-Context Learning and Instruction Tuning. All datasets consist of classification tasks or true/false questions, so accuracy is used to evaluate the performance. The performance of the base model and our proposed model is largely comparable.**
| Model | Method | In-Context Learning | Instruction-Tuning |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| SST2 | BoolQ | AG News | Avg. | SST2 | BoolQ | AG News | Avg. |  |  |
| Gemma-2 | Base | 87.16 | 56.12 | 16.47 | 53.25 | 96.21 | 69.03 | 47.03 | 70.76 |
| OCEAN | 89.33 | 55.72 | 13.14 | 52.73 | 96.56 | 68.66 | 60.08 | 75.10 |  |
| Phi-3.5 | Base | 41.28 | 60.06 | 31.89 | 44.41 | 96.44 | 68.13 | 86.43 | 83.67 |
| OCEAN | 40.48 | 59.11 | 32.37 | 43.98 | 96.44 | 68.81 | 86.24 | 83.83 |  |

### 6.2 Evaluation of Generation Quality Post Alignment

Figure: (a) comparison on Self-BLEU
Refer to caption: x4.png

To further evaluate the generation quality of post-alignment LLMs,
we use the Self-BLEU (Zhu et al., ) and Distinct-2 (Li et al., ) scores to evaluate the diversity of the generation,
concerning the similarity between generated texts and the uniqueness of generated 2-gram phrases respectively.
In addition, AlignScore (Zha et al., ) is used to evaluate the faithfulness of the generated answer given the question context.
The results are presented in Figure 2, which show that post-alignment LLMs achieve comparable or better performances in terms of generation diversity and faithfulness.
This demonstrates that while OCEAN aligns chain-of-thought reasoning with KGs, we maintain the text generation qualities of LLMs.

### 6.3 Case Study

In the previous Section 5.2 and 5.3 we observe efficient chain-of-thought alignment with improvement on the estimated policy value $\hat{V}(\theta)$.
To further understand the effects of the alignment, we choose backbone LLMs, Gemma-2 and Llama-3, with significant improvements on $\hat{V}(\theta)$, and perform a sample analysis by comparing the outputs of the base model and OCEAN on the same set of questions.
Our findings demonstrate that the application of our method enhances the precision and conciseness of the chain of thought in the generated responses.
Some illustrative examples are provided in Figure 3.
Specifically, in the first example, the base Llama-3 model incorrectly claims that singing is not a primary action associated with playing the guitar, which leads to an erroneous solution to the question. In contrast, our method enables the model to recognize that singing is a common activity when playing the guitar, while also understanding that making music serves as a broader term. In the second example, although both the base model and OCEAN on Gemma-2 provide reasonable answers to the question, our model demonstrates a more concise chain of thought, streamlining the reasoning process and arriving at the solution with greater simplicity.

Figure: Figure 3: Sample comparison between the base model and OCEAN on Llama-3 and Gemma-2. Our method enables a more precise and concise Chain of thought.
Refer to caption: x7.png

## 7 Conclusion

In this work, we propose OCEAN to address the challenge of offline chain-of-thought evaluation and optimization of LLMs.
By modeling the knowledge-graph preference and deriving feedback by developing a policy that verbalizes knowledge-graph trajectories,
we propose KG-IPS estimator to estimate policy values in the alignment of reasoning paths with knowledge graphs.
Theoretically, we proved the unbiasedness of the KG-IPS estimator and provided a lower bound on its variance.
Empirically, our framework effectively optimizes chain-of-thought reasoning while maintaining LLMs’ general downstream task performance,
offering a promising solution for enhancing reasoning capabilities in large language models.

## 8 Mean Analysis

To prove that the KG-IPS estimator is unbiased, we need to demonstrate that the expected value of the IPS estimator equals the true expected reward under $\pi_{\theta}$.

###### Proof.

The value function of policy $\pi_{\theta}$ can be defined as:

$$ $\displaystyle V(\pi_{\theta})$ $\displaystyle=\mathbb{E}_{a_{t}\sim\pi(\cdot\mid s_{t})}\left[r(s_{t},a_{t})\right]$ $\displaystyle=\frac{1}{N}\sum_{i=1}^{N}\frac{1}{T_{i}|c_{t}^{(i)}|}\sum_{t=1}^ {T_{i}}\mathbb{E}_{e\sim\pi_{\theta}(\cdot\mid s_{t}^{(i)})}\left[r(s_{t}^{(i) },e)\right]$ $$

where $r(s_{t},a_{t})$ represent the reward obtained by taking action $a_{t}$ under state $s_{t}$.

Given that our value function consists of two cases: the first case considers the reward derived from the entity tokens under the knowledge graph preference policy $\mu_{\theta}$, and the second case focuses on the reward derived from the non-entity tokens under the base LLM policy $\pi_{0}$.
We separately prove the unbiasedness by showing that the expected value of either the entity-based or non-entity-based estimators is equal to the true expected reward under their respective policies.

The expected value of the entity tokens in the knowledge graph is:

$$ $\displaystyle\mathbb{E}\left[\hat{V}_{KG}(\theta)\right]$ $\displaystyle=\frac{1}{N}\sum_{i=1}^{N}\frac{1}{T_{i}|c_{t}^{(i)}|}\sum_{t=1}^ {T_{i}}\mathbb{E}_{e\sim\mu_{\phi}(\cdot|s_{t}^{(i)}),e\sim\mathcal{P}(e)} \left[\frac{\pi_{\theta}(e|s_{t}^{(i)})}{\mu_{\phi}(e|s_{t}^{(i)})}\log\pi_{0} (e|s_{t}^{(i)})\right]$ $\displaystyle=\frac{1}{N}\sum_{i=1}^{N}\frac{1}{T_{i}|c_{t}^{(i)}|}\sum_{t=1}^ {T_{i}}\mathbb{E}_{e\sim\mu_{\phi}(\cdot|s_{t}^{(i)})}\left[\frac{\pi_{\theta} (e|s_{t}^{(i)})}{\mu_{\phi}(e|s_{t}^{(i)})}\mathbb{P}(e=\hat{y}|s_{t}^{(i)},e)\right]$ $\displaystyle=\frac{1}{N}\sum_{i=1}^{N}\frac{1}{T_{i}|c_{t}^{(i)}|}\sum_{t=1}^ {T_{i}}\mathbb{E}_{e\sim\pi_{\theta}(\cdot|s_{t}^{(i)})}\left[\mathbb{P}(e= \hat{y}|s_{t}^{(i)},e)\right]$ $\displaystyle=\frac{1}{N}\sum_{i=1}^{N}\frac{1}{T_{i}|c_{t}^{(i)}|}\sum_{t=1}^ {T_{i}}\mathbb{E}_{e\sim\pi_{\theta}(\cdot\mid s_{t}^{(i)})}\left[r(s_{t}^{(i) },e)\right]=V(\pi_{\theta})$ $$

For non-entity tokens, the proof is similar:

$$ $\displaystyle\mathbb{E}\left[\hat{V}_{base}(\theta)\right]$ $\displaystyle=\frac{1}{N}\sum_{i=1}^{N}\frac{1}{T_{i}|c_{t}^{(i)}|}\sum_{t=1}^ {T_{i}}\mathbb{E}_{e\sim\pi_{0}(\cdot|s_{t}^{(i)}),e\sim\mathcal{P}(e)}\left[ \frac{\pi_{\theta}(e|s_{t}^{(i)})}{\pi_{0}(e|s_{t}^{(i)})}\log\pi_{0}(e|s_{t}^ {(i)})\right]$ $\displaystyle=\frac{1}{N}\sum_{i=1}^{N}\frac{1}{T_{i}|c_{t}^{(i)}|}\sum_{t=1}^ {T_{i}}\mathbb{E}_{e\sim\pi_{0}(\cdot|s_{t}^{(i)})}\left[\frac{\pi_{\theta}(e| s_{t}^{(i)})}{\pi_{0}(e|s_{t}^{(i)})}\mathbb{P}(e=\hat{y}|s_{t}^{(i)},e)\right]$ $\displaystyle=\frac{1}{N}\sum_{i=1}^{N}\frac{1}{T_{i}|c_{t}^{(i)}|}\sum_{t=1}^ {T_{i}}\mathbb{E}_{e\sim\pi_{\theta}(\cdot|s_{t}^{(i)})}\left[\mathbb{P}(e= \hat{y}|s_{t}^{(i)},e)\right]=V(\pi_{\theta})$ $$

This completes the proof.
∎

###### Proof.

## 9 Variance Analysis

In this section, we derive a confidence bound from the confidence interval and calculate the lower bound of the variance for the KG-IPS estimator.

###### Proof.

Let $M$ be the maximum value of $\frac{\pi_{\theta}(e|s_{t})}{\mu_{\phi}(e|s_{t})}$ for entity tokens or the maximum value of $\frac{\pi_{\theta}(e|s_{t})}{\pi_{0}(e|s_{t})}$ for non-entity tokens, which represents the largest discrepancy between the target policy and the behavior policy.

In the KG-IPS estimator, each reward is bounded in $[0,1]$. Based on Hoeffding’s Lemma, it is specified to be sub-Gaussian with variance proxy $\frac{(1-0)^{2}}{4}=\frac{1}{4}$.

When scaling the reward by $M$, the resulting variance proxy becomes $\frac{M^{2}}{4}$-sub-Gaussian. The KG-IPS estimator is an average of n sampled terms, applying that the average reward becomes $\frac{M^{2}}{4n}$-sub-Gaussian.

For any sub-Gaussian random variable $X$ with variance $\sigma^{2}$, the concentration inequality holds:

$$ $\left|\hat{X}-\mathbb{E}[X]\right|\leq\sigma\sqrt{2\log\left(\frac{1}{\delta} \right)}.$ $$

Plugging the variables above into the inequality, we get the following bound for the KG-IPS estimator:

$$ $\left|\hat{V}_{\text{KG-IPS}}(\theta)-V(\theta)\right|\leq M\sqrt{\frac{\log(1 /\delta)}{2n}}=O(M\sqrt{\log(1/\delta)/n}),$ $$

with probability at least $1-\delta$.

The variance of a sub-Gaussian random variable is close to its variance proxy, which means the lower bound on the variance of the KG-IPS estimator is $\frac{M^{2}}{4n}$. In addition, by standard concentration inequalities, we can get $O(M\sqrt{\log(1/\delta)/n})$ confidence intervals on our estimator for policy $\pi_{\theta}$.
∎

###### Proof.

## 10 Theoretical Analysis

The value function of policy $\pi_{\theta}$ is defined as:

$$ $V^{\pi_{\theta}}(s_{t},a_{t})=\mathbb{E}_{a_{t}\sim\pi(\cdot\mid s_{t})}\left[ r(s_{t},a_{t})\right].$ $$

Based on our settings, we optimize the target policy for entity tokens aligning with knowledge graph preference. The reward function is formulated as:

$$ $r^{\text{KG}}(s_{t},a_{t})=\sum_{e\in a_{t}}\frac{\pi_{\theta}(e|s_{t})}{\mu_{ \phi}(e|s_{t})}\log\pi_{0}(e|s_{t}),$ (5) $$

where $\mu_{\phi}$ is the knowledge graph preference policy.

To reduce variance, the logged rewards for non-entity tokens under the base LLM policy $\pi_{0}$
are incorporated as a regularization term in the reward function, formulated as:

$$ $r^{\text{reg}}(s_{t},a_{t})=\sum_{v\in c_{t}\backslash a_{t}}\frac{\pi_{\theta }(v|s_{t})}{\pi_{0}(v|s_{t})}\log\pi_{0}(v|s_{t}),$ (6) $$

where $\pi_{0}$ is the base LLM policy.
This helps mitigate disturbances, ensuring the LLM’s behavior on non-entity tokens remains stable and preventing model degeneration.

The final reward is:

$$ $r(s_{t},a_{t})=r^{\text{KG}}(s_{t},a_{t})+r^{\text{reg}}(s_{t},a_{t}),$ (7) $$

Since both $r^{\text{KG}}(s_{t},a_{t})$ and $r^{\text{reg}}(s_{t},a_{t})$ are reweightings of the log-based reward $\log\pi_{0}(v|s_{t})$, they belong to the same equivalence class. By leveraging Lemma 2 from DPO (Rafailov et al., ), we show that the optimal policy for the task-specific reward $r^{\phi}$ aligns with the optimal policy for the final reward $r$. This implies that both rewards induce the same optimal policy.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Could you elaborate on the specific challenges encountered in balancing entity and non-entity token weights in the KG-IPS estimator? How sensitive is the estimator to variations in these weights across different types of reasoning tasks?

2. Since your framework relies on structured knowledge graphs, how does OCEAN perform on CoT reasoning tasks that are less fact-based and more abstract or conceptual? Have you considered extending the model to handle these types of reasoning, and what might be the challenges?

3. Given the framework's reliance on multiple components (e.g., MDP modeling, KG-IPS estimator), what is the additional computational overhead compared to baseline CoT methods? Could you provide some insights into how to make OCEAN more computationally efficient?

4. How well does the OCEAN framework adapt to different types of knowledge graphs (e.g., smaller or domain-specific KGs)? What, if any, modifications are necessary when using a KG that differs substantially in structure or domain from Wikidata5M?

5. In some tasks, you noted the possibility of knowledge conflicts when aligning LLM outputs with knowledge graphs. How does OCEAN address or mitigate conflicts between an LLM's parametric knowledge and KG-based feedback, especially in cases where both knowledge sources might offer valid, but differing, information?

6. The paper mentions using sub-Gaussian concentration inequalities to establish confidence intervals for the KG-IPS estimator. Could you explain why this choice was made and how these intervals compare to other possible statistical measures?
soundness: 3
strengths: 1. The concept of aligning CoT reasoning with knowledge graphs via a KG-IPS estimator is well-motivated, particularly in addressing the challenges associated with human feedback limitations and online deployment costs. The offline nature of the OCEAN framework allows for feasible and efficient model alignment using KGs as a structured feedback source.

2. The paper provides a thorough theoretical foundation for the KG-IPS estimator, proving its unbiasedness and offering a variance lower bound. This adds robustness to the evaluation framework and offers a sound basis for its application in LLM alignment.

3. The authors conduct a wide range of experiments across multi-hop question answering, knowledge-intensive tasks, and commonsense reasoning. These diverse evaluations demonstrate OCEAN's ability to align reasoning paths effectively while maintaining or enhancing the generalization capabilities of the LLMs across domains.
summary: The paper explores a novel approach for enhancing the chain-of-thought (CoT) reasoning in large language models (LLMs) through an offline evaluation and optimization framework. The proposed OCEAN framework utilizes knowledge graphs (KGs) to provide feedback and improve alignment in CoT tasks by modeling reasoning paths as a Markov Decision Process (MDP). By employing a KG-based inverse propensity score (KG-IPS) estimator, the framework enables unbiased offline evaluation, bridging the gap between LLMs and structured KG reasoning to ensure consistent, accurate CoT reasoning without undermining LLMs' performance in general tasks.
weaknesses: 1. While the OCEAN framework is well-conceived, the reliance on multiple components (MDP formulation, KG-IPS estimator, verbalized KG policy) might limit its accessibility for practitioners. Simplifying or modularizing the system, potentially with more explicit guidance for implementation, could improve its applicability.

2. Limited Scope in KG-Based Reasoning Tasks: Although the alignment with KGs enhances factual consistency, the paper does not deeply explore its limitations in more abstract CoT reasoning, where LLMs rely on conceptual, rather than fact-based, reasoning. This could be an area for future work, as abstract reasoning is a critical facet of human-like reasoning in LLMs.

3. The paper focuses on large-scale LLMs, which are often resource-intensive to train and evaluate. While the experiments are comprehensive, further validation on smaller LLMs could strengthen the claim of OCEAN's adaptability across various scales.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: Some questions were raised above. 

1. Why did you choose to 'verbalize' the trajectories generated on the knowledge graphs instead of inversely trying to reformat the LLM reasoning chain as a trajectory using entity and relationship linking? This approach seems initially more intuitive and less likely to cause errors?
soundness: 3
strengths: 1. The paper considers a practical approach for the offline evaluation of LLMs, which is an important area of research given that the dominant approach of RLHF is highly labor intensive and expensive.
2. The notion of using knowledge graphs for LLM alignment is very interesting, and their approach consisting of verbalizing the KG trajectories is novel to the best of my knowledge.
3. The paper contains very detailed and thorough experiments to demonstrate their results empirically.
4. Theoretical proofs look sound.
5. The paper is very well-written and enjoyable to read.
summary: The paper proposes a novel technique for performing offline evaluation of the chain-of-thought capabilities of LLMs using trajectories generated on knowledge graphs. Specifically they propose a well thought out framework, coined OCEAN, which models CoT reasoning in language models as an MDP, where they leverage on-policy reinforcement learning methods to learn a knowledge graph policy which generates likelihood distributions for LLM reasoning paths. They then define and introduce the KG-IPS estimator to evaluate the policy, which they prove is unbiased while also providing a lower bound on its variance. They then discuss how this estimator can be used for policy optimization which would align LLMs to favor the reasoning paths generated in the knowledge graph. Finally, they conduct experiments to evaluate their approach in knowledge-intensive, multi-hop and common sense reasoning tasks compared to various baseline and supervised fine-tuned language models.
weaknesses: 1. One aspect of the methodology that was not clear is the case where there are multiple trajectories that lead to the correct answer, which I believe can often be the case in reasoning tasks solved by LLMs. How will this be represented in the knowledge graph preference policy? Additionally, does the current methodology hold the potential to penalize reasoning paths that are correct but just not identified in a knowledge graph?

2. In a similar vein, it is not clear that the knowledge graph used in this model (which seems to be Wikidata5M) is extensive enough to capture all possible reasoning paths given I believe it only contains entities and relations captured on wikipedia. I think it is imperative that the authors provide more information on the specific knowledge graph they are using and why it is sufficient to model reasoning queries given to LLMs.

3. It is not clear from the tables provided in Section 5 that the OCEAN model significantly outperforms the baseline. A more detailed discussion explaining this seems necessary.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: See weakness.
soundness: 3
strengths: Overall, I find the paper’s idea and methodology compelling. The provided theoretical analysis is nice and the experimental results can support their statements.
summary: This paper focuses on evaluating Chain-of-Thought (CoT) reasoning with guidance from a Knowledge Graph (KG). The authors propose training a KG model as a reference and using KG-IPS to assess the target language model's performance based on the reference model’s output. They conduct extensive experiments to support this approach.
weaknesses: However, improvements in writing would enhance readability and comprehension.

Please use the full name when introducing KG-IPS for the first time.

What role does the base LLM policy $\pi_0$ play? Including a figure to illustrate the entire pipeline—showing the knowledge graph reference policy $\mu_\phi$, target policy $\pi_\theta$, and base LLM policy $\pi_0$—would improve clarity and help readers follow the workflow more easily. 

Could you expand on Equation 3, especially explaining the term $T_i | c_t^{(I)}$ and the intuition behind defining this term after the third summation symbol?

What does the x-axis in Figures 1(b) and 1(c) represent? This clarification is essential to verify the conclusions in the final paragraph of Section 4.

Readers might also wonder about the performance of the reference models and how the architecture of these models impacts the evaluation of the target model.

Finally, how would the reference model perform if it included process-based reward signals rather than relying solely on binary outcome signals defined by answer accuracy? Is a binary signal sufficient, or would process-based feedback add value?

If the author can address my concern, I would like to raise my rating.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: - How does the proposed IPS approach compare to simply doing the direct method (DM) for evaluation?
- What is the offline policy evaluation performance without performing updated according to equation 4? How does the offline policy evlauation change as you update the policy using equation 4?
- How does SFT + normal COT perform?
- How are you computing the value function estimate for SFT and the base model?
- Did you start training OCEAN with the SFT model?
- Did you also train SFT on the reasoning steps, or just on the questions and answers? If it was only on questions and answers, how does it perform when trained on reasoning steps similar to OCEAN?
- How does the index k depend on the question q and answer y for the policy gradient?
- How does SFT + normal CoT perform?
- What does context mean in w/ctx and w/o ctx?

[1] Rewarding Progress: Scaling Automated Process
Verifiers for LLM Reasoning by Setlur et al.
[2] Quiet-STaR: Language Models Can Teach Themselves to
Think Before Speaking by Zelikman et al.
[3] V-STaR: Training Verifiers for Self-Taught Reasoners by Hosseini et al.
soundness: 3
strengths: - The authors are focused on addressing a very important problem.
- The authors performed experiments across three types of datasets and four different models.
- The idea of evaluating multi-reasoning with a multi-hop knowlege graph is interesting.
summary: Evaluating the chain-of-thought (CoT) reasoning of large language models (LLMs) is challenging due to the difficulty in gathering accurate labels. Since CoT reasoning involves a sequential decision-making process, using traditional feedback methods like reinforcement learning from human feedback (RLHF) becomes impractical. The author leverages knowledge graphs to address the challenge of evaluating LLMs CoT offline. The authors propose aligning LLMs' multi-step reasoning with multi-hop knowledge graph trajectories using inverse propensity scores (IPS). Specifically, the author proposes OCEAN, an estimator that evaluates a model's reasoning ability using CoTs from knowledge graphs. The authors demonstrated that their proposed offline evaluation method performed well across several reasoning tasks, including multi-hop question answering, knowledge-intensive question answering, and commonsense reasoning.
weaknesses: - Reasoning only requires the model to generate sequences of thoughts with itself, which means we can optimize these thoughts online rather than offline [1,2,3].
- The ability to evaluate a model of multi-step reasoning based only on a multi-hop knowledge graph is very limiting.
- The paper lacks ablation experiments to clearly explain the performance gains of OCEAN, which come from the objective in equation 3 and training on GPT-4's verbalized structured chain of thoughts.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: # Task formulation
- Line 156, what is the "surrounding context $c_t  \backslash a_t$"? Could you please elaborate on why $a_t$ "is a sub-sequence in the reasoning path $a_t \in c_t$,"? Providing some intuitive examples may help improve understanding.
- In line 169, "the current edge $e_{t-1}$" implies that $e$ represents edges in graphs. However, in line 172, "the entity $e_{t-1}$" implies that it represents entities. Could you please clarify the meaning of the notation $e$? The notation $e$ also appears in Equation 3. Please see related questions below.

# Methodology
- As most of readers would be not familiar with the Wikidata5M dataset, could you please present some cases of sampled trajectories $h$ from the knowledge graph $\mathcal{G}$? In addition, some cases of the verbalizing transformation from $h$ to $c$ using GPT-4 would also be very interesting to see.
- What are the transition dynamics in the exploration of the knowledge graph?
- Could you please provide more detailed information about the resulting offline dataset used for knowledge graph preference modeling?
- Could you please clarify whether the training of the knowledge graph preference policy is conducted in an online or offline manner?
- About Equation 3:
  - $e \in c_t$ implies that $e$ is a sub-set of $c_t$. Could you please explain a little bit on the relation between these three notations: $e$, $a_t$, and $c_t$? Could we simply understand that $e$ represents a sequence of tokens?
  - $c_t^{(i)}$ is outside of the summation $\sum\limits_{t=1}^{T_i}$. The subscript $t$ for this variable is unassigned. 
  - Should $e \in c_t$ be $e \in c_t^{(i)}$?
- Line 221, is it correct that the reward function is defined as $r_t^{(i)} = \log \pi_0(e|s_t^{(i)})$ for the training of the target policy? If so, why we consider the reward in the token level, instead of the reward in the reasoning step level? Utilizing the reward function at the token level conflicts with the formulation in line 159, where $r_t = r(s_t, c_t)$. 

# Minor suggestions
- Some background knowledge about the conventional inverse propensity scores (IPS) method could be interesting for readers.
soundness: 3
strengths: - The research goal of aligning LLMs with chain-of-thought reasoning is both interesting and important.
- The approach of using graphs as supplementary information to align LLMs with chain-of-thought reasoning is a promising avenue for exploration.
- The paper is clearly presented and easy to follow, making it a pleasure to read overall.
- The experiments are thorough and extensive.
summary: This paper introduces a novel approach for aligning LLMs with chain-of-thought reasoning by leveraging offline knowledge graph data. Specifically, a knowledge graph preference policy is first trained using reinforcement learning on question-answering data derived from verbalized exploration trajectories within the offline knowledge graph. This policy, which serves as a proxy for the knowledge graph’s preferences, is then used to fine-tune the target LLM policy through a novel KG-IPS estimator that assesses the difference between the proxy knowledge graph preference policy and the target policy. Extensive experiments are conducted to evaluate both the chain-of-thought reasoning alignment performance and the impact of this alignment on the base LLM's general capabilities.
weaknesses: Several key concepts and notations in the task formulation and methodology are not clearly defined, and certain details about the dataset are missing. Please refer to the questions for further details.

Meta-review: This paper proposes a way to do an offline evaluation of the chain-of-thought capabilities of LLM using a knowledge graph and an inverse-propensity-based metric. Reviewers are generally positive on this one and found the idea of aligning knowledge graph reasoning and chain-of-thought to be interesting. Reviewers also appreciated the theoretical claims behind the proposed IPS approach. The main concerns were whether knowledge-graphs can capture abstract reasoning (reviewer QYsv), key concepts not explained (reviewer 3w4e), and whether experiments are convincing (reviewer wXum).

Personally, I found the paper hard to parse and agree with reviewer 3w4e's initial assessment. To begin with, the data generation process isn't clearly explained. Reading the unbiasedness proof of KG-ILP, it appears that certain tokens are generated by the knowledge graph policy $\mu_\phi$ and others by the LLM policy $\pi_0$, however, exactly how isn't very clear.  The notations don't make sense in Line 867 and Line 881 which show e being sampled twice. Further, why is the probability of token $v \in c^{(i)}_t$ in Equation 3 not dependent on the previous tokens in $c^{(i)}_t$. It isn't explained on line 211 why log-likelihood of a token under the base policy is used as the reward. Overall, the structure of the unbiasedness and variance is a standard analysis for IPS metrics. The result for variance shows, if anything, that it can be high if $M$ is high which is undesirable, although high variance is a general problem with IPS metrics. One way to fix this would be to do doubly-robust estimation using a direct method along with IPS (e.g., see Dudik et al., 2011 which is cited).

Overall, this one is hard to place since I think the paper isn't well presented despite generally positive sentiment from the reviewers. I think the biggest reason for acceptance is the innovativeness of the idea and positive results. For this reason, I am recommending a weak accept but wouldn't mind if the paper is not accepted this time. If the paper does get accepted, I would request authors to include a figure or explain in detail how the offline data is generated and motivate why different decisions are made (e.g., why is reward function given by log-likelihood under $\pi_0$).

Decision: Accept (Poster)