Multimodal Chain-of-Thought Reasoning in Language Models

In plain words

Large language models can reason step by step, writing out intermediate steps before the final answer. Small models often write steps that sound plausible and are wrong, and that makes their answers worse. This paper asks what happens when the question also comes with a picture, such as a school science question with a diagram.

The authors split the job into two steps, each done by its own model of the same design. The first model writes the reasoning from the text and the picture. The second model reads the question together with that written reasoning and picks the answer. The picture enters through features from a frozen image encoder, so the image encoder itself is not trained.

The authors trace most of the failures to made-up reasoning, and they link that to the missing picture. With picture features, the reasoning gets better and the answers get more accurate. The models are small, under a billion parameters, so they can be fine-tuned and run on consumer-grade graphics cards. Treat the benchmark numbers as the authors' own.

TL;DR

• Two T5 models, both initialised from FLAN-Alpaca and trained separately. Stage 1 maps question, context, options and image to a rationale; stage 2 maps the same input plus that rationale to the answer §4.1 tab. 4.
• The image enters as frozen ViT-L/32 patch features (145 tokens of width 1024). They are projected to the model width, attended to by the text tokens with one attention head, and mixed back in through a sigmoid gate §4.2 tab. 12.
• Author-reported ScienceQA test accuracy: 85.31% (Base) and 90.45% (Large). Without picture features, the two-stage text-only model reaches 78.57% tab. 3.
• Picture features correct 60.7% of the hallucinated-rationale mistakes §3.3 fig. 3.
• Read the Pitfalls: the text, appendix and released scripts disagree on learning rate, epochs and input length, and the released answer-stage command calls a file that is not in the repository.

Problem

Chain-of-thought (CoT) prompting helps large language models, but the paper notes that models under 100 billion parameters often produce illogical reasoning chains that lead to wrong answers §2. It targets models under 1 billion parameters, which can be fine-tuned and deployed on consumer-grade GPUs (e.g. 32 GB of memory) §3.

In a pilot with a 200M-parameter text-only model, asking for the rationale before the answer lowers accuracy from 81.63% (no CoT) to 69.32%, and asking for it after the answer gives 69.68% tab. 2. In a two-stage text-only setup, the authors read 56% of 50 sampled error cases as rationales that mislead the answer step §3.2 fig. 3. Adding captions gives a small gain (78.57% to 79.37%); adding ViT picture features gives 85.31% tab. 3.

The method needs annotated rationales for training. ScienceQA provides them §5.1 app. B.2. The paper does not discuss a setting without them, beyond one experiment with rationales generated by larger models §6.2 tab. 7.

Method

The model runs two stages, each with the same architecture but separate weights and separate training data §4.1. Stage 1 (rationale generation) takes the text input X_language and the image X_vision and generates the rationale R. Stage 2 (answer inference) appends R to the text, keeps the same image, and generates the answer A §4.1. At test time, the stage-1 model produces rationales for the test set, and the stage-2 model consumes them §4.1.

The core architecture has three parts: encoding, interaction, and decoding §4.2.

1. Encode the text with the T5 encoder; use the last layer's hidden states as H_language, shape n by d §4.2 eq. 2.
2. Encode the image with a frozen ViT and project its patch features to the text width d with a learnable matrix W_h, giving H_vision, shape m by d §4.2 eq. 3. Questions without an image use an all-zero block of the same shape §4.2.
3. Interaction: a single-head attention uses H_language as queries and H_vision as keys and values, scaled by the square root of d §4.2.
4. Gated fusion: a sigmoid gate, computed per token and per dimension from both streams, decides how much of the attended picture goes into each position §4.2 eq. 4 eq. 5.
5. Decoding: the T5 decoder reads the fused states and generates the target Y, which is the rationale in stage 1 and the answer in stage 2 §4.2.

Training maximises the likelihood of the target tokens, one stage at a time eq. 1 §4.1. Both stages use the annotated elements from the training set: X to R in stage 1, and X with R to A in stage 2 §4.1.

p(Y|X_{\text{language}},X_{\text{vision}}) = \prod_{i=1}^{N} p_{\theta}\left(Y_{i} \mid X_{\text{language}}, X_{\text{vision}}, Y_{<i}\right) eq. 1

H_{\text{language}} = \text{LanguageEncoder}(X_{\text{language}}) eq. 2

H_{\text{vision}} = W_{h} \cdot \text{VisionExtractor}(X_{\text{vision}}) eq. 3

H^{\text{attn}}_{\text{vision}} = \text{Softmax}\left(\frac{Q K^{\top}}{\sqrt{d_k}}\right) V, \quad Q = H_{\text{language}},\; K = V = H_{\text{vision}} §4.2, unnumbered in the compiled paper

\lambda = \text{Sigmoid}(W_{l} H_{\text{language}} + W_{v} H^{\text{attn}}_{\text{vision}}) eq. 4

H_{\text{fuse}} = (1 - \lambda) \cdot H_{\text{language}} + \lambda \cdot H^{\text{attn}}_{\text{vision}} eq. 5

The LaTeX source contains an algorithm box, but its include line is commented out, so it does not appear in the compiled paper and is not cited as alg. 1. The sketches below are ours, derived from §4.1 and §4.2.

# Two-stage inference  [§4.1]
# Input:  text (question, context, options), image or None
# Output: rationale text, answer

def multimodal_cot(text, image):
    # Stage 1: write the rationale from text and picture  [§4.1]
    rationale = rationale_model.generate(text, image)
    # Stage 2: append the rationale to the text; the image is unchanged  [§4.1]
    text_with_rationale = text + " " + rationale
    # Stage 2 answer, parsed as "The answer is (X)"  [code: main.py#L154]
    answer = answer_model.generate(text_with_rationale, image)
    return rationale, answer

# One pass of the fused encoder  [§4.2]
# text_tokens: (n,)   patch_feats: (m, 1024), frozen ViT output  [tab. 12]
# Matrix products use row vectors; W_h, W_l, W_v are learnable  [§4.2]

def fused_encoder(text_tokens, patch_feats):
    # Text encoder: last-layer hidden states  [eq. 2]
    language = t5_encoder(text_tokens)                    # (n, d)
    # Project patches to the text width d  [eq. 3]
    vision = patch_feats @ W_h                            # (m, d)
    # Single-head attention; text tokens are the queries  [§4.2]
    scores = language @ vision.T / sqrt(d)                # (n, m)
    vision_attn = softmax(scores, axis=-1) @ vision       # (n, d)
    # Per-token, per-dimension gate from both streams  [eq. 4]
    gate = sigmoid(language @ W_l + vision_attn @ W_v)    # (n, d)
    # Convex mix of text and attended picture  [eq. 5]
    fused = (1 - gate) * language + gate * vision_attn    # (n, d)
    return fused

# Supervised training, one model per stage  [§4.1]
# Stage 1 pairs: (text, image) -> rationale
# Stage 2 pairs: (text + gold rationale, image) -> answer
# Code's stage-2 input format  [code: utils_prompt.py#L63-L67]

def train(model, pairs, epochs, learning_rate):
    # Epochs: up to 20 in text; 20 (Base) or 50 (Large) in scripts
    #   [§5.2] [code: run_training.sh#L7] [code: run_training.sh#L29]
    for epoch in range(epochs):
        # Batch size 8 in text; per-device 2 or 4 for Large in scripts
        #   [§5.2] [code: run_training.sh#L29]
        for batch in batches(pairs):
            # Teacher-forced token negative log-likelihood  [eq. 1]
            loss = mean_token_nll(model, batch)
            update(model, loss, learning_rate)

Hyperparameters

Results

Author-reported. Accuracy in percent. The main table uses the ScienceQA test split (4k questions) §5.1 app. B.2. A-OKVQA is evaluated on its validation split because the test split is hidden, using the multiple-choice setting app. B.2 §5.1.

Main ScienceQA results tab. 4. NAT, SOC and LAN are the subject groups; TXT, IMG and NO are the kinds of context a question has; G1-6 and G7-12 are school-grade bands. Dagger marks concurrent work, as labelled by the authors. The source underlines the prior best average (86.54) and bolds the best average of this paper.

A-OKVQA validation accuracy, multiple-choice, Base model only tab. 5. Baseline rows come from the sources the authors cite.

Pilot: effect of CoT on a one-stage text-only FLAN-Alpaca Base model tab. 2.

Two-stage pilot tab. 3. Rationale generation is scored by ROUGE-L, answer inference by accuracy.

Ablation tab. 6.

Rationales from larger models instead of human annotation tab. 7. The authors use InstructBLIP for questions with images and ChatGPT for those without, zero-shot, as stage-1 targets §6.2 app. C.1.

Backbones, Base-sized models tab. 8.

Vision features, Base model tab. 9. Feature shapes from tab. 12.

Alignment strategy, Base model tab. 10.

Zero-shot transfer to MMMU, without training on it tab. 11.

Other findings, in words. The two-stage model with picture features starts higher and keeps improving over epochs, while the text-only two-stage model levels off early §6.1 fig. 5. Of 50 sampled incorrect answers, the authors classify 80% as commonsense mistakes, 14% as logical mistakes, and 6% as wrong answers with an empty or correct rationale §6.7 fig. 6. Vision features correct 60.7% of the hallucinated-rationale mistakes from the pilot §3.3 fig. 3.

Paper vs code

Code: amazon-science/mm-cot, Apache-2.0, at commit 8dd4ac0. The repository is linked from the paper and abstract and belongs to the authors' organisation. The repository contains Python and shell scripts for ScienceQA only. No A-OKVQA code is present, and no MMMU code is present.

• Learning rate and epochs. §5.2 gives 5e-5 and up to 20 epochs; the appendix allows {5e-5, 8e-5} app. B.3. The Base script uses 8e-5 for 20 epochs run_training.sh; the Large scripts use 5e-5 for 50 epochs run_training.sh run_training.sh.
• Batch size. §5.2 gives 8. Large scripts use 2 per device for stage 1 and 4 for stage 2 on 4 GPUs run_training.sh run_training.sh. The output directory name multiplies the batch by the GPU count main.py.
• Attention. The paper's equation is plain scaled dot-product attention with no projections §4.2. The code uses a one-head nn.MultiheadAttention over the projected patches, which has learned query, key and value projections and an output projection model.py model.py.
• Gate. The paper's gate has learnable W_l and W_v and no bias stated eq. 4. The code applies one linear layer with a bias to the concatenation of the two streams model.py model.py.
• Fusion location. The code fuses after the encoder's final layer norm and dropout, just before the decoder model.py. The paper does not describe this position.
• Blank images. The paper uses an all-zero block for questions without an image §4.2. The code passes that block through the same projection, so the encoder sees the projection's bias instead of zeros utils_data.py model.py.
• Rationale target. The paper calls the stage-1 target an "annotated rationale" §4.1. The code uses only the ScienceQA solution field and excludes the lecture utils_prompt.py.
• Answer-stage input. The paper says stage 2 takes the rationale from stage 1 §4.1. The code trains stage 2 on the gold lecture and solution utils_prompt.py and tests it on generated rationales utils_data.py. The two inputs differ.
• Missing entry point. run_training.sh and run_inference.sh call main_central.py for the Base rationale and answer steps and for the Large answer step run_training.sh run_training.sh run_training.sh. That file is not in the repository at this commit; only main.py is.
• Invalid answers. In the test-prediction path, an unparsed answer or an invalid letter is replaced by a random option main.py. In the in-training metric, an unparsed answer is simply wrong main.py. The paper does not say how invalid answers are scored.
• Input length for stage 2. The appendix says 512 for stage 1 and 64 for stage 2 app. B.3. In the code, 64 is the output length; the input length stays at its default of 512 for both stages main.py run_training.sh.
• Data and weights. The repository ships captions and a name map but not ScienceQA problems, the ViT feature file, or trained checkpoints; the README points to external downloads README.md.

The vision encoder matches the paper: timm vit_large_patch32_384, pretrained, and used frozen as precomputed features extract_features.py.

Pitfalls

• Conflicting hyperparameters. Learning rate (5e-5 in text, {5e-5, 8e-5} in the appendix, 8e-5 for Base in the script), epochs (up to 20 in text, 50 for Large in the script), and batch size (8 in text, 2 and 4 per device in Large scripts) disagree. Pick one and state it. The selection criterion for the learning rate is not stated.
• Model size. The text says 200M and 700M; Table 4 says 223M and 738M. The difference is not explained.
• Length labels. The appendix calls 512 and 64 the input lengths of the two stages. The code keeps the input at 512 and uses 64 as the output length of stage 2 main.py run_training.sh.
• Algorithm 1 is not in the paper. Its include line is commented out in the LaTeX source (line 437 of main-final.tex). Do not cite alg. 1 as a paper element.
• Cross-references in the appendix are off by two. "Section 6.2" for the vision feature comparison is §6.4, and "Section 6.1" for backbones is §6.3 app. B.1 §6.4.
• Train/test rationale mismatch. Stage 2 sees gold rationales in training and generated ones at test time utils_prompt.py. The paper does not discuss the gap, and its robustness to generation errors is not measured.
• Rationale scoring details. The code computes ROUGE-L with stemming and sentence splitting main.py. The paper does not state these, so ROUGE-L values may not be comparable across implementations.
• Blank-image bias. Questions without images still produce a fixed nonzero input to the fusion, from the projection bias model.py.
• Image fallback. Feature extraction uses image.png, and falls back to choice_0.png when it is missing extract_features.py. The paper does not state this; it can change which picture a question gets.
• Numbers from the authors only. Every result is a single run in the paper. No seeds, variance or confidence intervals are reported. The code default seed is 42 main.py.
• State-of-the-art claim. The abstract says the under-1B model reaches the best ScienceQA result, but the Large model (90.45) is below LLaVA-13B (90.92), a concurrent work the table marks with a dagger tab. 4. The Base model (85.31) is below Chameleon (GPT-4), the prior best at 86.54 tab. 4.
• Captions. The appendix says captions are appended app. B.3, and the scripts pass --use_caption run_training.sh. The paper does not say whether every reported number uses captions.
• A-OKVQA has no code here. Results on A-OKVQA cannot be reproduced from this repository, and A-OKVQA is evaluated on its validation split app. B.2.
• Optimizer and decoding. Optimizer, schedule, warmup, precision and decoding strategy are not stated in the paper. The code uses Trainer defaults for these under the pinned transformers version main.py.

Lineage

• Builds on: chain-of-thought prompting (cited in §2 as Wei et al. and Kojima et al.) and the ScienceQA benchmark (cited in §5.1). Neither has a card here yet.
• Compared against: InstructBLIP, used to caption images app. B.3 and as a comparison model tab. 4.
• Superseded by: not stated in this wiki's content.