Can Large Language Models Learn Independent Causal Mechanisms? (2024)

Abstract: 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 a lack of generalisation ability. 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 show that such causal constraints can improve out-of-distribution performance on abstract and causal reasoning tasks. We also investigate the level of independence and domain specialisation and show that LLMs rely on pre-trained partially domain-invariant mechanisms resilient to fine-tuning.

## introduction the latest generation of large language models (llms) with over several billion parameters has demonstrated impressive performance on an extensive range of in-context language and reasoning tasks @xcite @xcite and an even greater range when fine-tuned for a specific task @xcite . however, these observations do not hold for tasks that fall outside the training data distribution, sometimes even when the task is only slightly perturbed. in particular, standard llms perform poorly on complex reasoning tasks, such as abstract, causal, or logical reasoning @xcite @xcite @xcite . @xcite ; @xcite ; @xcite showed that finetuning llms can increase their in-distribution performance, but the improvement does not transfer to different distributions, highlighting that llms do not generalise as we might expect a person to when applied to domains requiring complex reasoning. several hypotheses have been proposed to explain this flaw, such as the lack of abstract or symbolic representations within the latent space of llms @xcite @xcite . these claims are supported by the brittleness that llms can exhibit; when changing the wording of a question, the performance of an llm can vary drastically @xcite . this observation hints that llms may rely on domain-specific information or spurious correlations in the training data that do not generalise to other distributions. causal models rely on the concept that causal mechanisms invariant under changes in environment exist. the independent causal mechanisms principle further states that "the causal generative process of a system's variables is composed of autonomous modules that do not inform or influence each other. " @xcite . these principles are applied in diverse ways in the field of causality, either in the structure of the model, which may be built in a modular fashion to respect causal relationships, as in structural causal models @xcite , or in the distribution of the data, which may be rendered independent and identically distributed (i.i.d) from an unbalanced distribution by division into subgroups @xcite . integrating these methods into the architecture of a large language model could increase its robustness and out-of-distribution (o.o.d) generalisation. we investigate this idea in this work: we aim to better understand how llms reason in and outof-distribution and whether they can behave as models of independent causal mechanisms un-der certain constraints and with fine-tuning. to this end, we propose an llm architecture integrating the concept of mechanisms as independent, self-contained llm modules. this model is summarised in figure 1 . we aim to answer the following questions: (i) can llms be used as self-routers for specialised mechanisms, and does it improve their performance? (ii) can llms capture domaininvariant abstractions with information-based regularisation? (iii) how useful is domain-specific knowledge on reasoning tasks? and (iv) can our proposed architecture approximate independent causal mechanisms? our contributions can be summarised as follows: our code is available at https://github.com/ strong-ai-lab/modular-lm . ## causal information routing for llms we now describe our proposed modular architecture: independent causal language models (iclm), where each module is an llm fine-tuned for a specific specialisation or generalisation objective. we aim to build a system that can adapt to changing distributions and capture better abstractions. our architecture is separated into n + 2 llm modules connected by three main components. the llm modules are composed of a router that generates embeddings of the inputs, a domaininvariant module trained to learn abstractions and n domain-specific modules trained to specialise on a single task or domain. the other components are the routing strategy, mutual information loss and the aggregation scheme. the routing strategy uses the embeddings from the llm router to redirect the inputs to a specific module. specifically, routing is performed in an unsupervised fashion: each embedding is projected into a clustered space. the centroid of each cluster is associated with a domain-specific module to which it assigns a binary activation weight for a given input. if the input belongs to a cluster, the corresponding module is activated. by contrast, the domain-invariant module processes all inputs. the mutual information loss induces abstraction within the domain-invariant module; minimising this loss reduces shared information between the domainspecific and domain-invariant modules. i.e. it is intended to cause the domain-invariant module to gain domain-invariant knowledge and the domainspecific modules domain-specific knowledge. finally, the aggregation scheme combines the output of the activated domain-specific module and the domain-invariant module to produce the final output. figure 1 shows an overview of our method. we describe the routing strategy in section 3.1, the mutual information minimisation in section 3.2 and the aggregation scheme in section 3.3. section 4 discusses how the architecture reflects indepen-dent causal mechanisms. ## theoretical perspective in this section, we provide theoretical evidence on how our model approximates independent causal mechanisms and under what assumptions. independent causal mechanisms consist of autonomous modules that work independently. in our case, all domain-specific modules are trained for specific tasks/distributions. the domain-invariant module is trained only to use domain-invariant knowledge. the router module is tasked to split the input distribution into n more balanced distributions. we aim to verify that the modules are not causally related. more formally, we aim to study under what conditions the following holds: h r , h i and h sn ∀@xmath0, n ] are the respective representations generated by the router, domaininvariant and n domain-specific modules. equations 5 and 6 are verified. the proof is provided in appendix b; the main idea is that the use of a separate loss function for training the router prevents the other modules from causally acting on the router, either in the forward or backward passes. it can be verified in figure 2 . however, if an invariant module is part of the model, equations 7 and 8 do not hold. the domain-specific modules do not directly influence each other because the routing mechanism allows a single module to go through the forward and backward passes. nevertheless, a causal path can be drawn through the domain-invariant module as it is always activated. for example, assuming a model with two domainspecific modules, s 0 and s 1 , activated one after the other, a path exists and can be represented in a simplified version as h s 1 ... again, details of the proof are given in appendix b. as h i and h sn are causally related, we need to reduce the dependency between the two quantities using a regularisation term. minimising the mutual information between h i and h sn amounts to reducing the mutual dependence between the variables. i(h i , h sn ) = 0 if and only if h i and h sn are independent. if verified, the loss l y can be divided into two independent components and equations 7 and 8 hold. we verify experimentally in appendix e.1 that the mutual information is close to zero after ∼ 50 training steps. ## routing and independence analysis 6.1 evolution of module independence we study the independence of the module hidden states during training (figure 3 ) and inference (figure 4 ). we focus on two complementary measures: mutual information and the pearson correlation coefficient that measures linear correlation between variables. the latter is limited to linear dependence but is more easily interpretable. the shared mutual information as well as the pearson correlation coefficient between modules are effectively reduced by the regularisation scheme during fine-tuning. however, the module hidden states remain correlated, in particular at inference time. further investigation in appendix e.5 further shows that this correlation is maintained across most layers, which indicates the presence of a general domain-invariant mechanism shared by all modules and composing the basis of their reasoning abilities. its influence is reduced via the fine-tuning procedure that develops domainspecific knowledge but it remains the main mechanism used, particularly at test time. ## conclusion performing strong out-of-distribution reasoning is a challenging task, and despite their impressive performance on a wide range of problems, llms have not demonstrated this ability yet. combining this popular model with causal models could help bridge this gap. this work presents a modular architecture yielding llms to behave as independent causal mechanisms. we show theoretically that the proposed model generates causally-independent modules. we perform experiments on abstract and