Decoupling Knowledge and Reasoning in Transformers: A Modular Architecture with Generalized Cross-Attention

TL;DR

This paper proposes a modular Transformer architecture that explicitly separates knowledge retrieval from reasoning, using a generalized cross-attention mechanism to improve interpretability, adaptability, and scalability in neural models.

Contribution

It introduces a novel modular Transformer design with a generalized cross-attention mechanism that decouples knowledge and reasoning, supported by a mathematical derivation relating FFNs to this framework.

Findings

Demonstrates that FFNs are a special case of generalized cross-attention.

Provides a theoretical foundation for understanding and improving Transformer interpretability.

Lays groundwork for future scalable and adaptable Transformer architectures.

Abstract

Transformers have achieved remarkable success across diverse domains, but their monolithic architecture presents challenges in interpretability, adaptability, and scalability. This paper introduces a novel modular Transformer architecture that explicitly decouples knowledge and reasoning through a generalized cross-attention mechanism to a globally shared knowledge base with layer-specific transformations, specifically designed for effective knowledge retrieval. Critically, we provide a rigorous mathematical derivation demonstrating that the Feed-Forward Network (FFN) in a standard Transformer is a specialized case (a closure) of this generalized cross-attention, revealing its role in implicit knowledge retrieval and validating our design. This theoretical framework provides a new lens for understanding FFNs and lays the foundation for future research exploring enhanced interpretability, adaptability, and scalability, enabling richer interplay with external knowledge bases and other systems.