💻 computer science

From Table to Cell: Attention for Better Reasoning with TABALIGN

The paper introduces TABALIGN, a novel reasoning framework that leverages a masked diffusion language model planner and a cell-grounding attention verifier to overcome the limitations of autoregressive models in multi-step table reasoning, achieving significant accuracy and efficiency gains across eight benchmarks.

Original authors: Tung Sum Thomas Kwok, Zeyong Zhang, Xinyu Wang, Chunhe Wang, Xiaofeng Lin, Hanwei Wu, Lei Ding, Guang Cheng, Zhijiang Guo

Published 2026-05-15

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Tung Sum Thomas Kwok, Zeyong Zhang, Xinyu Wang, Chunhe Wang, Xiaofeng Lin, Hanwei Wu, Lei Ding, Guang Cheng, Zhijiang Guo

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Left-to-Right" Trap

Imagine you are trying to solve a math problem written on a spreadsheet. The spreadsheet has rows of data (like "Obs 1," "Obs 2") and columns (like "Revenue," "Cost").

Current AI models are like students who are forced to read the paper strictly from left to right, top to bottom. They can't look ahead or jump around.

The Issue: If you shuffle the order of the rows in the spreadsheet (putting "Obs 3" before "Obs 1"), the student gets confused. They might pick the wrong numbers because they are stuck on the order they were taught to read, not the meaning of the data.
The Result: The AI often picks the wrong cells to look at, even if it eventually guesses the right answer. It's like finding the right answer by accident rather than by understanding the logic.

The Solution: TABALIGN

The authors built a new system called TABALIGN to fix this. Think of it as a two-person team working together to solve the spreadsheet puzzle: a Planner and an Executor.

1. The Planner (The "Diffusion" Architect)

Instead of the old "left-to-right" student, the team uses a Planner who is like an architect using a diffusion model (a type of AI that works like a sculptor chipping away at a block of marble to reveal a shape, rather than writing a sentence one word at a time).

How it works: The Planner looks at the entire spreadsheet at once. It doesn't care if the rows are shuffled. It sees the whole picture and draws a "blueprint" (a plan) of exactly which cells need to be used.
The Magic: Because it sees everything at once, it creates a much more stable and accurate map of where the important information is, regardless of how the table is arranged.

2. The Executor (The "Reasoner")

The Executor is the worker who actually does the math. It takes the blueprint from the Planner and starts clicking on the cells to get the answer.

The Problem: Even with a good blueprint, the Executor might get distracted and click on the wrong cell (like clicking "Total" instead of "Obs 1").
The Fix: The team added a Verifier (called TABATTN).

3. The Verifier (The "Spotter")

Imagine a coach standing next to the Executor. The coach has the Planner's blueprint in hand.

Every time the Executor clicks a cell, the coach checks: "Did you click the cell the blueprint said to click?"
If the Executor clicks the right cell, the coach says, "Good job, keep going."
If the Executor clicks the wrong cell (even if the final number looks okay), the coach says, "Stop! You're looking at the wrong place. Try again."

This ensures the AI doesn't just get lucky; it actually follows the correct path.

Why This Matters (The Results)

The paper tested this team (Planner + Executor + Coach) on eight different types of spreadsheet puzzles.

The Score: The TABALIGN team scored 15.76 points higher on average than the best existing open-source AI models of the same size.
Speed: Because the Planner creates a cleaner, more accurate map, the Executor doesn't waste time wandering around. The whole process became 44% faster in the actual reasoning steps.
Stability: If you shuffle the rows of the spreadsheet, the old AI gets confused and its performance drops. The TABALIGN team stays steady and performs the same, no matter how the table is organized.

The Core Insight

The paper discovered two main things:

Looking at the whole picture is better: Models that can look at all the data at once (like the Planner) understand tables much better than models that read line-by-line.
Checking the "where" is better than checking the "what": Instead of just asking, "Is the final answer right?", it's much more reliable to ask, "Did you look at the right specific cells to get that answer?"

Summary Analogy

Old AI: A robot that reads a menu from top to bottom. If the menu is rearranged, it orders the wrong food.
TABALIGN: A robot with a Chef (Planner) who looks at the whole kitchen and points to the exact ingredients needed, and a Sous-Chef (Executor) who grabs them. A Food Critic (Verifier) watches the Sous-Chef to make sure they are grabbing the ingredients the Chef pointed to, not just grabbing whatever is closest.

This system ensures the AI is not just guessing the right answer, but actually reasoning correctly through the table.

Technical Summary: From Table to Cell: Attention for Better Reasoning with TABALIGN

1. Problem Statement

Multi-step reasoning over structured tables using Large Language Models (LLMs) frequently fails due to a lack of an explicit cell-grounding contract between the planning and execution stages. Existing approaches suffer from two primary bottlenecks:

Planning Bottleneck (Left-to-Right Bias): Autoregressive (AR) models generate plans via left-to-right factorization. This imposes an order-sensitive prior that conflicts with the permutation invariance of tables (where row/column order should not alter semantics). Consequently, AR planners often produce globally inconsistent plans where early filtering steps are invalid because the necessary context (e.g., aggregation results) has not yet been scoped.
Execution Bottleneck (Unverified Attention): Existing Process Reward Models (PRMs) typically score intermediate states based on generated content (lexical or semantic quality) or final answer correctness. These methods overlook whether the model actually attended to the correct cells designated by the plan. A model can generate a correct final answer by "shortcutting" (attending to wrong cells but deriving the right value) or by hallucinating content, leading to unfaithful reasoning trajectories.

The core issue is the absence of a mechanism to verify that the reasoner's cell-level attention aligns with the planner's designated evidence.

2. Methodology: TABALIGN

The authors propose TABALIGN, a planned table reasoning framework that operationalizes the cell-grounding contract through two main components: a Masked Diffusion Language Model (DLM) Planner and the TABATTN Verifier.

2.1 The Planner: Bidirectional DLM

Instead of an AR generator, TABALIGN employs a masked DLM (specifically LLaDA-8B-Instruct) to generate the reasoning plan.

Bidirectional Denoising: The model denoises a fully masked plan canvas bidirectionally. This allows requirements in later steps to constrain earlier ones, enforcing inter-step dependencies and global consistency that AR models cannot achieve.
Cell Mask Generation: Each reasoning step in the plan concludes with a [target: col] tag. These tags are parsed into a binary cell mask ( $m_t \in \{0, 1\}^C$ ), explicitly designating which cells the step must operate on.
Rationale: A pilot study (Section 3) demonstrated that bidirectional attention models produce more human-aligned and permutation-stable cell attention than AR models, particularly when relevant cells are sparse.

2.2 The Verifier: TABATTN

To enforce the contract, TABALIGN introduces TABATTN, a lightweight verifier trained on 1,600 human-verified attention standards.

Attention Overlap Scoring: For each reasoning step $t$ , TABATTN computes the overlap between the reasoner's (an AR multimodal model, Qwen3-VL-8B) per-cell attention ( $a_t$ ) and the plan's designated mask ( $m_t$ ).
Reward Signal: The reward $R_{attn}$ is calculated as:
$R_{attn}(a_t; m_t) = \frac{\sum_{c=1}^C a_t^{(c)} m_t^{(c)}}{\max(\sum_{c=1}^C a_t^{(c)}, \epsilon)}$
This score serves as an in-loop feedback signal. If the overlap is low, the system flags the step, potentially triggering a resample or halting unproductive iterations.
Labelability: The authors argue that cell-relevance (is this cell necessary?) is a more labelable supervision target than step-correctness (is this step correct?), as the former admits confident binary decisions while the latter often requires ambiguous "unsure" judgments.

2.3 Execution Flow

Planning: The DLM generates a step-by-step blueprint with explicit cell masks.
Execution: The AR reasoner executes the plan step-by-step, selecting modalities (text or image) based on the operation.
Verification: After each step, TABATTN scores the attention overlap. If the score stagnates or is low, the system intervenes.
Termination: A stagnation halt mechanism stops the process if the attention score fails to improve by a threshold ( $\delta=0.02$ ) for two consecutive steps, emitting the best-scoring answer found.

3. Key Contributions

Cell-Level Attention Pilot Study: The authors demonstrate that bidirectional attention mechanisms (via diffusion) are more human-aligned and permutation-stable than AR attention. They show that cell-relevance labels are significantly more labelable (93.0% agreement) than step-correctness labels (73.2% agreement) under identical LLM-as-a-judge pipelines.
TABALIGN Framework: A novel architecture pairing a bidirectional DLM planner with an attention-overlap verifier (TABATTN). This framework explicitly captures grounding failures missed by content-based rewards and utilizes optimized DLM plans to reduce column hallucinations.
Empirical Validation: Consistent performance gains across eight diverse table reasoning benchmarks (covering Table QA and Fact Verification) at a comparable 8B-parameter scale.

4. Experimental Results

Evaluated on eight benchmarks (WTQ, MMQA, TabMWP, TAT-QA, HiTab, FeTaQA, TabFact, InfoTabs):

Accuracy Gains: TABALIGN achieves an average accuracy of 90.62% across six accuracy-based benchmarks. This represents a 15.76 percentage point (pp) improvement over the strongest open-source baseline (TableDART on Ovis2, 74.86%) and a 16.35 pp gain over the strongest single-modality baseline (Qwen3-VL-8B, 74.27%).
Ablation Studies:
- Planner Choice: Replacing the AR planner with the DLM planner on a fixed reasoner backbone yields a 2.87 pp gain, attributing part of the success to the DLM's superior planning capabilities.
- Reward Signal: Using TABATTN (attention-grounded) improves accuracy by 4.8 pp over content-based rewards (e.g., TABROUGE), with the largest gains on numerically and structurally dense tasks.
- Efficiency: While DLM planning is slower (78.0s vs 6.0s for AR planning), the cleaner plans reduce downstream reasoning execution time by 44.64% (from 102.6s to 56.8s), resulting in a net 24% increase in total wall-clock time for a significant accuracy boost.
Error Reduction: The framework reduces schema hallucinations (column hallucination rates drop from 75.0% to 45.9%) and compresses trajectory lengths by 22%.

5. Significance and Claims

The paper claims that TABALIGN addresses the fundamental "cell-grounding contract" gap in table reasoning. By decoupling planning (via bidirectional DLM) and execution (via AR with attention verification), the framework ensures that reasoning steps are grounded in the intended evidence rather than relying on spurious content overlaps.

The authors emphasize that:

Bidirectional Attention is Critical: The permutation stability and human alignment of bidirectional pre-training are essential for handling the order-invariant nature of tables.
Attention is a Better Supervision Signal: Verifying where the model looks (attention overlap) is more effective and labelable than verifying what the model says (content similarity) for intermediate steps in table reasoning.
Scalability: The framework operates effectively at the 8B scale, outperforming larger proprietary models and complex trajectory-optimization agents without requiring massive parameter scaling.

The work positions itself as a step toward more interpretable and auditable AI reasoning over structured data, reducing trajectory drift and hallucination compounding in multi-step agents.