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

Junda Wu, Xintong Li, Ruoyu Wang, Yu Xia, Yuxin Xiong, Jianing Wang, Tong Yu, Xiang Chen, Branislav Kveton, Lina Yao, Jingbo Shang, Julian McAuley

TL;DR

OCEAN is a novel framework for evaluating and aligning the "Chain-of-Thought" (CoT) reasoning of Large Language Models (LLMs) using Knowledge Graphs (KGs) as an offline source of truth. By treating reasoning as a Markov Decision Process and using a specialized "Inverse Propensity Scoring" estimator, OCEAN allows researchers to fine-tune models for better factual accuracy and more concise reasoning without the need for expensive human feedback or risky online testing.

What Problem Does This Paper Address?

While LLMs are famous for "Chain-of-Thought" reasoning—breaking down complex problems into steps—they often suffer from hallucinations or "unfaithful" reasoning where the intermediate steps are factually wrong even if the final answer is correct.

Aligning these models typically requires Reinforcement Learning from Human Feedback (RLHF), which is incredibly expensive and difficult to scale for multi-step reasoning. Furthermore, evaluating models "online" (interacting with them in real-time) is costly and potentially risky. There is also a "heterogeneity" problem: Knowledge Graphs (like Wikidata) are structured and precise, while LLM reasoning is unstructured and "free-form." Bridging the gap between these two formats to provide automated, factual feedback is the core challenge OCEAN tackles.

Methods

The authors propose the OCEAN (Offline Chain-of-thought Evaluation and AligNment) framework, which consists of three main stages:

Extracted figure from page 5 of OCEAN: Offline Chain-of-thought Evaluation and Alignment in Large Language Models
  1. Verbalized KG Reasoning: To bridge the gap between structured data and natural language, the researchers use a "verbalizer" (GPT-4) to turn Knowledge Graph triplets into natural language reasoning paths. This makes the KG "speak the same language" as the LLM.
  2. Knowledge Graph Preference Modeling: They train a small language model (GPT2-Medium) on these verbalized paths to act as a "behavior policy." This model represents the "ideal" factual reasoning preferred by the Knowledge Graph.
  3. The KG-IPS Estimator: This is the mathematical heart of the paper. They use Inverse Propensity Scoring (IPS)—a technique from causal inference—to estimate how well a target LLM aligns with the KG's preferences using only offline, logged data.
  4. Innovation: They split the scoring. "Entity tokens" (factual names/relations) are weighted against the KG preference model, while "non-entity tokens" (connective text) are weighted against the base LLM. This prevents the model from "degeneration" (losing its ability to speak naturally while trying to be factually correct).

Finally, they use this estimator to perform offline policy optimization, essentially fine-tuning the LLM to maximize its alignment with the factual paths found in the KG.

Datasets

The study utilizes a wide array of benchmarks to test different types of reasoning:

  • Knowledge Source: Wikidata5M.
  • Training Data: Complex Web Questions (CWQ).
  • Multi-hop Reasoning: HotpotQA, MuSiQue, and StrategyQA.
  • Knowledge-Intensive Tasks: ARC (science questions), PubMedQA (biomedical), and SciQA.
  • Commonsense Reasoning: CSQA, CSQA2, OpenBookQA, and WinoGrande.

Key Insights and Ideas

  • Unbiased Offline Evaluation: The authors provide a theoretical proof that their KG-IPS estimator is unbiased. This means we can reliably predict how a model will perform on factual reasoning tasks without ever actually running it on those tasks during the alignment phase.
  • Avoiding the "Alignment Tax": Often, when you fine-tune a model to be better at one thing (like facts), it gets worse at others (like general conversation). OCEAN’s split-weighting approach (entity vs. non-entity) allows the model to gain factual precision without losing its general linguistic capabilities.
  • Conciseness as a Virtue: A surprising qualitative finding was that OCEAN-aligned models didn't just become more accurate; they became more concise. They learned to stop "rambling" and get to the factual point faster, mimicking the efficiency of a graph traversal.

Key Results

The researchers tested OCEAN on four major backbone models: Gemma-2 (2B), Llama-3 (8B), Phi-3.5-mini (3.8B), and Mistral-0.2 (7B).

  • Superiority over SFT: Standard Supervised Fine-Tuning (SFT) often caused models to "forget" general knowledge (catastrophic forgetting). OCEAN consistently outperformed SFT across multi-hop and knowledge-intensive tasks.
  • Accuracy Gains: On the StrategyQA dataset, OCEAN-aligned Llama-3 achieved significantly higher estimated policy values compared to the base model, indicating much tighter alignment with factual truth.
  • Robustness: In commonsense tasks (where KGs aren't always helpful), OCEAN maintained the base model's performance, whereas SFT often caused performance to crater. For example, Gemma-2's accuracy on CSQA dropped from nearly 58% to under 15% with SFT, but improved to over 67% with OCEAN.
Extracted table from page 8 of OCEAN: Offline Chain-of-thought Evaluation and Alignment in Large Language Models

Limitations and Future Work

  • Knowledge Graph Dependency: OCEAN is only as good as the Knowledge Graph it uses. If a fact isn't in Wikidata5M, the system cannot provide feedback on it.
  • Initial Setup Cost: While the evaluation is offline and efficient, the initial step of "verbalizing" the KG using a high-end model like GPT-4 carries an upfront computational or API cost.
  • Future Directions: The authors suggest extending this to multi-agent systems, where different LLM "agents" could use OCEAN to evaluate each other's reasoning paths in complex, collaborative environments.
BibTeX
@inproceedings{wu2025ocean,
  author = {Junda Wu and Xintong Li and Ruoyu Wang and Yu Xia and Yuxin Xiong and Jianing Wang and Tong Yu and Xiang Chen and Branislav Kveton and Lina Yao and Jingbo Shang and Julian McAuley},
  title = {{OCEAN: Offline Chain-of-thought Evaluation and Alignment in Large Language Models}},
  booktitle = {International Conference on Learning Representations},
  year = {2025},
  url = {https://openreview.net/forum?id=rlgplAuN2p},
}