📊 statistics

Transformers with RL or SFT Provably Learn Sparse Boolean Functions, But Differently

This paper theoretically demonstrates that while both reinforcement learning with process rewards and supervised fine-tuning enable one-layer transformers to provably learn sparse Boolean functions via Chain-of-Thought reasoning, they differ fundamentally in their learning dynamics, with RL acquiring the entire reasoning chain simultaneously while SFT learns it step-by-step.

Original authors: Bochen Lyu, Yiyang Jia, Xiaohao Cai, Zhanxing Zhu

Published 2026-05-27

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Bochen Lyu, Yiyang Jia, Xiaohao Cai, Zhanxing Zhu

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

Imagine you have a very smart but slightly confused robot (a Transformer) that needs to solve a complex puzzle. The puzzle is a Boolean function, which is just a fancy way of saying a logic problem where the answer is either "Yes" (+1) or "No" (-1). Specifically, the paper looks at "sparse" puzzles, meaning the answer only depends on a few specific pieces of information hidden among a lot of noise.

To solve these puzzles, the robot uses a strategy called Chain-of-Thought (CoT). Instead of jumping straight to the answer, it breaks the problem down into a series of small, intermediate steps, like a human thinking through a math problem step-by-step on scratch paper.

The paper investigates two different ways to teach this robot to use CoT effectively: Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL). The authors prove that both methods work, but they teach the robot in fundamentally different ways.

Here is the breakdown using simple analogies:

1. The Puzzle: Recursive Decomposition

Imagine the puzzle is a giant tree. To find the answer at the top, you have to solve small 2-piece logic problems at the bottom, then combine those answers to solve slightly bigger 2-piece problems, and so on, all the way up to the top.

The Goal: The robot needs to learn to look only at the two specific pieces of information (the "relevant" leaves) needed for each step and ignore the rest of the noise.

2. The Two Teachers

Teacher A: The Strict Drill Instructor (SFT)

Supervised Fine-Tuning (SFT) is like a teacher who gives the robot the perfect answer key for every single step of the puzzle.

How it works: The teacher says, "For step 1, the answer is X. For step 2, the answer is Y."
The Catch: The robot has to generate the answer for step 2 based on what it just wrote for step 1.
The Result (Step-by-Step Learning): The paper proves that this robot learns one step at a time.
- Analogy: Imagine trying to learn a dance routine. If you mess up the first move, you can't learn the second move because your starting position is wrong. The robot must master Step 1 perfectly before it can even start learning Step 2. It takes one training session to fix Step 1, then another session to fix Step 2, and so on. It's a slow, linear process.

Teacher B: The Process Coach (RL with Process Rewards)

Reinforcement Learning (RL) is like a coach who doesn't just look at the final score but gives feedback on every single move the robot makes.

How it works: The robot tries to solve the puzzle. If it gets a small step right, the coach gives it a "good job" reward immediately. If it gets it wrong, it gets a penalty.
The Result (Simultaneous Learning): The paper proves that this robot learns the entire chain of steps at once.
- Analogy: Imagine a coach shouting, "Good footwork on move 1! Good hand position on move 5! Bad elbow on move 3!" all at the same time. Because the robot gets specific feedback for every single step regardless of whether the previous steps were perfect, it can adjust its entire routine in a single training session. It learns the whole dance simultaneously.

3. The Big Discovery: "Process" vs. "Outcome"

The paper highlights a crucial difference in how these teachers give feedback:

SFT relies on the robot's own previous output. If the robot is wrong at the start, the "ground truth" for the next step becomes confusing noise. This forces the step-by-step learning.
RL (specifically with process rewards) gives the robot the correct "ground truth" for every step independently. It doesn't matter if the robot messed up step 1; the coach still knows what step 2 should have been and rewards/punishes accordingly. This allows for the "all-at-once" learning.

4. What About the "Hard" Puzzles?

The paper tested this on three specific types of logic puzzles:

k-PARITY: Like checking if a group of switches has an even or odd number of "on" positions. (This is notoriously hard for AI to learn without help).
k-AND: Checking if all specific switches are "on".
k-OR: Checking if at least one specific switch is "on".

The paper mathematically proves that for all three of these puzzles, both teaching methods work, provided the robot can distinguish between the "important" pieces of information and the "noise."

Summary of the Findings

Both work: You can teach a Transformer to do complex reasoning using either SFT or RL.
They are different:
- SFT is like a student who must master the basics before moving on. It learns step-by-step.
- RL (with process rewards) is like a student getting instant feedback on every specific part of the problem. It learns the whole chain simultaneously.
The Warning: If you compare SFT and RL in real life, you have to be careful. If you change how the teacher gives feedback (e.g., using a "final reward" only at the end instead of "process rewards" at every step), the learning behavior changes completely. The paper suggests that comparing these two methods requires controlling for how the rewards are designed, not just the method itself.

In short, the paper shows that while both methods can teach a robot to think logically, they do so with different "learning speeds" and "teaching styles," and understanding these differences is key to building better AI.

Technical Summary: Transformers with RL or SFT Provably Learn Sparse Boolean Functions, But Differently

Problem Statement
Large Language Models (LLMs) based on the Transformer architecture have demonstrated remarkable reasoning capabilities through Chain-of-Thought (CoT) prompting. While Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL) are the primary methods for enhancing these capabilities, the theoretical mechanisms underlying their learning dynamics remain poorly understood. Specifically, there is a lack of rigorous analysis comparing how RL (particularly with process rewards) and SFT (specifically without teacher forcing) enable Transformers to learn complex, sparse Boolean functions that are otherwise difficult to learn in an end-to-end manner. This work addresses this gap by investigating the learnability of $k$ -sparse Boolean functions—functions determined by a small subset of $k$ input bits—using a one-layer Transformer equipped with intermediate reasoning steps.

Methodology
The authors propose a unified theoretical framework to analyze the learning dynamics of both RL and SFT on a specific class of $k$ -sparse Boolean functions. These functions are recursively decomposable into fixed 2-sparse Boolean functions (e.g., parity, AND, OR), forming a reasoning chain akin to CoT.

Model Architecture: The study utilizes a simplified one-layer Transformer consisting of positional encoding, a self-attention layer, and a feedforward layer. The model is designed to iteratively generate intermediate reasoning steps ( $y^{(t)}$ ) before producing a final answer.
"Pretrained" Mask: To ensure theoretical tractability and simulate the structural priors gained from pretraining, the authors impose a specific "pretrained" mask on the attention weights. This mask restricts the attention of reasoning tokens to their relevant child nodes (as defined by the recursive decomposition) and the input sequence, effectively controlling error propagation.
Learning Paradigms:
- RL with Process Reward: The Transformer is optimized via policy gradient to maximize a reward function where rewards are assigned to each intermediate reasoning step ( $r_t$ ), rather than just the final outcome.
- SFT without Teacher Forcing: The Transformer is trained to minimize a loss (hinge loss) over generated reasoning chains compared to ground-truth labels. Crucially, the model generates its own intermediate steps during training (autoregressive generation) rather than being fed ground-truth prefixes (teacher forcing), addressing the train-inference mismatch.
Core Analytical Tool: The authors introduce a "critical gradient component" ( $\gamma$ ). They demonstrate that the separation of this component—where the gradient signal for relevant positions (child nodes) is strictly greater than for irrelevant positions—serves as a sufficient condition for provable learnability.

Key Contributions and Results
The paper establishes sufficient conditions under which a one-layer Transformer can provably learn $k$ -sparse Boolean functions via both RL and SFT. The primary findings are:

Unified Learnability Condition: Both RL and SFT succeed if the critical gradient component separates relevant from irrelevant positions. This condition is verified for three fundamental examples: $k$ -PARITY, $k$ -AND, and $k$ -OR.
Distinct Learning Dynamics:
- RL (Process Reward): Under the separation condition, RL learns the entire CoT reasoning chain simultaneously in a single gradient update. The process reward provides informative signals for all steps regardless of the current state of other steps, allowing the model to acquire the full reasoning capability immediately.
- SFT (No Teacher Forcing): In contrast, SFT exhibits a step-by-step learning behavior. Due to the autoregressive nature of generation, the model cannot reliably learn step $t$ until step $t-1$ is correctly generated. Consequently, SFT requires $T$ gradient updates (where $T$ is the chain length) to learn the entire chain, solving one step per update.
Hardness of Final Reward: The paper proves that RL optimized with only a final (outcome-based) reward fails to learn these functions efficiently. The policy gradient variance is exponentially small, drowning the signal in noise, which explains the necessity of process rewards for complex reasoning tasks.
Extension to Serial CoT: The analysis is extended beyond the binary-tree decomposition to serial CoT structures (common in practical LLMs), confirming that the same learning dynamics (simultaneous for RL, step-wise for SFT) hold under appropriate masking.

Significance and Claims
The authors claim that this work provides the first theoretical characterization of the learning dynamics of RL with process rewards for Transformers, filling a critical gap in the literature. By removing teacher forcing and data augmentation strategies often used in prior theoretical works (e.g., Kim & Suzuki, 2025), the study offers a more realistic comparison between RL and SFT.

The significance lies in revealing that while both methods can achieve learnability, they do so through fundamentally different mechanisms:

RL with process rewards mimics the efficiency of SFT with teacher forcing by learning the whole chain simultaneously.
SFT without teacher forcing naturally induces a step-wise curriculum, an intrinsic property of autoregressive training that does not require external curriculum learning or data filtering.

The paper concludes that empirical comparisons between RL and SFT must carefully control for factors such as reward design (process vs. outcome) and training regimes (teacher forcing vs. autoregressive generation), as these factors significantly alter learning dynamics and may confound the comparison of the optimization methods themselves.