Bilinear Convolution Decomposition for Causal RL Interpretability

TL;DR

This paper introduces bilinear convolutional models in reinforcement learning that enhance interpretability through analytic decomposition and causal validation, enabling better understanding of model decisions.

Contribution

It proposes replacing nonlinearities with bilinear variants in ConvNets for RL, enabling analytic decomposition and causal validation of model interpretability.

Findings

Bilinear models perform comparably to traditional models in RL tasks.

Decomposition techniques reveal interpretable low-rank structures.

Methodology allows causal validation of concept-based probes.

Abstract

Efforts to interpret reinforcement learning (RL) models often rely on high-level techniques such as attribution or probing, which provide only correlational insights and coarse causal control. This work proposes replacing nonlinearities in convolutional neural networks (ConvNets) with bilinear variants, to produce a class of models for which these limitations can be addressed. We show bilinear model variants perform comparably in model-free reinforcement learning settings, and give a side by side comparison on ProcGen environments. Bilinear layers' analytic structure enables weight-based decomposition. Previous work has shown bilinearity enables quantifying functional importance through eigendecomposition, to identify interpretable low rank structure. We show how to adapt the decomposition to convolution layers by applying singular value decomposition to vectors of interest, to separate the channel and spatial dimensions. Finally, we propose a methodology for causally validating concept-based probes, and illustrate its utility by studying a maze-solving agent's ability to track a cheese object.