Math-PUMA: Progressive Upward Multimodal Alignment to Enhance Mathematical Reasoning
Math-PUMA uses 692K alignment pairs and three-stage training to reduce the text–vision gap in mathematical reasoning.
Key Findings
Methodology
Math-PUMA implements Progressive Upward Multimodal Alignment in three stages. It first strengthens the language model with 200K textual math problems, then aligns text-rich and vision-rich representations using 692K equivalent-information pairs and forward/reverse KL divergence over next-token distributions, and finally applies multimodal instruction tuning on 996K high-quality examples. Its objective is L=λKL(αLFKL+(1−α)LRKL)τ²+(1−λ)Lhard, combining soft distribution matching with answer supervision.
Key Results
- On MATHVERSE, Math-PUMA-Qwen2-7B reaches 33.6% overall accuracy versus 22.9% for Math-LLaVA, an improvement of about 10.7 percentage points. It still obtains 26.0% on vision-only problems, indicating substantially stronger visual mathematical reasoning.
- On MATHVISTA, the 7B model scores 48.1% GPS, 47.7% ALG, 47.3% GEO, and 42.6% SCI. Relative to same-scale InternLM-XComposer2-VL, the reported gains are 16.4, 15.7, 16.8, and 4.9 percentage points.
- Ablations identify Stage 2 as the main mechanism reducing the text–vision gap. Removing Stage 3 sharply harms conversational ability, while swapping Stages 2 and 3 substantially reduces performance, demonstrating that training order matters.
Significance
The work addresses a persistent failure mode of MLLMs: performance decreases as mathematical information moves from text into diagrams, even though diagrams often help humans. By separating mathematical capability acquisition from cross-modal alignment and then recombining them, Math-PUMA gives relatively small 7B systems competitive results on MATHVERSE, MATHVISTA, and WE-MATH. The approach is relevant to educational tutoring, scientific question answering, and controllable visual reasoning.
Technical Contribution
The central technical contribution is not merely more image–question data, but paired representations of the same problem with different modality allocations. A stronger text-rich input supplies teacher logits to a weaker vision-rich input. The corpus contains 120K automatically generated pairs and 572K augmented pairs; training combines bidirectional KL distillation with hard answer loss using αKL=0.2, τ=1.0, and λKL=0.1. This differs from Geo170K, MathV360K, and similar methods that primarily expand questions or images without explicitly enforcing prediction consistency.
Novelty
Math-PUMA’s fundamental novelty is to model multimodal mathematical competence as a hierarchy from text-only to vision-only inputs and align it progressively from stronger to weaker levels. Unlike most data-augmentation approaches, it directly constrains token-level predictive distributions for equivalent problems. The paper therefore turns modality invariance into an explicit optimization objective rather than an incidental consequence of instruction tuning.
Limitations
- The automatic diagrams are produced by predefined plotting tools, so the method may underrepresent messy handwriting, unusual layouts, occlusion, and real scientific figures. The reported benchmarks do not fully establish robustness in those conditions.
- Evaluation emphasizes accuracy and WE-MATH average/Rote Memorization scores; process faithfulness, calibration, error categories, and multilingual transfer are not systematically studied.
- Stage-2 distillation assumes that text-rich teacher logits are reliable; incorrect textual solutions can propagate errors into the vision branch.
Future Work
Future research should extend the pair construction pipeline to dynamic diagrams, handwritten figures, and scientific charts, while adding theorem provers, executable solvers, and visual grounding feedback. Adaptive difficulty, multilingual alignment, and process-level evaluation are important directions. Reducing the data and compute burden—692K alignment pairs trained with 32 A100 GPUs—will be essential for broad deployment.
AI Executive Summary
Multimodal large language models can solve many textual mathematics problems, yet often fail when crucial information is embedded in a diagram. The paradox is that visual aids usually help people, but can hurt models. Existing systems such as LLaVA, Qwen-VL, and DeepSeek-VL were largely exposed to natural images, while methods such as Geo170K and MathV360K mainly enlarge multimodal datasets without explicitly requiring the same problem to yield the same reasoning behavior across representations.
Math-PUMA addresses this gap through Progressive Upward Multimodal Alignment. Stage 1 strengthens the language model with 200K textual mathematics examples. Stage 2 constructs 692K pairs containing equivalent information but different proportions of text and vision, then matches their next-token distributions with forward and reverse KL divergence, supplemented by hard answer loss. Stage 3 performs instruction tuning on 996K high-quality multimodal examples, including an expanded geometric subset of MathV360K. The underlying idea is simple: teach the model mathematics first, then teach it that changing the presentation should not change the solution.
The results are strong for relatively compact models. Math-PUMA-Qwen2-7B reaches 33.6% on MATHVERSE, compared with 22.9% for Math-LLaVA, and scores 26.0% on vision-only items. On MATHVISTA it achieves 48.1% GPS, 47.7% ALG, 47.3% GEO, and 42.6% SCI. On WE-MATH, the 7B model surpasses several open-source systems with more than 20B parameters. Ablations show that Stage 2 narrows modality gaps, whereas Stage 3 preserves conversational behavior. Remaining challenges include synthetic-figure bias, costly training, unreliable teacher solutions, and insufficient testing on messy real-world diagrams.
Deep Analysis
Background
Text-based LLM reasoning has advanced through large mathematical corpora and chain-of-thought supervision. MLLMs such as LLaVA, Qwen-VL, and DeepSeek-VL add visual encoders and projectors, but their image training is dominated by natural scenes. Mathematical diagrams contain abstract symbols, geometric relations, and layout-dependent meaning, creating a domain gap. Geo170K, MathV360K, and GeoGPT4V improve data coverage, but do not directly enforce modality-invariant reasoning for identical problems.
Core Problem
A single mathematical fact may appear as complete text, text plus a diagram, diagram-dominant content, or an almost visual-only prompt. Current MLLMs often perform well in text-only settings and degrade sharply as visual information increases. The challenge is to preserve mathematical reasoning while aligning different modality-conditioned distributions, without allowing noisy solutions or modality-specific shortcuts to become supervision.
Innovation
Math-PUMA contributes four linked ideas. First, it represents capability as a hierarchy from text-only to vision-only inputs. Second, it creates equivalent-information pairs instead of merely independent image questions. Third, it uses forward and reverse KL divergence to align strong and weak modality predictions, while hard labels stabilize optimization. Fourth, it orders textual math enhancement, multimodal alignment, and instruction tuning into a single pipeline, producing the Math-PUMA-1M resource.
Methodology
- �� Stage 1: fine-tune the initialized LLM on 200K textual mathematical examples.
- �� Pair generation: Question Designer, Plotter, Solver, and Pair Constructor create up to four modality forms for each problem.
- �� Automatic data: generate 40K pairs each for plane geometry, solid geometry, and functions, totaling 120K.
- �� Augmentation: rephrase 80K public problems, apply scaling, stretching, and gamma transformation to reach 310K, and add 262K VisualWebInstruct examples for 572K pairs.
- �� Stage 2: obtain logits z from strong and weak inputs, compute softmax(z/τ), then combine forward KL, reverse KL, and hard answer loss.
- �� Stage 3: add MathV360K, expand its geometry subset from 40K to 120K, balance textual data, and train on 996K multimodal instructions.
Experiments
The study evaluates Math-PUMA-Qwen2-1.5B, Qwen2-7B, and DeepSeek-Math-7B. Qwen models use SigLIP-so400m-patch14-384, with a LLaVA-like architecture; DeepSeek compatibility follows DeepSeek-VL. Benchmarks are MATHVERSE, MATHVISTA, and WE-MATH. GPT-4o-mini extracts and verifies answers for the first two, while WE-MATH reports average and Rote Memorization scores. AdamW uses β1=0.9, β2=0.999; learning rates are 3e-5, 5e-5, and 3e-5, with one epoch, batch sizes 256/512/256, and 32 A100 80GB GPUs.
Results
On MATHVERSE, the 1.5B, Qwen2-7B, and DeepSeek-7B systems achieve 29.6%, 33.6%, and 31.8% overall accuracy; Qwen2-7B obtains 26.0% on vision-only items. On MATHVISTA, it reaches 47.9% overall and 48.1%, 47.7%, 47.3%, and 42.6% in GPS, ALG, GEO, and SCI. The authors report open-source SOTA or near-SOTA performance and competitiveness with GPT-4V. Removing Stage 2 widens modality gaps; removing Stage 3 sharply reduces conversational performance.
Applications
The method can support educational systems that read geometry diagrams and problem text, then produce structured explanations. It can also assist scientific question answering over function plots, experiment schematics, and basic engineering drawings. Practical deployment requires reliable visual parsing, answer verification, uncertainty reporting, and domain restrictions. The results suggest that careful alignment may reduce dependence on very large proprietary models for constrained mathematical tasks.
Limitations & Outlook
The approach depends on automatic plotting, GPT-4o-mini-generated explanations, and textual teacher distributions, so synthetic style and teacher errors may limit generalization. KL agreement does not guarantee genuine understanding of geometric relations; a model may imitate token patterns. The 692K-pair and 996K-instruction pipeline is computationally expensive, using 32 A100 GPUs. Future work should include real hand-drawn and complex scientific figures, executable verification, process-level metrics, multilingual evaluation, and more efficient pair selection.
Plain Language Accessible to non-experts
Imagine a student learning mathematics in a classroom. First, the teacher gives the student 200,000 problems written entirely in words. The student learns how to calculate, compare, and prove things. Next, the teacher takes the same problems and changes their packaging: one version has full sentences, another has sentences plus a picture, another hides most clues in the picture, and the last is almost all picture. The student must reach the same answer each time.
The teacher does more than check the final answer. They compare what the student is likely to write next in each version. If the student begins reasoning differently just because the information moved into a diagram, the teacher gently pulls the two approaches together. This is the practical meaning of Math-PUMA’s alignment stage. Finally, the student practices 996,000 richer picture-and-text conversations.
The result is a smaller AI student that handles both formats better: its 7B version scores 33.6% on MATHVERSE and 47.9% on MATHVISTA overall. But the student mainly practiced clean computer-made diagrams. Messy handwriting, confusing drawings, or an incorrect teacher explanation can still cause mistakes.
ELI14 Explained like you're 14
Think of an AI playing a math puzzle game. In one mode, every clue is typed out. In another, the clues are hidden in a map, triangle, graph, or diagram. Many AIs are like students who ace the written mode but suddenly get lost when the game switches to pictures. Why? They have seen tons of ordinary photos, but not enough mathematical drawings.
Math-PUMA is a three-level training camp. Level one gives the AI 200,000 text math problems so it learns the actual math. Level two shows the same problem in different formats and tells the AI, “New outfit, same answer!” It compares what the AI predicts next in the text and picture versions, then corrects it with the real solution. Level three gives it 996,000 multimodal instruction examples so it can explain answers in conversation.
The results are impressive for a 7B model: 33.6% on MATHVERSE, compared with 22.9% for Math-LLaVA, and about 47% on several MATHVISTA categories. In other words, the AI is no longer limited to reading the question; it has learned to use the diagram too.
Is it unbeatable? Definitely not! Clean computer diagrams are easier than blurry photos or messy homework sketches. And if the teacher’s solution is wrong, the AI might copy the mistake. The next challenge is to make it check its own geometry, understand real-world drawings, and explain exactly which line, angle, or formula supports its answer.
Glossary
Progressive Upward Multimodal Alignment
A training strategy that moves from stronger text-based mathematical reasoning toward increasingly visual inputs. It aims to preserve equivalent problem-solving ability when information changes modality.
This is the central mechanism of Math-PUMA Stage 2.
KL Divergence
A measure of discrepancy between probability distributions; forward and reverse forms emphasize different directions of mismatch. The paper applies both to next-token distributions from paired inputs.
It provides soft supervision for cross-modal alignment.
Text-rich / Vision-rich
Text-rich samples place a larger share of problem information in language, whereas vision-rich samples place more information in an image. The paired samples should preserve the same facts and answer.
These categories define the 692K alignment pairs.
Hard Loss
A direct loss against the ground-truth solution tokens, usually cross-entropy. It supplies an explicit target and stabilizes soft-distribution matching.
It is combined with forward and reverse KL losses in Stage 2.
MATHVERSE
A multimodal mathematics benchmark that evaluates five conditions ranging from text-dominant to vision-only inputs. It is designed to expose performance differences caused by modality changes.
It is the main benchmark for the paper’s comparisons and ablations.
MATHVISTA
A benchmark requiring fine-grained visual understanding and compositional mathematical reasoning across domains such as geometry, algebra, and science.
Math-PUMA reports category-level GPS, ALG, GEO, and SCI results on it.
Open Questions Unanswered questions from this research
- 1 Does matching token distributions imply genuine understanding of diagram relations, or can the model imitate surface language? Process-level proofs, visual grounding, and counterfactual diagram tests are needed.
- 2 Can synthetic training diagrams transfer to handwriting, low resolution, occlusion, and culturally different notation? The current benchmarks do not answer this systematically.
- 3 How should unreliable textual teachers be handled? Verified solvers, confidence filtering, and teacher ensembles may be required to prevent error propagation.
Applications
Immediate Applications
Multimodal mathematics tutoring
Educational platforms could provide a problem statement, geometry diagram, and student work to a 7B-scale model, which would generate structured steps and identify relevant visual relations. Deployment should include answer checking and human review because synthetic-data bias can produce confidently wrong instruction.
Scientific figure question answering
Research assistants could use the model to read elementary plots, experimental schematics, and engineering sketches together with captions. Early use should focus on domains with standard answers and executable verification, while exposing uncertainty rather than treating every response as authoritative.
Long-term Vision
Verifiable visual reasoning
A future system could combine visual grounding, symbolic calculation, theorem proving, and language explanation. It would identify the exact line, angle, or formula supporting each step and automatically verify the derivation, transforming multimodal models from answer generators into auditable reasoning tools.
Abstract
Multimodal Large Language Models (MLLMs) excel in solving text-based mathematical problems, but they struggle with mathematical diagrams since they are primarily trained on natural scene images. For humans, visual aids generally enhance problem-solving, but MLLMs perform worse as information shifts from textual to visual modality. This decline is mainly due to their shortcomings in aligning images and text. To tackle aforementioned challenges, we propose Math-PUMA, a methodology focused on Progressive Upward Multimodal Alignment. This approach is designed to improve the mathematical reasoning skills of MLLMs through a three-stage training process, with the second stage being the critical alignment stage. We first enhance the language model's mathematical reasoning capabilities with extensive set of textual mathematical problems. We then construct a multimodal dataset with varying degrees of textual and visual information, creating data pairs by presenting each problem in at least two forms. By leveraging the Kullback-Leibler (KL) divergence of next-token prediction distributions to align visual and textual modalities, consistent problem-solving abilities are ensured. Finally, we utilize multimodal instruction tuning for MLLMs with high-quality multimodal data. Experimental results on multiple mathematical reasoning benchmarks demonstrate that the MLLMs trained with Math-PUMA surpass most open-source MLLMs. Our approach effectively narrows the performance gap for problems presented in different modalities. The code and data are available at: \url{https://github.com/wwzhuang01/Math-PUMA}.