Are Transformers Able to Reason by Connecting Separated Knowledge in Training Data?

Abstract: Humans exhibit remarkable compositional reasoning by integrating knowledge from various sources. For example, if someone learns ( B = f(A) ) from one source and ( C = g(B) ) from another, they can deduce ( C=g(B)=g(f(A)) ) even without encountering ( ABC ) together, showcasing the generalization ability of human intelligence. In this paper, we introduce a synthetic learning task, "FTCT" (Fragmented at Training, Chained at Testing), to validate the potential of Transformers in replicating this skill and interpret its inner mechanism. In the training phase, data consist of separated knowledge fragments from an overall causal graph. During testing, Transformers must infer complete causal graph traces by integrating these fragments. Our findings demonstrate that few-shot Chain-of-Thought prompting enables Transformers to perform compositional reasoning on FTCT by revealing correct combinations of fragments, even if such combinations were absent in the training data. Furthermore, the emergence of compositional reasoning ability is strongly correlated with the model complexity and training-testing data similarity. We propose, both theoretically and empirically, that Transformers learn an underlying generalizable program from training, enabling effective compositional reasoning during testing.

• The paper demonstrates that few-shot chain-of-thought prompting significantly improves Transformers' ability to deduce causal relationships by revealing the correct vertex order.
• Key results reveal a phase transition in performance at a relative knowledge ratio of 0.3, highlighting the necessity of multi-layer attention mechanisms.
• Empirical and theoretical findings indicate that Transformers learn an underlying program, using induction heads to copy and integrate fragmented data efficiently.

Are Transformers Able to Reason by Connecting Separated Knowledge in Training Data?

Introduction

The paper investigates whether Transformer-based models can demonstrate compositional reasoning by integrating knowledge fragments from separate training instances. The authors introduce a synthetic dataset, "Fragmented at Training, Chained at Testing" (FTCT), to evaluate this capability. This dataset allows them to explore when Transformers can perform compositional reasoning, how training factors affect this ability, and what internal mechanisms are involved.

FTCT Dataset

The FTCT dataset is designed to test compositional reasoning by simulating causal relationships in a graph structure:

• Causal Structure: Knowledge relationships are represented as a directed graph where vertices denote knowledge points, and edges define causal relationships.
• Fragmented Training: The training data consists of subgraphs (child chains) with additional noise vertices to mimic incomplete knowledge.
• Chained Testing: The testing phase uses complete causal chains to assess compositional reasoning ability.
• Few-Shot Learning: Examples for training and testing are concatenated into sequences to facilitate few-shot learning.

  Figure 1: Overview of the FTCT Data Generation Process.

Key Findings

Few-Shot Chain-of-Thought (CoT) Prompting

Few-shot CoT prompting effectively enables Transformers to perform compositional reasoning. The study finds that Transformers struggle with zero-shot prompts but significantly improve with few-shot examples. This is primarily because few-shot CoT prompts reveal the correct order of vertices, allowing the model to deduce their values based on preceding reasoning paths.

Impact of Data Distribution and Model Complexity

Compositional reasoning emerges when there is sufficient similarity between training and testing data, measured by the relative knowledge ratio ($M/N$). A critical phase transition occurs when this ratio reaches 0.3. Additionally, multi-layer attention mechanisms with a minimum of 2 layers and 2 heads are essential for achieving compositional reasoning.

Learning an Underlying Program

The authors propose that Transformers learn an underlying program during training, which minimizes the loss for both training and testing by capturing common latent structures. This is corroborated by theoretical and empirical evidence showing the use of induction heads for in-context learning and specific attention patterns for retrieving and integrating knowledge fragments.

Empirical Evidence

Empirical evaluation shows that:

• Zero and Few-Shot Performance: Testing performance improves with the number of shots due to the model's ability to copy vertex order (induction heads) and correctly deduce values (attention assignment).
• Attention Patterns: Attention heatmaps reveal that induction heads in different layers facilitate copying patterns. Layer-specific attention to past vertices and values assists in retrieving necessary information for reasoning tasks.

Figure 2: Left: Zero and few-shot testing performance with various causal depths and child chain lengths. Right: Relationship between the relative knowledge ratio and compositional reasoning ability.

Limitations and Future Work

While the FTCT dataset facilitates controlled experimentation, the conclusions drawn may not directly apply to real-world tasks involving natural language. Future work could explore the integration of more complex causal structures and noise types to further validate these findings. Additionally, applying these insights to larger and more complex LLMs remains an open area for exploration.

Conclusion

The research validates the potential of Transformers to perform compositional reasoning through careful experimentation with the FTCT dataset. Few-shot CoT prompts play a crucial role, and the Transformers' ability to learn an underlying program is instrumental in reasoning with fragmented training data. The study highlights a phase transition in reasoning capabilities and underscores the importance of model architecture in achieving these results.