Transformers with RL or SFT Provably Learn Sparse Boolean Functions, But Differently
This provides theoretical insights into fine-tuning mechanisms for transformers, addressing a gap in understanding for researchers in machine learning and AI, though it is incremental as it builds on existing methods for specific functions.
The paper tackles the problem of understanding how transformers learn Chain-of-Thought capabilities via reinforcement learning (RL) and supervised fine-tuning (SFT) for sparse Boolean functions, showing that both methods provably learn functions like k-PARITY, k-AND, and k-OR, but RL learns the chain simultaneously while SFT learns step-by-step.
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.