Reframing attention as a reinforcement learning problem for causal discovery

RLC 2025 Workshop CRLW Submission6 Authors

Keywords: Causal RL, 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. This allows us
to reformulate the attention mechanism popularized by Transformer networks within an
RL setting with the goal 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 original 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, and uniquely recover graphs of dynamic
causal processes.

Submission Number: 6