SlowGuess/NeurIPS_2023_OCR huggingface dataset:
(a) Environment & task structure
Figure 3: The odd-one-out meta-learning environments, where the goal is to pick up an object that is unique along the "correct" dimension, which cannot be directly observed. (a) An example of the observations and task structure. At top is an example experimentation trial. The agent is the white square (the view is egocentric, i.e. the agent is always in the center); it transforms one object to a unique shape, then it will "pick up" that object to check whether shape is the relevant dimension for this episode. After three experimentation trials, the agent proceeds to an exploit trial (bottom), where it can no longer transform objects, but where each object is unique along a different dimension; it must pick up an object according to the latent causal dimension discovered through experimentation. In training, the environment provides explanations of the objects and the results of experiments, which support learning. (b) Generalization to novel feature combinations. (c) Generalization from training data where 0-2 dimensions have a hard initialization (see text) to the case when all 3 dimensions do. (Panel (a) adapted with permission from Lampinen et al. [41]. 3 seeds per condition.)
(b) Feature set generalization
(c) Hard dimension generalization
The above experiments used simple environments, where the agent directly observes the variables and the outcomes of its interventions. Here, we extend to a more complex set of environments involving
a causal structure that can only be understood by interpreting relations among objects in high-dimensional pixel observations. Specifically, we experiment on the causal intervention/meta-learning versions of the odd-one-out environments [41]. See Fig. 3a for some example observations.
In these environments, objects vary along multiple dimensions (color, shape, and texture). In each episode there is one dimension that is "correct"—this "correct" dimension is a latent causal factor that the agent cannot directly observe. The agent is rewarded for picking up an object that is unique (the odd one out) along the "correct" feature dimension. The agent is first given a series of trials in which to experiment and determine which dimension is currently "correct." Specifically, the agent is given a magic wand that allows it to transform features of objects, thus allowing it to make an object unique along one feature dimension. It can then pick up this object, and it can infer from whether it receives reward whether the dimension it tried is correct. These experiments allow the agent to succeed on a final exploit trial, where it no longer has access to a magic wand, but instead is faced with multiple objects that are each unique along a different dimension. To achieve a reward, the agent must use the results of its earlier experiments to choose the object that is unique in the correct way.
We evaluated two types of generalization. First, we simply split possible feature combinations into train and test sets. We evaluate whether the agent can learn, from passive data, an experimentation and exploitation strategy that generalizes to novel combinations. In Fig. 3b, we show that the agent is able to learn these tasks well, and to generalize equally well to the evaluation conditions. Note that the tasks are designed such that simpler strategies — such as relying on feature values — will not yield above chance performance.
We then evaluated a more challenging generalization condition. The original paper [41] considered harder object initializations that make the experiments more difficult, by forcing the agent to transform one object, and then choose a different object that is made unique by that transformation. We create a more challenging relational generalization split by training the agent in settings where 0-2 of the feature dimensions (randomly sampled in each training episode in the dataset) have this harder initialization. We then evaluate the agent's ability to generalize to evaluation in a harder setting where all 3 feature dimensions have this difficult initialization. This condition tests the agent's generalization to settings more challenging than those observed in the data. In Fig. 3c, we show that the agent is able to learn and generalize these more difficult tasks.
Explanations support learning: In these experiments, we followed the method of Lampinen et al. [41]; the agent is trained to predict natural language explanations (provided by the environment during training) as an auxiliary loss. For example, when the agent performs a successful experiment that identifies the correct dimension as shape, during training this will be accompanied by an explanation like "correct, the latent feature is shape, and this object is uniquely square." In Supp. Fig. 12 we show that these explanations are beneficial even when learning from expert data; without predicting explanations, the agents do not learn or generalize these tasks as readily. Furthermore, explanations can enable other kinds of generalization from passive learning, as we explore in the next section.
5.2.1 Explanations can deconfound for passive learners too
The original odd-one-out work also explored a setting of learning from perfectly confounded experiences, in which a rewarding object is unique along three feature dimensions, so it is unclear which feature(s) are relevant. The authors showed that RL agents can generalize out of distribution to follow any particular dimension in a deconfounded test (where different objects are unique along each dimension), if they are trained to predict language explanations that focus on that particular feature dimension. Although these experiments do not directly involve causal intervention, they are relevant to the broad question of what can be learned from confounded data — and are particularly relevant to LMs, because explanations are a common feature of the human language from which these models learn. However, the prior work only experimented with active RL agents. Can passive learners also generalize out-of-distribution from predicting explanations?
Figure 4: Explanations can deconfound for passive learners trained on confounded data — the agents generalize OOD in accordance with the explanations they were trained to predict.
We reproduced these deconfounding experiments in our passive learning setup, using the same expert as above. We find that passive learners also generalize consistently OOD according to the explanations they predicted during training (Fig. 4). Passive learners that were trained to predict explanations about shape would tend to choose the uniquely-shaped object in the deconfounded evaluation, and agents trained to predict explanations about texture or shape would analogously generalize according to those features. Thus, natural language explanations encountered during passive training can shape the way a system generalizes out-of-distribution at test time.
5.3 Pretrained language models
Finally, we tested whether pretrained language models could generalize causal strategies from a few-shot prompt. To do so, we adapt the odd-one-out intervention experiments from above to a language setting, and evaluate whether a large language model [70B parameter Chinchilla, from 27] can generalize a causal strategy in this environment from a few-shot prompt. Specifically, we give the LM natural language observations of the objects in a scene, and pose choices between actions described in natural language. For the experimentation trials, we give a choice of which object to intervene on and which feature to transform. In the final trial, we simply give a choice among the objects, each of which is unique along one dimension.
We give the LM four example shots, in all of which the latent "correct" dimension is sampled from only two of the three possible dimensions (color, shape, and texture), and evaluate the ability of the LM to generalize to an episode where the relevant dimension is the third, held-out one. Note that this is a more difficult generalization setting than was investigated with the agents — generalizing a causal relation to a dimension for which it was never observed in training, from only a handful of examples of other dimensions. We also experiment with holding out each of the dimensions, to ensure that the results are robust.
LM performance can be increased by incorporating factors like reasoning traces/chain-of-thought or explanations [55, 71, 69, 77, 40, 35] in the prompt, and these strategies are particularly useful in challenging reasoning cases. We therefore include both reasoning traces and explanations in our prompt examples; at test time we allow the LM to generate its own explanations and reasoning traces at the appropriate places in the episode. We also choose the example shots for the prompt by optimizing for performance on new samples of tasks involving the two trained dimensions [45, 40]. For more details, including an example episode, see Appx. A.2.1.
In Fig. 5 we show the results. Given a strong four-shot prompt, the language model can generalize correctly to new in
stances of the dimensions that were relevant in the prompt; critically, it can also generalize to new instances of the held-out dimension in each case (and even to a somewhat harder test with an added object; Appendix B.9.) However, explanations are critical for this generalization. In Appendix B.7 we show that either explanations or chain-of-thought reasoning suffice to provide some benefit, but results are most reliable when both are provided together. Thus, the language model is able to build on its strong prior knowledge to generalize, but only when supported by some form of explanation.
Figure 5: LMs can generalize the odd-one-out tasks to new dimensions that were not correct in the example shots (right panel), but only if explanations/reasoning are included in the prompt. On new samples of the dimensions included in the examples (left), LMs can achieve some performance with or without explanations. (3 prompts per condition; errorbars are  $95\%$  -CIs.)
What can be learned from passive data about strategies for experimenting to discover and exploit causal structure? Our work illustrates that a non-trivial amount can, at least when the passive data contains the results of experiments. From passive imitation, agents can learn strategies for experimentation that can be generalized to correctly experiment on novel causal structures, and learn
strategies for exploiting the results of their own experiments to achieve goals under those novel structures. Agents can even generalize in complex environments where they receive high-dimensional observations; likewise, passively-trained language models can generalize from a high-quality few-shot prompt. Furthermore, our results highlight the role that natural language explanations can play in drawing out latent causal structure to enable more generalizable learning. Explanations can even enable passive learners to generalize nearly perfectly from otherwise perfectly-confounded data. Passive learning can be surprisingly powerful under the right circumstances.
These observations provide perspective on the kinds of understanding of causality and experiments that can be acquired by language models. For example, consider the following, from Pearl [60]:
"This hierarchy, and the formal restrictions it entails, explains why statistics-based machine learning systems are prevented from reasoning about actions, experiments and explanations. [...] the hierarchy denegrades [sic] the impressive achievements of deep learning to the level of Association [...] Unfortunately, the theoretical barriers that separate the three layers in the hierarchy tell us that the [...] objective function does not matter. As long as our system optimizes some property of the observed data, however noble or sophisticated, while making no reference to the world outside the data, we are back to level-1 of the hierarchy with all the limitations that this level entails."
Naively, it might seem that these arguments doom language models to never acquire real knowledge of causality [cf. 72]. Our results provide useful formal and empirical nuance. While there are certainly restrictions to learning from passive data (see below for further discussion), it is possible to learn generalizable strategies for causal experimentation and intervention from passive data alone — at least if the data include examples of an expert intervening, and perhaps explanations.
Other work has engaged with these issues. In particular, Ortega et al. [56] provide a theoretical account of the effects of confounding on passively trained sequence models, and propose a counterfactual training method to enable an adaptive system to learn effective intervention strategies in the presence of confounders. However, this method relies on having active access to an expert teacher rather than simply offline data; the authors note that "in general, we cannot meta-train an agent only from expert demonstrations to imitate the expert at deployment time" and "[t]o address the problem, it is necessary to incorporate assumptions (e.g. a set of hypotheses) about the confounding variable, but currently it is unclear how to incorporate this seamlessly into a meta-learning scheme." Our work complements theirs, by illustrating some of the cases in which an agent can successfully generalize causal experimentation and goal-seeking strategies from passive data alone — even in the presence of confounders — and by highlighting the role that natural language explanation can play in incorporating expert knowledge to allow active generalization from passive learning.
Explanations & causal structure: Explanations were important for effective learning or generalization in the odd-one-out tasks; explanations can even allow generalizing from otherwise totally-confounded data. These results fit with prior work showing that explanations and reasoning traces help agents [41, 53] and LMs [52, 24, 71, 40, 29]. More broadly, explanations are key for human learning because they explicitly highlight generalizable, causal structures [47, 46, 34]. For example, explanations can constrain the hypothesis space of causal graphs that a system needs to consider [cf. 51]. Alternatively, language<sup>3</sup> may obviate the need for explicit causal graphs. That is, explanations could serve as the “reference to the world outside the data” that Pearl [60] sought in the quote above.
Language models: Our experiments have obvious implications for the causal abilities of language models. Prior works have suggested that LMs can only parrot causal structures observed in training [72]. However, our experiments illustrate that this should not be assumed; the implications of passive learning for causal strategies can be subtle. Passive data is not necessarily strictly observational. In particular, observing other's interventions, especially with explanations, can enable generalization. LM training data will contain many examples describing interventions and outcomes together with explanations thereof — whether scientific articles, or just descriptions of a debugging process. Just as the data in a scientific paper do not become observational if we read the paper rather than writing it, the data on the internet do not become observational just because an LM consumes it passively. Thus, sequences of interventions, outcomes, and explanations may in some cases be sufficient to allow an
$^{3}$ Or other forms of abstraction.
LM to generalize to proposing experiments and using the results to draw appropriate conclusions, at least if given an appropriate prompt.
Tool use & agency: These findings are of particular interest given the recent focus on language models using tools [e.g. 66]. Tool use is a clear case of active interaction in the world, which one might assume requires some knowledge of causality. Our results may help explain the successes of these methods with models trained purely from passive supervised learning (though note that many applications of LMs are also actively fine-tuned, which may also contribute, see below). A related question is whether LMs, despite passive learning, can acquire an ability to act as goal-directed agents. Andreas [2] notes that effective language modeling requires modeling the goals of the communicator that produced the language, and thus suggests that through this process LMs learn to at least model some aspects of agency. Our results may provide some support for these arguments, as generalizing goal-directed behavior to new causal structures is a key aspect of agency.
Experimental limitations: Our experiments are limited in a number of important ways. First, the causal structures we used are small, finite acyclic graphs. While this limitation (i.e., that the causal graph is a DAG over a handful of variables) is assumed by much of the prior causal literature, many real world causal structures are much more complex. For example, they may include cycles, at least over time. $^{4}$  Changing these assumptions changes the nature of the causal discovery problem [6, 19]; thus it is an open question whether agents could similarly learn causal strategies in these more complex settings. More generally, even our most complex environments have relatively simple and constrained observation and action spaces compared to the real world, and the space of possible experiments is generally small and heavily constrained. Our agents are also limited to fit with their environments (e.g. having the right number of observation inputs and action outputs for the number of variables in the DAG environment). It is an open question whether agents could in practice learn causal strategies from passive data in environments that relaxed these simplifying assumptions, though our results suggest it would be possible in principle.
What this paper does not imply: We do not intend to suggest that passive learning is as beneficial as active learning, or that confounding is not a challenge for passive learners. Active learning can be substantially more effective or efficient than passive learning in humans [23, 49], animals [25], and RL agents [58, 5, 57, 8]. It can be difficult to learn from purely positive examples, without some active correction of off-optimal-policy behavior [e.g. 65, 15]. Indeed, while the LM we evaluated was only trained passively, many deployments of LMs incorporate interactive learning through RL from Human Feedback [12; e.g. 59]; one key benefit of this may be the value of training the model on active, on-policy data. We therefore expect that agents or LMs that are fine-tuned in an active setting, or with more sophisticated causal corrections [56], will generally develop more robust and generalizable causal strategies than those trained in purely passive settings.
Furthermore, confounding is clearly a challenge for any system that attempts to discover and use causal structure. The ongoing debates about confounding in the medical literature [e.g. 37, 75] illustrate that this is a challenge even for causal experts such as human scientists. However, our work contributes to understanding the regimes in which confounding during passive learning can be overcome by experimentation at test time. We consider a range of confounds; even our train-test splits effectively correspond to strong confounds which are not directly observable, and for which the agent/LM only experiences training data sampled under one value, and then is tested on situations sampled under another. Nevertheless, agents can in some cases generalize despite these confounds, particularly when supported by explanations. These results highlight the role that language and explicit descriptions of the quest for causal structure can play in acquiring causal strategies.
Conclusions: We illustrate the potential for purely passive learning of active causal strategies across various settings, particularly with the support of explanations. These results have implications for understanding the causal limitations of passive learning for a system that can only interact at test time, and have particular relevance to understanding the capabilities and behaviour of language models.
We thank Nan Rosemary Ke, Tobias Gerstenberg, and the anonymous reviewers for helpful comments and suggestions.
Though these can be accommodated in a DAG by "unrolling" the graph.
[69] Mirac Suzgun, Nathan Scales, Nathanael Scharli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc V Le, Ed H Chi, Denny Zhou, et al. Challenging big-bench tasks and whether chain-of-thought can solve them. arXiv preprint arXiv:2210.09261, 2022.
[70] Jane X Wang, Michael King, Nicolas Pierre Mickael Porcel, Zeb Kurth-Nelson, Tina Zhu, Charlie Deck, Peter Choy, Mary Cassin, Malcolm Reynolds, H Francis Song, et al. Alchemy: A benchmark and analysis toolkit for meta-reinforcement learning agents. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
[71] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 2022.
[72] Moritz Willig, Matej Zečević, Devendra Singh Dhami, and Kristian Kersting. Causal parrots: Large language models may talk causality but are not causal. preprint, 2023.
[73] Kevin Xia, Kai-Zhan Lee, Yoshua Bengio, and Elias Bareinboim. The causal-neural connection: Expressiveness, learnability, and inference. Advances in Neural Information Processing Systems, 34:10823-10836, 2021.
[74] Junzhe Zhang, Daniel Kumor, and Elias Bareinboim. Causal imitation learning with unobserved confounders. Advances in neural information processing systems, 33:12263-12274, 2020.
[75] Jinhui Zhao, Tim Stockwell, Tim Naimi, Sam Churchill, James Clay, and Adam Sherk. Association between daily alcohol intake and risk of all-cause mortality: A systematic review and meta-analyses. JAMA Network Open, 6(3):e236185-e236185, 2023.
[76] Chunting Zhou, Pengfei Liu, Puxin Xu, Srini Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, Lili Yu, Susan Zhang, Gargi Ghosh, Mike Lewis, Luke Zettlemoyer, and Omer Levy. Lima: Less is more for alignment. arXiv preprint arXiv:2305.11206, 2023.
[77] Hattie Zhou, Azade Nova, Hugo Larochelle, Aaron Courville, Behnam Neyshabur, and Hanie Sedghi. Teaching algorithmic reasoning via in-context learning. arXiv preprint arXiv:2211.09066, 2022.