Source-linked AI summary
Zebra-CoT: A Dataset for Interleaved Vision Language Reasoning
Ang Li, Charles Wang, Deqing Fu, Kaiyu Yue, Zikui Cai, Wang Bill Zhu, Ollie Liu, Peng Guo, Willie Neiswanger, Furong Huang, Tom Goldstein, Micah Goldblum
TL;DR
Visual CoT training is limited by weak native visual reasoning and a shortage of diverse, logically coherent interleaved text-image data. Zebra-CoT addresses this gap with a 182,384-sample dataset spanning four broad task categories, and fine-tuning improves benchmark performance while enabling Bagel-7B to generate visual reasoning chains.
Problem
High-quality, diverse datasets with strong logical connections between textual and visual reasoning are lacking, limiting general visual-CoT training.
Method
Zebra-CoT curates interleaved text-image reasoning traces from real-world and synthetic sources across scientific, 2D, 3D, and visual-logic tasks.
Results
4.9% average improvement across seven challenging benchmarks is achieved after Anole-7B fine-tuning, while Bagel-7B learns to generate visual CoT inherently.
Takeaways & Limitations
Zebra-CoT provides a broad training resource for developing and evaluating multimodal models that reason with visual aids.
Abstract
from arXiv · showhide
Humans often use visual aids, for example diagrams or sketches, when solving complex problems. Training multimodal models to do the same, known as Visual Chain of Thought (Visual CoT), is challenging due to: (1) poor off-the-shelf visual CoT performance, which hinders reinforcement learning, and (2) the lack of high-quality visual CoT training data. We introduce $\textbf{Zebra-CoT}$, a diverse large-scale dataset with 182,384 samples, containing logically coherent interleaved text-image reasoning traces. We focus on four categories of tasks where sketching or visual reasoning is especially natural, spanning scientific questions such as geometry, physics, and algorithms; 2D visual reasoning tasks like visual search and jigsaw puzzles; 3D reasoning tasks including 3D multi-hop inference, embodied and robot planning; visual logic problems and strategic games like chess. Fine-tuning the Anole-7B model on the Zebra-CoT training corpus results in an improvement of +12% in our test-set accuracy and yields up to +13% performance gain on standard VLM benchmark evaluations. Fine-tuning Bagel-7B yields a model that generates high-quality interleaved visual reasoning chains, underscoring Zebra-CoT's effectiveness for developing multimodal reasoning abilities. We open-source our dataset and models to support development and evaluation of visual CoT.
1 Introduction
Zebra-CoT addresses the challenge of training models to reason with interleaved text and visual aids by providing a diverse, logically coherent dataset. It spans broad task categories and improves multimodal reasoning through model fine-tuning and visual-CoT generation.
- Motivation: Visual Chain of Thought remains challenging because models predominantly produce textual reasoning and high-quality, diverse interleaved training data is scarce.These limitations hinder general foundation models’ visual-CoT capabilities and reinforcement-learning approaches.
- Dataset: Zebra-CoT contains 182,384 interleaved text-image reasoning traces spanning four major categories and 18 subcategories.The categories include scientific questions, 2D visual reasoning, 3D visual reasoning, and visual logic and strategic games.
- Evaluation: 31.51% average accuracy for three frontier models rises to 47.99% with one multimodal CoT step and 56.70% with two steps.The improvements are +16.48 points and +25.19 points overall, with gains up to +43.77 points in specific domains.
- Evaluation: 4.9% average improvement across seven challenging benchmarks is achieved after fine-tuning Anole-7B, with a maximum 13.1% gain on a visual logic benchmark.The in-distribution test accuracy increases from 4.2% to 16.9%.
- Model Capability: Fine-tuning Bagel-7B enables it to generate high-quality visual CoT during its own reasoning process, unlike the base model.The resulting capability is positioned as useful for future reinforcement-learning training.
2 Related Work
Prior work includes visual-programming agents, visual reasoning datasets, and large interleaved corpora, but these resources generally do not provide broad, logically connected visual reasoning traces.
- Visual Chain of Thought: Visual CoT research has largely used visual programming and external tools to generate sketches, bounding boxes, and other visual outputs.Visual Sketchpad is described as a versatile open-source agent handling a wide range of tasks.
- Visual Reasoning Datasets: Most visual reasoning datasets keep multimodal information in the input while leaving reasoning traces purely textual.VISUAL-COT is identified as the only open-source dataset with interleaved text-image reasoning among the cited visual reasoning datasets.
- Interleaved Datasets: Existing interleaved text-image datasets primarily use images for recognition, captioning, or supplementary purposes rather than explicit visual reasoning.This limits their suitability as sources of logically structured visual-CoT traces.
3 Data Curation Details and Compositions
Zebra-CoT is curated to connect text and images through coherent reasoning across scientific, 2D, 3D, and visual-logic tasks. Its examples include visual aids such as diagrams, bounding boxes, zooms, plans, and intermediate scene transformations.
- Curation Motivation: Existing datasets often lack strong logical connections between text and images, while current visual-CoT datasets cover too few domains for broad generalization.The curation process therefore targets both modality coherence and category diversity.
- Data Sources: Real-world problems are sourced from mathematics, physics, coding, chess, and other resources, then their multimodal reasoning traces are extracted and cleaned.Raw traces may require repair because image references do not clearly map to the corresponding images.
- Data Processing: Frontier VLMs fill template placeholders, enhance reasoning traces, and produce textual narratives containing image placeholders before invalid cases are filtered.Filtering removes cases such as duplicate references to one image and unreferenced image placeholders.
- Scientific Questions: Scientific visual reasoning uses auxiliary lines, free-body diagrams, and sketches across geometry, physics, chemistry, algorithms, and graph problems.These visual aids support abstraction, interpretation, and step-by-step problem solving.
- 2D Visual Reasoning: 2D visual reasoning includes visual search and jigsaw tasks using bounding-box drawing, zooming, and visual comparison of image regions.Jigsaw examples require matching candidate pieces to missing regions in a complete image.
- 3D Visual Reasoning: 3D visual reasoning covers embodied planning, robot action sequences, viewpoint transformations, and multihop object-counting operations.ALFRED is reformulated as image-goal-conditioned planning, while examples show iterative scene updates after object-removal instructions.
- Visual Logic and Strategic Games: Visual logic and strategic-game tasks include IQ matrices, Tetris, ciphers, ARC-AGI, and chess, where visual reasoning can preserve spatial relationships lost through verbalization.Chess examples evaluate candidate moves while tracking pins, attacks, and defensive relationships.
4 Analysis of Zebra-CoT and the Value of visual CoT
Zebra-CoT tasks challenge frontier multimodal models in zero-shot evaluation, while supplying early interleaved reasoning steps substantially improves accuracy. Gains vary by task type, with especially large improvements on maze and graph reasoning.
- Frontier-model evaluation: Frontier models perform poorly on Zebra-CoT tasks despite their advanced multimodal reasoning capabilities.The evaluation uses question-only zero-shot inputs before scaffolding is introduced.
- Scaffolding experiment: Zebra-CoT tasks can require up to 20 images across multiple multimodal reasoning steps, so two-step scaffolding still leaves substantial reasoning to the model.The provided steps are intended as guidance rather than revealing the solution.
- Scaffolding experiment: 47.99% average accuracy follows one multimodal reasoning step, rising to 56.70% after two steps, for gains of +16.48% and +25.19%.The scaffolding settings provide the first one or two multimodal reasoning steps as context.
- Task-level variation: Maze accuracy increases from 52.59% to 76.60% and then 96.36%, corresponding to gains of +24.01% and +43.77% on average.Maze tasks show the most dramatic improvements among the reported task types.
- Task-level variation: Graph reasoning improves from 3.92% to 22.03% with two multimodal reasoning steps, while multihop object counting reaches 67.65% from a 59.40% baseline.Performance gains vary across task types but generally follow an improvement trend.
5 Training Models on Zebra-CoT
The paper fine-tunes Anole-7B and Bagel-7B on Zebra-CoT to test whether interleaved visual reasoning can improve benchmark performance and become an intrinsic generation capability. Both experiments report substantial benefits, including stronger benchmark results for Anole and native visual CoT generation for Bagel.
- Anole-Zebra-CoT: Training Anole-7B on Zebra-CoT increases in-distribution accuracy from 4.2% to 16.9%.This corresponds to a 4 times relative performance improvement and a 12% gain in accuracy.
- Bagel-Zebra-CoT: Bagel-7B is fine-tuned end-to-end on Zebra-CoT despite lacking native interleaved text-image generation in its original implementation.The training uses packed sequences and caps images at a minimum-side resolution of 512.
- Bagel-Zebra-CoT: After fine-tuning, Bagel-Zebra-CoT inherently generates high-quality visual CoT during reasoning, including on tasks outside its training distribution.The paper presents representative generated reasoning traces qualitatively.
6 Conclusion & Future Directions
The paper concludes that Zebra-CoT provides a large, diverse foundation for training models that reason with interleaved text and images. It identifies Bagel-Zebra-CoT as a promising initialization for future reinforcement-learning research on visual reasoning.
- Conclusion: Zebra-CoT contains 182K interleaved text-image reasoning traces spanning four major categories, 18 domains, and over 50 distinct tasks.The dataset is presented as a large-scale resource for visual chain-of-thought research.
- Conclusion: Fine-tuning demonstrates that Anole-7B improves across seven challenging benchmarks, while Bagel-7B learns to generate visual aids during problem solving.The reported Anole result averages a 4.9% gain and reaches up to 13.1% on visual logic tasks.
- Future directions: Bagel-Zebra-CoT is proposed as a strong initialization for future reinforcement-learning fine-tuning with verifiable or fine-grained rewards.The future-work discussion connects native visual-thought generation with reinforcement-learning research.
- Conclusion: The authors position Zebra-CoT as a step toward AI systems that use diagrams, graphs, and spatial reasoning during complex problem solving.This conclusion stays within the paper’s stated goal for visually oriented reasoning systems.
A.1 Data Statistics.
The appendix introduces ZEBRA-COT’s data statistics and points readers to a dedicated statistics table.
- Data Statistics: The appendix provides detailed statistics about ZEBRA-COT’s categories.
- Data Statistics: The category statistics are presented as part of the appendix’s data-statistics material.
- Data Statistics: Table 3 is identified as the statistics table for ZEBRA-COT.
A.2 Scientific Questions
Scientific Questions uses visual artifacts to ground reasoning in geometry, physics, chemistry, algorithms, and graphs. The section combines rendered diagrams and interleaved explanations across diverse problem sources.
- Geometry: Geometry samples replace plotting code with rendered problem and solution images, making visual reasoning an explicit training chain.The collection includes 1,061 MATH train-split samples, including geometry and other mathematical subcategories.
- Physics: Physics samples cover sketch-friendly classical mechanics problems generated from parametric templates and openly licensed educational resources.Examples include free-body, motion, circuit, and ray diagrams, with physically plausible parameters sampled from predefined ranges.
- Chemistry: Chemistry traces depict each reactant separately, combine both reactants, and show the resulting product for 4,700 USPTO-50K reactions.The subset contains two-to-one reactions filtered for distinct reactants and single products.
- Algorithmic problem solving: Algorithmic traces use compact visual scaffolds and an iterative visual-sketchpad loop to connect problem statements, diagrams, and explanations.The corpus includes 1,200 algorithm-based problems spanning competitive-programming abstractions such as logistics, routing, and flow optimization.
- Graph problems: Graph reasoning covers connectivity, shortest path, minimum spanning tree, and topological sort with coupled problem and reasoning images.The collection contains 10,000 graph problems, about 2,500 per task, with up to 19 reasoning images per sample.
- 2D visual reasoning: Visual search traces use bounding boxes or zooming across chart, document, relation, and general VQA tasks, while jigsaw traces iteratively test candidate missing-piece sets.The jigsaw procedure compares reconstructed images or imagined completions against the original puzzle.
A.5 Visual Logic and Strategic Games
Visual logic and strategic-game tasks represent transformations, spatial patterns, navigation, and competitive decisions as interleaved visual reasoning traces. The section spans Tetris, mazes, chess, Connect Four, checkers, ciphers, and grid-based geometry problems.
- Tetris: Tetris tasks require selecting tetromino sets, completing partially occupied grids, or applying spatial transformations to shapes.The transformation task includes scaling, rotation, and translation operations.
- Mazes: Maze tasks separate visual perception from planning through raster inputs and vector overlays for paths, trajectories, heat maps, and landmarks.Evaluation covers topology, pathfinding, navigation planning, and coverage problems.
- Chess: Chess traces compare independently rendered legal move options, with Stockfish identifying the optimal move from mid-game Lichess positions.The example evaluates Ba2 as strongest among four options.
- Strategic games: Connect-Four and checkers traces use minimax-based game play to supervise short-horizon planning and tactical motifs such as forced captures and multi-jump chains.Connect Four contains 2,100 traces and renders every drop on a 7 × 6 board.
- Ciphers: Cipher traces visualize encoding rules through grids, rails, columns, coordinates, and digraph geometry for Scytale, Polybius, Rail Fence, and Playfair ciphers.Text explains the applied symbolic rule alongside each visual transformation.
- Visual logic: Grid-based geometry examples reason over symmetry lines and cube angles by decomposing visual structures into explicit geometric relationships.The cube example identifies four right angles and sums them to 360°.
Visual CoT Example: MathVista
The MathVista examples show visual CoT combining diagram inspection with geometric deductions and focused image examination. Reasoning traces derive answers from rendered structures and targeted visual attention.
- Geometry: The circle example uses the diameter condition to establish a right angle, then derives the remaining angle as 50 degrees.The reasoning introduces the radius segment OC and uses triangle relationships.
- Visual inspection: The image-counting example zooms into the bench area before concluding that four children are sitting there.The answer depends on focused inspection of the relevant image region.
D Bagel Performance Analysis
Bagel fine-tuning enables high-quality visual CoT generation, but benchmark performance does not substantially improve and can slightly decline because generated images may hallucinate visual changes.
- Performance limitation: Bagel shows no substantial benchmark improvement over its pure-text original and slightly declines on tasks such as MathVista.The analysis attributes the issue to generated-image hallucinations being re-encoded as inaccurate visual information.
E Scaffolding Results Breakdown
The scaffolding evaluation compares question-only prompts with prompts that include the first one or two multimodal reasoning steps. The supplied passages also specify how these traces are constructed, emphasizing coherent text-image progression and preservation of original answers where applicable.
- Scaffolding Results Breakdown: Scaffolding compares zero-shot question-only evaluation with prompts containing the first multimodal reasoning step or first two steps.The table caption labels these settings Q, 1MT, and 2MT, with delta columns measuring absolute improvement over Q.
- Trace Structure: The trace format presents a problem statement, an interleaved reasoning trace, and a final answer.Reasoning steps may include image placeholders between textual thoughts, followed by a final calculated answer or preserved original answer.
- Trace Construction: Trace enhancement preserves the original correct reasoning and final answer while adding detail, explanations, and multimodal coherence.The instructions distinguish enhancement from generating a new trace and require the original images to be used.
- Multimodal Reasoning Flow: Each visual reasoning step should explain why an image is needed and what insight it contributes to the solution.The guidelines require each step to build logically on previous textual and visual reasoning steps.
- Image References: Image placeholders identify original problem images and generated reasoning images, with each unique image represented once.Explanations should refer to images naturally rather than repeating placeholder tags.