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

TL;DR

This paper provides a theoretical comparison of reinforcement learning and supervised fine-tuning in transformers, showing they learn sparse Boolean functions differently, with implications for understanding Chain-of-Thought reasoning.

Contribution

It offers the first theoretical analysis of how RL and SFT enable transformers to learn sparse Boolean functions, revealing their distinct learning dynamics.

Findings

RL learns the entire CoT chain simultaneously

SFT learns the CoT chain step-by-step

Both approaches can learn k-sparse Boolean functions under certain conditions

Abstract

Transformers can acquire Chain-of-Thought (CoT) capabilities to solve complex reasoning tasks through fine-tuning. Reinforcement learning (RL) and supervised fine-tuning (SFT) are two primary approaches to this end, yet their underlying mechanisms and differences remain theoretically unclear. In this work, we examine these aspects specifically for learning $k$-sparse Boolean functions with a one-layer transformer and intermediate supervision that is akin to CoT. In particular, we consider $k$-sparse Boolean functions that can be recursively decomposed into fixed 2-sparse Boolean functions. We analyze the learning dynamics of fine-tuning the transformer via either RL or SFT with CoT to identify sufficient conditions for it to provably learn these functions. We verify that these conditions hold for three basic examples, including $k$-PARITY, $k$-AND, and $k$-OR, thus demonstrating the learnability of both approaches. Notably, we reveal that RL and SFT exhibit distinct learning behaviors: RL learns the whole CoT chain simultaneously, whereas SFT learns the CoT chain step-by-step. Overall, our findings provide theoretical insights into the underlying mechanisms of RL and SFT as well as how they differ in triggering the CoT capabilities of transformers.