Can Large Language Models Learn Independent Causal Mechanisms?
Gaël Gendron Bao Trung Nguyen Alex Yuxuan Peng Michael Witbrock Gillian Dobbie
Despite impressive performance on language modelling and complex reasoning
tasks, Large Language Models (LLMs) fall short on the same tasks in uncommon
settings or with distribution shifts, exhibiting some lack of generalisation
ability. This issue has usually been alleviated by feeding more training data
into the LLM. However, this method is brittle, as the scope of tasks may not be
readily predictable or may evolve, and updating the model with new data
generally requires extensive additional training. By contrast, systems, such as
causal models, that learn abstract variables and causal relationships can
demonstrate increased robustness against changes in the distribution. One
reason for this success is the existence and use of Independent Causal
Mechanisms (ICMs) representing high-level concepts that only sparsely interact.
In this work, we apply two concepts from causality to learn ICMs within LLMs.
We develop a new LLM architecture composed of multiple sparsely interacting
language modelling modules. We introduce a routing scheme to induce
specialisation of the network into domain-specific modules. We also present a
Mutual Information minimisation objective that trains a separate module to
learn abstraction and domain-invariant mechanisms. We show that such causal
constraints can improve out-of-distribution performance on abstract and causal
reasoning tasks.