Rlad Discovering Abstractions

RLAD: Two-Player RL for Abstraction Discovery

Long reasoning traces often become verbose and circular. A 50-token chain-of-thought might say the same thing three times before reaching an answer. The issue is that raw step-by-step reasoning doesn't isolate what's actually useful—the key insights or procedural knowledge that solve the problem.

RLAD approaches this differently. Instead of generating long traces, it trains two players: one discovers abstractions (concise summaries of useful knowledge), and one uses those abstractions to solve problems. Both learn through RL, creating a natural curriculum where abstractions must be genuinely useful to receive reward.

Core Concept

RLAD is a two-player system:

1. Abstraction Generator: Proposes 1-3 concise descriptions of key insights (e.g., "Fibonacci numbers grow exponentially," "Even numbers are divisible by 2")
2. Solution Generator: Receives abstractions as context, uses them to solve the problem
3. Reward signal: Solution generator gets reward for correctness; abstraction generator gets reward when its abstractions help solution generation

The key insight: abstractions must be useful. An abstraction only receives reward if it actually helps solve the problem, not for being verbose or clever.

Architecture Overview

• Abstraction generator: Encoder-decoder or autoregressive model generating concise descriptions
• Solution generator: Uses proposed abstractions as part of reasoning context
• Verifier: Checks if solution is correct
• Reward aggregator: Assigns credit to abstractions based on solution success
• Training loop: Two-player RL with shared rewards

Implementation Steps

First, implement the abstraction generation component:

import torch
import torch.nn.functional as F

class AbstractionGenerator:
    """
    Generate domain-specific abstractions (key insights).
    """
    def __init__(self, model, max_abstraction_length=30):
        self.model = model
        self.max_length = max_abstraction_length

    def generate_abstractions(self, problem, num_abstractions=3):
        """
        Generate multiple abstractions for a problem.

        Args:
            problem: Problem statement (text)
            num_abstractions: How many abstractions to generate

        Returns:
            abstractions: List of text descriptions
        """
        prompt = f"""Given this problem, what are the key insights or knowledge needed to solve it?

Problem: {problem}

Generate {num_abstractions} concise abstractions (1 sentence each) that capture important insights.
Format:
[ABSTRACTION 1]: <description>
[ABSTRACTION 2]: <description>
etc.
"""

        abstractions = []

        for _ in range(num_abstractions):
            # Generate one abstraction at a time for diversity
            output = self.model.generate(prompt, max_length=self.max_length)

            # Extract abstraction from output
            lines = output.split('\n')
            for line in lines:
                if '[ABSTRACTION' in line and ':' in line:
                    abstraction = line.split(':', 1)[1].strip()
                    abstractions.append(abstraction)
                    break

        return abstractions[:num_abstractions]

    def score_abstraction_usefulness(self, abstraction, problem, solution_success):
        """
        Simple heuristic: abstraction is useful if it's mentioned in correct solution.
        (In practice, would use more sophisticated metrics.)

        Args:
            abstraction: Generated abstraction text
            problem: Problem statement
            solution_success: Did the solution succeed?

        Returns:
            usefulness_score: 0-1 scalar
        """
        # This is a placeholder; real scoring would be more sophisticated
        # (e.g., measuring if solution leverages the abstraction)
        if solution_success:
            # Reward abstractions that are concise and conceptually sound
            score = min(1.0, 30.0 / len(abstraction.split()))  # Shorter is better
            return score
        else:
            return 0.0

Now implement the solution generator that uses abstractions:

class SolutionGenerator:
    """
    Generate solutions given problem and abstractions.
    """
    def __init__(self, model, verifier):
        self.model = model
        self.verifier = verifier

    def solve_with_abstractions(self, problem, abstractions):
        """
        Generate solution using provided abstractions as context.

        Args:
            problem: Problem statement
            abstractions: List of key abstractions

        Returns:
            solution: Generated solution
            reasoning: Full reasoning trace
        """
        # Build prompt that includes abstractions
        abstraction_text = "\n".join(
            [f"- {abs}" for abs in abstractions]
        )

        prompt = f"""Use these key insights to solve the problem:

{abstraction_text}

Problem: {problem}

Solve step by step:"""

        # Generate solution
        output = self.model.generate(prompt, max_length=200)

        return output

    def score_solution(self, solution, problem):
        """
        Check if solution is correct.

        Args:
            solution: Proposed solution
            problem: Original problem

        Returns:
            is_correct: Binary correctness
            confidence: 0-1 confidence score
        """
        is_correct = self.verifier.verify(solution, problem)
        confidence = self.verifier.get_confidence(solution, problem)

        return float(is_correct), confidence

Implement the two-player RL training loop:

class TwoPlayerRLTrainer:
    """
    Train abstraction and solution generators together via RL.
    """
    def __init__(self, abstraction_model, solution_model, verifier):
        self.abstraction_gen = AbstractionGenerator(abstraction_model)
        self.solution_gen = SolutionGenerator(solution_model, verifier)
        self.verifier = verifier

    def training_step(self, problem, verifier, num_abstraction_samples=3):
        """
        Single training step: generate abstractions and solutions, compute rewards.

        Args:
            problem: Problem to solve
            verifier: Correctness verifier
            num_abstraction_samples: Number of abstraction sets to try

        Returns:
            loss_abstractions: Loss for abstraction generator
            loss_solution: Loss for solution generator
        """
        all_losses_abs = []
        all_losses_sol = []
        best_solution_success = 0

        # Try multiple abstraction sets
        for _ in range(num_abstraction_samples):
            # Generate abstractions
            abstractions = self.abstraction_gen.generate_abstractions(problem)

            # Generate solution using these abstractions
            solution = self.solution_gen.solve_with_abstractions(problem, abstractions)

            # Check correctness
            is_correct, confidence = self.solution_gen.score_solution(solution, problem)

            # Reward signal: solution success
            solution_reward = is_correct

            # Abstraction reward: same as solution (they share credit)
            # In practice, could use more sophisticated credit assignment
            abstraction_reward = is_correct

            # Loss for solution generator (maximize solution reward)
            loss_sol = -torch.log(torch.tensor(solution_reward + 1e-6))
            all_losses_sol.append(loss_sol)

            # Loss for abstraction generator (maximize abstraction reward)
            # Make reward proportional to abstraction usefulness
            abstraction_usefulness = 0.0
            for abs_item in abstractions:
                abs_use = self.abstraction_gen.score_abstraction_usefulness(
                    abs_item,
                    problem,
                    is_correct
                )
                abstraction_usefulness += abs_use

            abstraction_usefulness /= len(abstractions)
            loss_abs = -torch.log(torch.tensor(abstraction_usefulness + 1e-6))
            all_losses_abs.append(loss_abs)

            best_solution_success = max(best_solution_success, is_correct)

        # Average losses across samples
        avg_loss_abs = torch.stack(all_losses_abs).mean()
        avg_loss_sol = torch.stack(all_losses_sol).mean()

        return avg_loss_abs, avg_loss_sol, best_solution_success

    def train(self, problems, num_epochs=5):
        """
        Train both generators on problem set.

        Args:
            problems: List of problems
            num_epochs: Training epochs

        Returns:
            trained_models: (abstraction_model, solution_model)
        """
        abs_optimizer = torch.optim.AdamW(
            self.abstraction_gen.model.parameters(),
            lr=1e-5
        )
        sol_optimizer = torch.optim.AdamW(
            self.solution_gen.model.parameters(),
            lr=1e-5
        )

        for epoch in range(num_epochs):
            total_loss_abs = 0
            total_loss_sol = 0
            total_success = 0

            for problem in problems:
                # Forward pass
                loss_abs, loss_sol, success = self.training_step(problem)

                # Backward pass for abstractions
                abs_optimizer.zero_grad()
                loss_abs.backward()
                abs_optimizer.step()

                # Backward pass for solutions
                sol_optimizer.zero_grad()
                loss_sol.backward()
                sol_optimizer.step()

                total_loss_abs += loss_abs.item()
                total_loss_sol += loss_sol.item()
                total_success += success

            print(f"Epoch {epoch+1}: "
                  f"loss_abs={total_loss_abs/len(problems):.4f}, "
                  f"loss_sol={total_loss_sol/len(problems):.4f}, "
                  f"success_rate={total_success/len(problems):.2%}")

        return self.abstraction_gen.model, self.solution_gen.model

Implement allocation of computation at test time:

def test_time_scaling_with_abstractions(
    problem,
    abstraction_gen,
    solution_gen,
    verifier,
    compute_budget=100  # Tokens or inference calls
):
    """
    At test time, decide whether to generate more abstractions or more solutions.

    Args:
        problem: Problem to solve
        abstraction_gen: Trained abstraction generator
        solution_gen: Trained solution generator
        verifier: Correctness verifier
        compute_budget: Total compute available

    Returns:
        best_solution: Best solution found within budget
    """
    # Split budget between abstraction generation and solution generation
    # Research finding: more abstractions often better than more solutions
    num_abstractions = max(1, compute_budget // 10)
    solutions_per_abstraction = max(1, (compute_budget - 10) // num_abstractions)

    best_solution = None
    best_score = 0

    # Generate multiple abstraction sets
    for _ in range(num_abstractions):
        abstractions = abstraction_gen.generate_abstractions(problem)

        # Generate multiple solutions for each abstraction set
        for _ in range(solutions_per_abstraction):
            solution = solution_gen.solve_with_abstractions(problem, abstractions)
            is_correct, _ = solution_gen.score_solution(solution, problem)

            if is_correct > best_score:
                best_solution = solution
                best_score = is_correct

    return best_solution

Practical Guidance

When to use RLAD:

• Reasoning tasks where abstractions significantly reduce complexity (math, logic, coding)
• Tasks requiring domain-specific procedural knowledge
• Settings where transparency matters (abstractions are human-interpretable)
• Problems where test-time compute scaling is beneficial

When NOT to use:

• Simple classification tasks (abstraction overhead exceeds benefit)
• Domains without clear abstract structure
• Extreme latency constraints (two-stage reasoning adds overhead)
• Tasks with single optimal solution paths

Performance improvements on reasoning benchmarks:

Key allocation insight: At larger compute budgets, allocating more compute to abstraction generation yields better returns than generating more solutions with fixed abstractions. Trade-off curve:

Common pitfalls:

• Abstraction reward too weak: If abstractions don't receive clear reward signal, generator learns nothing. Verify that abstraction-less solutions have <50% success before training.
• Two-player instability: Balancing two loss terms can oscillate. Use moving averages of losses; skip gradient steps if one player dominates.
• Abstraction trivialism: If generator learns to emit generic statements ("be careful, think step by step"), it's not learning domain structure. Periodically inspect generated abstractions.
• Verifier quality: Weak verifiers give conflicting reward signals. Validate verifier accuracy >90% before full training.

Integration checklist:

• Validate solution generator baseline (without abstractions) on problem set
• Generate 50 abstractions manually; verify they're concise and useful
• Train abstraction generator on these examples
• Run training step on 10 problems; monitor both loss curves
• Periodically sample abstractions during training; verify quality
• Evaluate solution generator with abstractions vs. without (measure accuracy gain)
• Test compute allocation strategy (vary num_abstractions, measure accuracy)

Reference: https://arxiv.org/abs/2510.02263