Transformers Learn Bayesian Networks Autoregressively In-Context

The main goal of this paper is to demonstrate how transformers can learn Bayesian networks in context. The paper provides a theoretical proof showing that transformers are capable of modeling Bayesian networks in context. Additionally, The paper provides an evidence that transformers trained on such tasks can make Bayes-optimal predictions for new samples in an autoregressive manner.

The paper studies the problem of in-context learning in transformers. In particular, it focuses on whether transformers are able to learn Bayesian networks in-context. In this setting, the model, given N different realisations of a specific graph and a query sample, is tasked to predict the probability distribution associated with a missing variable (see construction in Eq. 3.2). The assumption is that, if the model is able to infer the conditional probabilities associated with the Bayesian network, it can then use them to predict the value of the missing variable. In addition, once the model has captured such conditional probabilities, it is in principle able to generate new samples from the inferred graphical model (Algorithm 1).

The authors first provide a theoretical construction for a two-layer transformer which is capable of estimating the conditional probabilities of the Bayesian network according to the context, and autoregressively generating a new sample according to the Bayesian network with estimated conditional probabilities (Theorem 4.1 and Lemma 6.1 and Lemma 6.2).

The authors also conduct an empirical analysis to show the performance of trained transformers (with up to 6 layers) on the task of in-context learning three graph structures, namely a "general graph", a "tree" and a "chain". The performance of the model is studied by varying the number of in-context examples seen at training time and evaluating the model on different number of test in-context samples. The results show some evidence that transformers are capable of learning Bayesian networks in-context.

This paper theoretically constructs a simple transformer model that learns to sample from a Bayesian Network (BN) from in-context samples. A BN is an ordered set of variables that has a causal ordering among their variables and satisfy various conditional probabilities. The paper shows that for a BN of bounded maximum indegree, a simple 2-layer transformer can approximate the conditional probabilities and generate new samples. The proof is simple and basically constructs positional embeddings that mimic the parent structure of the BN and then applies MLE. Experiments are conducted to validate the theory on simulated BNs and also probe the number of layers needed. The target audience are people interested in theoretical machine learning.

This paper considers the ability of Transformers to estimate in-context the conditional probabilities of categorical variables. Theoretically, the paper seeks to prove that for any joint distribution over categorical variables and an ordering over them, there exists a two-layer Transformer that can represent the true conditional probabilities of each variable given those that are earlier in the ordering. Empirically, the paper considers experiments on synthetic data where Transformers are trained on samples from different Bayesian networks that all come from some family of graphs. The paper compares the probabilities estimated in-context to the ground truth as well as those estimated via naive Bayes and Bayesian inference, finding trends that suggest that Transformers have the capacity to estimate conditional probabilities in-context.

We thank the ACs for handling our paper and thank all the reviewers for the constructive and helpful comments. We realize that many of the reviewers' concerns are caused by unclear terminologies and insufficient background explanations, and that our work can benefit from a thorough revision. Therefore, we decide to withdraw our submission. We will carefully revise the paper based on the valuable feedback we receive, and submit the work to a future venue.

Dear Reviewers,

Thank you very much for your constructive and helpful comments. We realize that many of your concerns are caused by unclear terminologies and insufficient background explanations, and that our work can benefit from a thorough revision. Therefore, we decide to withdraw our submission. We will carefully revise the paper based on your valuable feedback before submitting it to future venues. We would like to respond to your major concerns as follows.

Q1. Our setting requires a priori knowledge of the causal order of variables.

A1. Our paper aims to demonstrate that transformers can autoregressively generate new samples according to an estimated Bayesian network. Having a pre-determined order of the variables is natural and is consistent with the practice, since we consider autoregressive generation.

Please also note that without a pre-determined order of variables, Bayesian networks are not identifiable because multiple Bayesian networks can equivalently describe the same joint distribution of variables. Specifically, it can be shown that for any given order of variables, there exists a Bayesian network that describes the same joint distribution, satisfying that each variable can only be a descendant of the preceding variables according to that order.

Q2. This is not a paper about learning Bayesian networks in-context; rather, it is a paper about whether Transformers can estimate conditional probabilities of discrete variables in-context.

A2. Thanks for pointing this out. We will clarify in our revision that our goal is to show that there exists a simple transformer model that can (i) estimate the conditional probabilities of the Bayesian network according to the context, and (ii) autoregressively generate a new sample according to the Bayesian network with estimated conditional probabilities. We have made clarifications in the abstract, and we will consider changing the title of the paper to further avoid confusion.

Please note that similar settings of estimating the conditional probabilities of the Bayesian network according to the context is standard and has been considered in recent work [1].

Q3. What is the precise objective by which the Transformer is trained?

A3. Our theory focuses on studying the expressive power of transformers. The nature of this kind of research is that it does not necessarily focus on any particular training objective. Instead, we just aim to show that there exists such a transformer model that can accomplish the desired task. Similar studies of the expressive power of transformers are also considered in previous works such as [2].

In our experiments, we train the transformers to minimize the cross-entropy loss of predicting the masked-out variables in queries. We will add clarifications to this together with a more detailed explanation of the curriculum.

Q4. Explanations of the 'naive Bayes' and 'Bayesian inference' methods we used for performance comparison with transformers in our experiments.

A4. We have realized that these are confusing terminologies, and we will replace them with clearer names in the revision.

Suppose that there are $M$ variables $X_1,\ldots,X_M$, and that we have $N$ independent groups of observations $(X\_{11},\ldots,X\_{M1}),\ldots,(X\_{1N},\ldots,X\_{MN})$. Further suppose that we have query observations $X\_{1q},\ldots,X\_{(m_0-1)q}$, and our goal is to estimate the conditional probabilities of the form:

$P( X\_{m_0} = j | X\_{1} = X\_{1q},\ldots, X\_{m_0-1} = X\_{(m_0-1)q} )\quad\quad (Eq1)$.

Denote by $\mathcal{P}(m_0)$ the parent set of the $m_0$-th variable according to the Bayesian network. In our current manuscript, we mean by 'naive Bayes' the method that estimates the conditional probability above in (Eq1) as

$\frac{ | \{ i\in [N]: X\_{m_0i} = j, \text{ and } X\_{mi} = X\_{mq} \text{ for all } m = 1,\ldots,m_0-1 \} | }{ | \{ i\in [N]: X\_{mi} = X\_{mq} \text{ for all } m = 1,\ldots,m_0-1 \} | }$.

We mean by ‘Bayesian inference' the method that estimates the conditional probability above in (Eq1) as

$\frac{ | \{ i\in [N]: X\_{m_0i} = j, \text{ and } X\_{mi} = X\_{mq} \text{ for all } m\in \mathcal{P}(m_0) \} | }{ | \{ i\in [N]: X\_{mi} = X\_{mq} \text{ for all } m\in \mathcal{P}(m_0) \} | }$.

Clearly, 'Bayesian inference' can utilize more observations to calculate frequencies and is therefore more efficient. According to our theory, the performance of transformers should be comparable to 'Bayesian inference' and better than 'naive Bayes', particularly when the ground truth Bayesian network is complex.

We hope that our response above can address most of your concerns. Thank you again for your valuable feedback.

Best regards,

Authors

[1] Eshaan Nichani, et al. "How Transformers Learn Causal Structure with Gradient Descent." ICML 2024.

[2] Yu Bai, et al. "Transformers as statisticians: Provable in-context learning with in-context algorithm selection." NeurIPS 2023.