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

AI summary

Overview

Research area: Theoretical machine learning — specifically the learning theory of transformers, Chain-of-Thought (CoT) reasoning, reinforcement learning (RL) fine-tuning, and supervised fine-tuning (SFT).

Technical level: Advanced. The paper is a theoretical analysis built on gradient dynamics, policy gradient derivations, and error-bound proofs for a one-layer transformer; it presumes comfort with attention mechanics and optimization theory.

Scope (one sentence): The paper gives a unified theoretical analysis of when a one-layer transformer provably learns recursively decomposable k-sparse Boolean functions via RL with verifiable process rewards versus SFT without teacher forcing, and shows the two methods learn the CoT chain differently.

What This Paper Is About

Transformers can be fine-tuned to produce intermediate reasoning steps (Chain-of-Thought) using either reinforcement learning or supervised fine-tuning, but there is little theory explaining how these two approaches actually acquire reasoning capability or how they differ. The authors study a class of hard-to-learn functions — k-sparse Boolean functions that decompose recursively into a fixed 2-sparse Boolean function — and ask under what conditions a one-layer transformer provably learns them through RL with verifiable process rewards and through SFT without teacher forcing. Their answer is a single sufficient condition, the "separation of the critical gradient component," that governs learnability in both cases while producing different learning dynamics.

Key Contributions

1. A unifying decomposition framework. For k-sparse Boolean functions that are recursively decomposable using a fixed 2-sparse Boolean function, the authors decompose learning into multi-step reasoning tasks, modeled as a complete binary tree of height T with 2k−1 nodes, where the t-th level has d_t = k/2^t nodes. This lets a whole class of problems be analyzed together rather than one function at a time.

2. A unified analysis of RL and SFT learning dynamics via a new "critical gradient component." The authors identify a critical gradient component and show that its separation across relevant versus irrelevant positions is a sufficient condition for learnability under both RL with verifiable process rewards and SFT.

3. A proven distinction in learning behavior. For RL optimized by policy gradient, the transformer learns the target function after a single gradient update under the separation condition (Theorem 3.1), acquiring the whole reasoning chain simultaneously. For SFT in the absence of both teacher forcing and data augmentation, learnability holds under a similar separation condition (Theorem 3.3), but the transformer learns the chain step-by-step, requiring one gradient update per step and therefore a total number of updates equal to the chain length.

4. Verification on three concrete functions. The authors verify that the separation condition of the critical gradient component is satisfied for k-PARITY, k-AND, and k-OR, for both RL and SFT, demonstrating learnability of these specific functions.

Main Findings

• RL learns the entire CoT chain at once. Theorem 3.1 shows that fine-tuning a one-layer transformer via RL optimized by the sign of the policy gradient, with initialization W(0) = 1 and learning rate η = Ω(ln(d/ε)) for ε > 0, achieves ‖softmax(W(1)) − softmax(W*)‖₁ ≤ ε after one update, provided the separation condition in Eq. (8) holds. The result is described as the first study of learnability of k-sparse Boolean functions for transformers with CoT through RL fine-tuning.

• SFT without teacher forcing learns step-by-step. Theorem 3.3 guarantees SFT learnability under a similar separation condition, but the transformer naturally exhibits step-wise learning: one gradient update per reasoning step, with total updates equal to the chain length.

• RL with process reward resembles SFT with teacher forcing. Because each sampled trajectory step receives feedback conditioned on previously generated steps, RL with verifiable process rewards can learn the whole chain simultaneously, matching the simultaneous learning that SFT with teacher forcing exhibits.

• RL with only a final reward is hard. Proposition 3.2 formalizes the hardness of RL with final reward r^final = y^(T) Φ_k(x), which provides supervision only for the final answer. For a function class H whose members are orthogonal (E_x[h(x)h′(x)] = 0 for distinct h, h′) under a bounded-gradient assumption E_x[‖∇W E{y~p_W(·|x)}[y^(T)]‖²] ≤ M, the policy gradient variance is bounded in a way that limits target-specific information.

• RL versus SFT is not the deciding factor. The paper states explicitly that the RL/SFT distinction alone does not determine learning dynamics (step-by-step versus simultaneous). What matters is whether each reasoning step receives useful feedback for the state the model is trained on — which depends on whether supervision is applied only to final answers, under teacher-forced prefixes, or along self-generated trajectories.

• The separation condition holds for three basic functions. For k-PARITY (Φ_k^parity(x) = ∏{i∈B} x_i), k-AND (Φ_k^and(x) = 2∏{i∈B}((x_i+1)/2) − 1), and k-OR (Φ_k^or(x) = 1 − 2∏_{i∈B}((1−x_i)/2)), the authors verify the separation condition for both RL and SFT.

• The framework extends beyond balanced trees. A serial CoT-style decomposition is also analyzed, in which y_1 = φ₂(x_{i₁}, x_{i₂}) and y_j = φ₂(y_{j−1}, x_{i_{j+1}}) for j ≥ 2, forming a length-(k−1) reasoning chain that does not require k = 2^T. The same gradient separation analysis applies with a different "pretrained" mask, so RL with process rewards and SFT with teacher forcing learn the serial chain simultaneously, while SFT without teacher forcing learns step by step under an appropriate filtering rule.

• Not reported in the available content: the truncated text cuts off inside Proposition 3.2 and does not contain the numeric results of the experiments referenced for Appendix D, nor the details of the FAQ (Appendix B) or the extended related-work discussion (Appendix C).

Methodology in Plain English

The authors construct a deliberately simplified setting so that learning can be analyzed mathematically.

First, they restrict attention to k-sparse Boolean functions: functions of a d-bit input x ∈ {+1,−1}^d that depend only on a fixed subset B of size k (with 2 ≤ |B| = k ≤ d), where the input follows a uniform distribution. They consider functions that can be broken into a chain of smaller sub-problems, each applying the same fixed 2-sparse Boolean function φ₂ to two inputs. Assuming k = 2^T, this yields a complete binary tree of height T with 2k−1 nodes, where each level of the tree is one reasoning step.

Second, they use a one-layer transformer consisting of a positional encoding, a single-head self-attention layer without residual connection, and a feedforward layer with activation ψ: [−1,1] → [0,1] whose form depends on φ₂. The input sequence has length d + k − 1, and intermediate reasoning tokens are initialized to 0 and generated autoregressively by feeding the transformer's own outputs back in — a CoT-style loop. A "pretrained mask" (a mask, not actual pretraining) restricts attention so that each reasoning step depends only on the previous step, capturing the notion that pretraining provides structural priors for fine-tuning.

Third, they set up RL as maximizing an expected reward over sampled trajectories, using a per-step process reward r_t = (1/(k−1)) Σ_{j=1}^{d_t} y_j^(t) ȳ_j^(t), where ȳ_j^(t) is the correct label given the earlier steps. They derive the policy gradient, simplify the analysis by using the sign of the policy gradient for the parameter update, and characterize the optimal parameters W*, where R(W*) = 1 and the attention softmax assigns weight 1/2 to each of the two relevant child positions and 0 elsewhere.

Fourth, they identify the critical gradient component γ, which depends on the activation derivative ψ′(ξ), the 2-sparse function φ₂, and the previous-step token. If this component is separated — larger at relevant positions (the child nodes) than at irrelevant ones, in expectation — then learning succeeds. They prove Theorem 3.1 (RL, one update) and Theorem 3.3 (SFT, step-by-step) from this condition, contrast it with Proposition 3.2 (hardness of final-reward RL), and then check the condition by hand for k-PARITY, k-AND, and k-OR.

Why This Matters

Impact on research. The paper supplies a theory-level account of why RL fine-tuning and SFT can both produce CoT reasoning yet behave differently during training, and it argues that the method label alone (RL vs. SFT) is the wrong explanatory variable — what matters is whether each intermediate step receives useful feedback. It generalizes earlier theoretical work on k-PARITY, which studied SFT with teacher forcing and SFT without teacher forcing but with data augmentation and self-verification filtering, to a broader class of recursively decomposable sparse Boolean functions, and it introduces a single diagnostic (separation of the critical gradient component) that applies across methods.

Real-world applications (implications drawn from the paper's framing, not benchmarks reported in the text):

• Designing process reward models and step-level verifiers for reasoning systems, given the theoretical contrast between process reward and final-answer reward.
• Deciding between RL and SFT pipelines when fine-tuning models for multi-step tasks such as mathematical problem solving, where correctness can be checked step by step.
• Diagnosing why some fine-tuning runs learn an entire reasoning chain quickly while others appear to learn one step at a time and stall.
• Structuring training data with intermediate supervision rather than final answers alone, since the analysis ties learnability to whether steps receive feedback along the trajectories the model actually visits.

Industry relevance. The comparison between RL and SFT is central to how reasoning-capable models are trained, and the "SFT memorizes but RL generalizes" intuition is widely cited. This paper's stance — that empirical comparisons should account for intermediate supervision rather than attribute behavior solely to RL or SFT — is directly relevant to teams choosing training objectives and reward designs.

Future Directions

• Relax the architectural simplifications. The analysis uses a one-layer transformer with a single attention head, no residual connection, and a shortened key-query parameterization W_KQ = W_K^T W_Q; extending the separation-condition analysis to deeper transformers, multiple heads, and residuals is a natural open problem.

• Move from sign-of-gradient updates to full stochastic policy gradients. Theorem 3.1 analyzes the sign of the policy gradient; whether the same one-update learnability holds for the unsimplified policy gradient update is left open by the summarized content.

• Bridge to the reported experiments. Experiments are stated to appear in Appendix D but are not included in the available content; follow-up work would test whether the predicted step-by-step versus simultaneous behaviors are observable in trained models at the scales used in practice.

• Broaden the function class. The paper analyzes functions recursively decomposable into a fixed 2-sparse function, including k-PARITY, k-AND, and k-OR. Extending the framework to other hard classes and to the serial CoT decomposition with alternative filtering rules are explicitly flagged as extensions.

Target Audience

The paper is aimed at machine learning theorists and researchers studying the foundations of reasoning in transformers — particularly those working on expressivity, learnability, gradient dynamics, and the theory of CoT. It is also relevant to practitioners who design RL fine-tuning pipelines with process rewards or SFT data with intermediate supervision and want a principled account of when each approach can be expected to learn an entire reasoning chain at once versus one step at a time. Readers need a background in optimization and attention mechanics; it is not an introductory paper.

Authors’ abstract

Transformers can acquire Chain-of-Thought (CoT) capabilities to solve reasoning tasks via fine-tuning. Reinforcement learning (RL) and supervised fine-tuning (SFT) are two primary approaches to this end. In this work, we examine RL with verifiable process rewards and SFT for learning $k$-sparse Boolean functions with a one-layer transformer through intermediate reasoning steps akin to CoT. In particular, we consider Boolean functions that can be recursively decomposed into fixed 2-sparse Boolean functions. We first analyze the learning dynamics of RL fine-tuning with verifiable process rewards and SFT in a unified way, allowing us to identify sufficient conditions under which the transformer provably learns these functions. We then verify that the conditions hold for three examples, including $k$-PARITY, $k$-AND, and $k$-OR, thus demonstrating their learnability via both RL and SFT. Notably, we reveal that RL and SFT exhibit distinct learning behaviors depending on supervision: RL learns the whole CoT chain simultaneously, whereas SFT without teacher forcing learns the CoT step-by-step. Overall, our findings provide insights on the mechanisms underlying RL and SFT and how they differ in triggering the CoT capabilities of transformers, and suggest that the comparison between RL and SFT should consider the intermediate supervision.