Reframing attention as a reinforcement learning problem for causal discovery

ICLR 2026 Conference Submission25201 Authors

Keywords: Causal World Models, Causal Reinforcement Learning, Causal Processes, Causal Representation Learning

Abstract: Formal frameworks of causality have operated largely parallel to modern trends in deep reinforcement learning (RL). However, there has been a revival of interest in formally grounding the representations learned by neural networks in causal concepts. Yet, most attempts at neural models of causality assume static causal graphs and ignore the dynamic nature of causal interactions. In this work, we introduce Causal Process framework as a novel theory for representing dynamic hypotheses about causal structure. Furthermore, we present Causal Process Model as an implementation of this framework. Leveraging the inherent causality of the RL framework, we reformulate the attention mechanism from Transformer networks within an RL setting to infer interpretable causal processes from visual observations. Here, causal inference corresponds to constructing a causal graph hypothesis, which itself becomes an RL task nested within the original RL problem. To create an instance of such hypothesis, we employ RL agents. These agents establish links between units similar to the Transformer attention mechanism. We demonstrate the effectiveness of our approach in an RL environment where we outperform current alternatives in causal representation learning and agent performance.

Primary Area: causal reasoning

Submission Number: 25201