Salesforce/cogalign-llava-ov-0_5b

Why Vision Language Models Struggle with Visual Arithmetic? Towards Enhanced Chart and Geometry Understanding

Vision Language Models (VLMs) have achieved remarkable progress in multimodal tasks, yet they often struggle with visual arithmetic, seemingly simple capabilities like object counting or length comparison, which are essential for relevant complex tasks like chart understanding and geometric reasoning. In this work, we first investigate the root causes of this deficiency through a suite of probing tasks focusing on basic visual arithmetic. Our analysis reveals that while pre-trained vision encoders typically capture sufficient information, the text decoder often fails to decode it correctly for arithmetic reasoning. To address this, we propose CogAlign, a novel post-training strategy inspired by Piaget's theory of cognitive development. CogAlign trains VLMs to recognize invariant properties under visual transformations. We demonstrate that this approach significantly improves the performance of three diverse VLMs on our proposed probing tasks. Furthermore, CogAlign enhances performance by an average of 4.6% on CHOCOLATE and 2.9% on MATH-VISION, outperforming or matching supervised fine-tuning methods while requiring only 60% less training data. These results highlight the effectiveness and generalizability of CogAlign in improving fundamental visual arithmetic capabilities and their transfer to downstream tasks.

Quick start

You must install llava-next (https://github.com/LLaVA-VL/LLaVA-NeXT/tree/main).

pip install git+https://github.com/LLaVA-VL/LLaVA-NeXT.git`

Then, you can load models easily with llava:

import copy
from llava.model.builder import load_pretrained_model
from llava.mm_utils import get_model_name_from_path, process_images, tokenizer_image_token
from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, IGNORE_INDEX
from llava.conversation import conv_templates, SeparatorStyle
model_name = "llava_qwen"
tokenizer, llava_model, image_processor, _ = load_pretrained_model(
    "Salesforce/cogalign-llava-ov-0_5b",
    None,
    model_name,
    device_map="cpu",
)
llava_model.cuda()

After that, we can try an example on the CHOCOLATE datset:

chocolate = load_dataset("khhuang/CHOCOLATE")["test"]
chocolate_df = pd.DataFrame(chocolate)
chocolate_df_lvlm = chocolate_df.loc[chocolate_df.split=="LVLM",:]

instance =  chocolate_df_lvlm.iloc[2]
caption = ' '.join(instance.sentences)
prompt = f"""
You are given a chart and a caption, you are tasked to detect whether the caption is factually
consistent with the chart. A caption is factually consistent with the chart if it describes the datapoints within the charts without factual errors (e.g. wrong label, value, trends).
[Start of Caption]
{caption}
[End of Caption]
For the above caption, you should respond 'Answer: Yes' if it is factually consistent with the chart. Otherwise, respond 'Answer: No'. Do not provide explanation or other thing.
"""

url = instance.image_path
image = Image.open(requests.get(url, stream=True).raw)
image_tensor = process_images([image], image_processor, llava_model.config)
image_tensor = [_image.to(dtype=torch.float16, device='cuda') for _image in image_tensor]

conv_template = "qwen_1_5"  # Make sure you use correct chat template for different models
question = DEFAULT_IMAGE_TOKEN + f"\n{prompt}"
conv = copy.deepcopy(conv_templates[conv_template])
conv.append_message(conv.roles[0], question)
conv.append_message(conv.roles[1], None)
prompt_question = conv.get_prompt()

input_ids = tokenizer_image_token(prompt_question, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to('cuda')
image_sizes = [image.size]


cont = llava_model.generate(
    input_ids,
    images=image_tensor,
    image_sizes=image_sizes,
    do_sample=False,
    temperature=0,
    max_new_tokens=4096,
)
response = tokenizer.batch_decode(cont, skip_special_tokens=True)[0]
print(f'User: {question}\nAssistant: {response}')

License information

This release is for research purposes only in support of an academic paper. This repository is licensed under the noncommercial license CC-BY-NC 4.0.

Citation

If you find CogAlign useful in your research, please consider citing:

@misc{huang-etal-2025-cogalign,
    title = "Why Vision Language Models Struggle with Visual Arithmetic? Towards Enhanced Chart and Geometry Understanding",
    author = "Huang, Kung-Hsiang  and
      Qin, Can  and
      Qiu, Haoyi  and
      Laban, Philippe  and
      Joty, Shafiq  and
      Xiong, Caiming  and
      Wu, Chien-Sheng",
    year = "2025",
    eprint={2502.11492},
    archivePrefix = "arXiv",
    primaryClass={cs.AI}
}