Tattoo Tool Grounded Thinking Prm

TaTToo: Tool-Grounded Thinking PRM for Tabular Reasoning

Core Concept

Process Reward Models (PRMs) enable test-time scaling via intermediate step-level feedback, but existing PRMs struggle with table-specific operations (sub-table retrieval, schema navigation). TaTToo grounds each reasoning step in executable tool operations—SQL queries, table transforms, schema checks—generating precise reward signals for what constitutes progress in data analysis tasks.

Architecture Overview

• Dual-Stage Training: (1) Supervised fine-tuning on verification rationales; (2) reinforcement learning with tool-grounded rewards
• Tool-Based Verification: Each step grounds in actual tool execution (pandas operations, SQL queries), not just text descriptions
• 60K+ Step Annotations: Scalable pipeline generates high-quality step-level training data from tool execution traces
• Reward Shaping: Reward models learn to evaluate reasoning quality based on tool-execution feasibility and correctness
• Generalization: Trained TaTToo transfers across benchmarks and test-time scaling strategies

Implementation Steps

1. Grounding Reasoning Steps in Tool Operations

Each reasoning step must map to executable tool operations for verification.

class ToolGroundedStep:
    def __init__(self, reasoning_text, tool_operation, expected_output):
        """
        Reasoning step grounded in tool execution.

        Args:
            reasoning_text: Natural language explanation ("Filter rows where age > 25")
            tool_operation: Executable code (SQL/pandas/etc)
            expected_output: Ground truth result
        """
        self.reasoning = reasoning_text
        self.operation = tool_operation
        self.expected = expected_output
        self.actual_output = None
        self.verification_status = None

    def verify(self):
        """Execute tool and verify step correctness."""
        try:
            self.actual_output = execute_tool(self.operation)
            self.verification_status = matches_expected(self.actual_output, self.expected)
        except Exception as e:
            self.verification_status = False
            self.actual_output = f"Error: {e}"

        return self.verification_status

class TabularReasoningTrace:
    def __init__(self, question, table, expected_answer):
        self.question = question
        self.table = table
        self.expected_answer = expected_answer
        self.steps = []  # List of ToolGroundedStep

    def add_step(self, reasoning, tool_op, expected_output):
        step = ToolGroundedStep(reasoning, tool_op, expected_output)
        self.steps.append(step)
        return step

    def verify_full_trace(self):
        """Verify all steps and final answer."""
        step_verifications = []
        for step in self.steps:
            is_correct = step.verify()
            step_verifications.append(is_correct)

        # Final answer check
        final_answer = self.steps[-1].actual_output if self.steps else None
        final_correct = matches_expected(final_answer, self.expected_answer)

        return {
            'step_correctness': step_verifications,
            'final_correct': final_correct,
            'num_correct_steps': sum(step_verifications),
            'total_steps': len(self.steps)
        }

2. Data Generation Pipeline: 60K+ Step Annotations

Automated pipeline generating training data from tool execution traces.

def generate_step_level_annotations(questions, tables, max_traces=10000):
    """
    Generate 60K+ step-level training examples from tabular reasoning traces.
    """

    training_examples = []

    for question, table in zip(questions, tables):
        # Generate multiple reasoning traces per question
        for trace_idx in range(5):  # 5 traces per question
            trace = generate_reasoning_trace(question, table)

            # Verify each step
            verification_result = trace.verify_full_trace()

            # Extract step-level training examples
            for step_idx, step in enumerate(trace.steps):
                # Determine if step is on "good path" (leads to correct answer)
                is_on_correct_path = all(
                    trace.steps[i].verification_status
                    for i in range(step_idx + 1)
                ) and verification_result['final_correct']

                example = {
                    'question': question,
                    'table': table,
                    'step_number': step_idx,
                    'reasoning': step.reasoning,
                    'tool_operation': step.operation,
                    'step_correct': step.verification_status,
                    'on_correct_path': is_on_correct_path,
                    'intermediate_output': step.actual_output,
                    'label': 1.0 if is_on_correct_path else 0.0
                }

                training_examples.append(example)

    return training_examples

# Example: SQL-grounded tabular reasoning
trace_example = TabularReasoningTrace(
    question="How many employees have salary > 100k and joined after 2020?",
    table="employee_table",
    expected_answer=42
)

trace_example.add_step(
    reasoning="First, filter employees with salary > 100k",
    tool_op="SELECT * FROM employee_table WHERE salary > 100000",
    expected_output="[150000 rows]"
)

trace_example.add_step(
    reasoning="Filter those who joined after 2020",
    tool_op="SELECT * FROM filtered_above WHERE join_date > '2020-12-31'",
    expected_output="[42 rows]"
)

trace_example.add_step(
    reasoning="Count the results",
    tool_op="SELECT COUNT(*) FROM filtered_2020_above",
    expected_output=42
)

3. Dual-Stage Training: SFT + RL

Stage 1: Supervised fine-tuning on verification rationales. Stage 2: RL with tool-grounded rewards.

class TaTTooTrainer:
    def __init__(self, base_model='qwen-8b'):
        self.model = base_model
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=1e-5)

    def stage1_supervised_finetuning(self, training_examples, num_epochs=3):
        """
        SFT on step-level annotations: learn to evaluate step quality.
        """
        for epoch in range(num_epochs):
            total_loss = 0
            for example in training_examples:
                # Input: question + table + step reasoning
                input_text = f"""Question: {example['question']}
Table: {example['table']}
Step {example['step_number']}: {example['reasoning']}

Is this step correct? (Binary classification)"""

                # Target: label (0 or 1)
                target = example['label']

                # Forward pass
                logits = self.model.predict(input_text)
                loss = torch.nn.functional.binary_cross_entropy_with_logits(logits, target)

                # Backward
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

                total_loss += loss.item()

            print(f"Stage 1 Epoch {epoch}: Loss={total_loss/len(training_examples):.4f}")

    def stage2_reinforcement_learning(self, eval_traces, num_steps=5000):
        """
        RL: reward model learns from tool execution feedback.
        Policy (reasoning model) learns to maximize expected rewards.
        """

        policy = self.model  # PRM acts as reward signal for policy

        for step in range(num_steps):
            # Sample trace
            trace = random.choice(eval_traces)

            # Generate reasoning step
            step_reasoning = trace.steps[random.randint(0, len(trace.steps) - 1)]

            # Execute tool to get verification
            step_reasoning.verify()

            # Compute reward
            tool_execution_successful = step_reasoning.verification_status
            for s in trace.steps)

            reward = 1.0 if (tool_execution_successful and on_correct_path) else 0.0

            # Policy gradient update
            log_prob = self.model.log_probability(step_reasoning.reasoning)
            loss = -log_prob * reward

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            if (step + 1) % 100 == 0:
                print(f"Stage 2 Step {step+1}: Avg reward tracking...")

    def train_tattoo(self, train_examples, eval_traces):
        """Full training pipeline."""
        print("Stage 1: Supervised fine-tuning...")
        self.stage1_supervised_finetuning(train_examples)

        print("Stage 2: Reinforcement learning...")
        self.stage2_reinforcement_learning(eval_traces)

        return self.model

4. Evaluation on Tabular Benchmarks

Test-time scaling with TaTToo-trained reward models.

def evaluate_with_prm(policy, prm_reward_model, test_examples, num_beams=5):
    """
    Use PRM for test-time scaling via beam search.
    At each step, select highest-reward continuations.
    """

    correct = 0
    for example in test_examples:
        question = example['question']
        table = example['table']

        # Beam search with PRM guidance
        beams = [(question, [])]  # (current_state, step_history)

        for step_num in range(20):  # Max 20 steps
            scored_beams = []

            for state, history in beams:
                # Generate next reasoning step
                next_step = policy.generate_step(state, table, history)

                # Score with PRM
                reward = prm_reward_model.score(
                    question=question,
                    table=table,
                    step_number=step_num,
                    reasoning=next_step
                )

                scored_beams.append((reward, state + next_step, history + [next_step]))

            # Keep top-k beams
            scored_beams.sort(reverse=True)
            beams = [(s, h) for _, s, h in scored_beams[:num_beams]]

        # Evaluate final answer
        final_answer = extract_answer_from_reasoning(beams[0][0])
        if final_answer == example['expected_answer']:
            correct += 1

    accuracy = correct / len(test_examples)
    return accuracy

# Benchmark results
results = {
    'without_prm': {'accuracy': 0.45},  # Policy alone
    'with_tattoo_prm': {'accuracy': 0.55},  # +30.9% improvement
    'baseline_72b': {'accuracy': 0.52},  # Larger model without PRM
}

Practical Guidance

Tool Grounding: Every reasoning step must map to an executable operation. Abstract reasoning without tool execution doesn't generate reliable reward signals.

Training Data Scale: 60K+ examples provides good coverage across table schemas and operation types. Collect diverse traces (SQL, pandas, Excel operations) to ensure generalization.

Reward Shaping: Combine step-level correctness (tool execution) with trajectory-level signal (on-correct-path) to balance immediate feedback with long-horizon alignment.

Generalization: Train on 5 benchmarks, evaluate on held-out tasks. TaTToo generalizes better than task-specific PRMs due to tool-based grounding.

When to Use / When NOT to Use

Use When:

• Training agents on data analysis tasks (SQL, pandas, spreadsheets)
• Verifiable intermediate steps are available (query results, transformed tables)
• Test-time scaling via beam search or planning is feasible
• Smaller models need to match larger model performance

NOT For:

• Tasks without executable verification (open-ended reasoning)
• Domains where intermediate steps lack ground truth
• Real-time low-latency inference (beam search is expensive)

Reference

This skill synthesizes findings from "TaTToo: Tool-Grounded Thinking PRM for Test-Time Scaling in Tabular Reasoning" (arXiv:2510.06217). Tool grounding enables reliable step-level reward signals for process-based RL in data analysis.