What Has a Foundation Model Found? Using Inductive Bias to Probe for World Models

Thank you for your careful and insightful review.

I will increase my score if the Framework section is sufficiently improved, and if the language is toned down a bit re: the results showing that the foundation model has learned a world model (rather than, to what degree its predictions are consistent with a given world model).

We really like the way you've described the framework and we've modified our paper to match (see below).

We also agree with your language change: it's overstating to say “has learned a world model”; given our test, it's much more accurate to say “consistent with a given world model”, which is the language we'll use going forward.

Your comments were very constructive and have improved the paper. We believe we've addressed your comments below; if we haven't, please let us know if you have any further questions.

Question 1: ...My first question is if this accurately captures what they are proposing...

We've rewritten the framework to incorporate your suggestions. We don't have space to include the full section here but the outline is:

• This paper is about measuring the degree to which a foundation model's predictions are consistent with a given world model.
• Note a mechanistic interpretation is difficult to measure.
• Instead we examine model predictions on synthetic tasks that obey the world model
• Consistency requires two properties:
  1. R-IB: When two inputs map to the same state in the world model, do they have the same predictions for each task?
  2. D-IB: When two inputs have the same predictions for each task, do they belong to the same state?
• R-IB and D-IB are different perspectives on consistency, analogous to type-I/type-II errors.

Question 2: It could be that the foundation model has learned the correct world model but its predictions cannot be described by prediction = f(state = WM(input))

Great question. This is why we have two metrics:

• If a model predicts f(state, input) != f(state, input') for same-state inputs, R-IB is low
• If a model predicts f(state, input) = f(state', input') for different-state inputs, D-IB is low

These metrics can be used independently (e.g. focusing only on R-IB). But there are many applications where we’d care if a model has extraneous information. For example, if a model carries extraneous information its predictions on small data may be subject to spurious correlations. Empirically we find that low D-IB scores suggest that models base their predictions not on true state but on next-token partitions (see Figure 2 and Table 8), hurting small-data performance.

Question 3: When you say about D-IB “If a learning algorithm’s effective foundation is based on the world model, it should now extrapolate in predictable ways between inputs values in different state”, what precisely do you mean?

You've helped us identify an unfortunate typo. The sentence should read "...it should not extrapolate in predictable ways...". So if the model learns according to not only the effective foundation but other inputs, D-IB is lower. We've corrected this.

Question 4: Physical laws involve two kinds of object; kinematics (states) and dynamics...

We want to make sure we're interpreting your question correctly:

• kinematics: mapping from inputs X to state space (i.e. given by phi)
• dynamics: how the state spaces evolve.

We agree the distinction is interesting. Our framework only addresses kinematics (e.g. state, not the laws that may govern dynamics on top of state). Therefore R-IB and D-IB tell us about kinematics. But the two are related; knowing kinematics often allows recovering dynamics from enough sequences.

We focus on states because even in physics there's value in studying kinematics independently from dynamics. Although orbital trajectories (dynamics) are fixed, many interesting functions depend on state (e.g. energy, angular momentum). Addressing dynamics explicitly is valuable future work.

[Does the result that] models are learning task-specific heuristics [come only] from these symbolic regressions?

Symbolic regressions are one tool for showing that models learn task-specific heuristics (we've added more experimental details in the appendix), and it's a common method for validating physics models [1, 2]. But it's not the only tool; our appendix has additional experiments with other methods (Table 8) in non-physics domains. We've moved this to the main text.

We've also added added new physics experiments to Table 8:

• D-IB for distinct states with same next tokens: 0.709
• D-IB for distinct states with different next tokens: 0.764

So poor D-IB can be partially attributed to models basing their predictions not on state but on next-token partitions. Figure 2 illustrates this point for Othello.

[1] Udrescu et al. "AI Feynman: a Physics-Inspired Method for Symbolic Regression." (2020).
[2] Cranmer et al. "Discovering Symbolic Models from Deep Learning with Inductive Biases." (2020).

Thank you for your positive review of our paper. We're glad that you appreciated our paper and findings, and that our results gave you "a deeper understanding [of] why foundation models fail to generalize".

The proposed metrics (R-IB and D-IB), while needed for a comprehensive evaluation, are somewhat abstract and sometimes difficult to understand in the context of the various domains (Orbital physics, Othello). It would be helpful to ground the concepts for these metrics with domain-specific examples when describing them in section 2.

This is a great point. We've revised our paper to make the framework clearer, and we've added more domain-specific examples.

Specifically:

• A foundation model's predictions are consistent if two properties hold:
  1. When two inputs map to the same state in the world model, do they have the same predictions for each synthetic task? This is measured by R-IB.
  2. When two inputs have the same predictions for each synthetic task, do they belong to the same state? This is measured by D-IB.
• So R-IB and D-IB are different perspectives on measuring consistency, analogous to type-I and type-II errors in classification.

Here's an implementation example for the lattice problem:

• X is a sequence of directions ("L", "R", "stay")
• 5 states (1-5)
• phi(X) maps movement in X starting at state 1
  • So phi("R", "R", "L")=2 (two states right, one state left)
• We create random synthetic datasets (X, Y) such that phi(x) = phi(x') implies y = y' for each (x,y), (x',y') pair.
• We fit a foundation model (e.g. by fine-tuning) on a very small subset of each dataset and make predictions for the held-out data points.
• R-IB: do two held-out points with the same state always have the same predictions across synthetic datasets?
• D-IB: do two held-out points with different states have unpredictable extrapolations across synthetic datasets?

Thank you for your positive review of our paper. We respond to your comments and describe new results below; to summarize the new results, we've added:

• New experiments showing robustness to predictor in R-IB and D-IB calculations
• A more intuitive explanation of our framework
• Clarification and new analysis of multiple representations in Othello

We hope our comments have addressed your concerns. If not, please let us know if you have any more questions we can address in the follow-up.

Can you provide a more intuitive explanation of the inductive bias definitions? Can you provide a toy example to go along with the Section 2: Implementation?

We've revised our paper to make this more intuitive. Broadly:

• A foundation model's predictions are consistent if two properties hold:
  1. When two inputs map to the same state in the world model, do they have the same predictions for each synthetic task? This is measured by R-IB.
  2. When two inputs have the same predictions for each synthetic task, do they belong to the same state? This is measured by D-IB.
• So R-IB and D-IB are different perspectives on measuring consistency, analogous to type-I and type-II errors in classification.

Here's an implementation example for the lattice problem:

• X is a sequence of directions ("L", "R", "stay")
• 5 states (1-5)
• phi(X) maps movement in X starting at state 1
  • So phi("R", "R", "L")=2 (two states right, one state left)
• We create random synthetic datasets (X, Y) such that phi(x) = phi(x') implies y = y' for each (x,y), (x',y') pair.
• We fit a foundation model (e.g. by fine-tuning) on a very small subset of each dataset and make predictions for the held-out data points.
• R-IB: do two held-out points with the same state always have the same predictions across synthetic datasets?
• D-IB: do two held-out points with different states have unpredictable extrapolations across synthetic datasets?

My impression is that if the approach fails, it could be that the predictor model is underpowered, rather than due to a lack of world model

We can think of two ways to interpret this question: one is about the predictor that's part of the foundation model; the other is about the predictor used to calculate R-IB and D-IB.

For the first interpretation: We're interested in understanding a foundation model's extrapolative properties in small data regimes (e.g. few shot learning) since this is how they're typically used. Thus, our metrics measure consistency with the world model, not predictor accuracy, so they wouldn't penalize an underpowered predictor.

For the second interpretation: R-IB and D-IB are based on a predictor that only uses a single binary feature. Therefore the Bayes optimal solution is nonparametrically identifiable. This table varies the number of data points used to form R-IB and D-IB; empirically we find little sensitivity, because the predictor converges to the optimal solution quickly.

I wonder how useful this method can be when there [are] multiple acceptable state representations. It would be helpful for the paper to analyze... what happens in Othello under the two representations described above.

This is an important question. Our framework depends on the state mapping, not how it's represented. The only important quantities for D-IB and R-IB are whether two inputs have the same or different state. This means the metrics aren't sensitive to different ways state can be represented, so the two Othello representations have identical metrics. In comparison, sensitivity to representation makes other metrics (e.g. probes) more fragile.

However we could still use different representations to help analyze poor IB metrics. D-IB measures how similar the predictions for input pairs with distinct states are. We consider two distance metrics for distinct-state pairs: standard ("cell is black") vs. relative ("cell has my color") Hamming distances. We find:

• Corr(D-IB, standard distance) = 0.094
• Corr(D-IB, relative distance) = 0.001

Interestingly this indicates states with similar standard representations are more easily confused.

I believe... this paper assumes the generalization is still in-distribution.

This is a good point. It's crucial to also test out-of-distribution (OOD). While our submission was unclear, we do test on OOD sequences -- the only restriction on synthetic tasks is that they be consistent with the world model, not the original sampling distribution. We've changed the writing to make this clear.

It is highly unlikely that the model will see the same state twice for two different trajectories [in Othello].

This is absolutely right. Going back to your earlier point, it shows the importance of testing distribution shifts. In practice, we sample inputs in Othello so that we have enough same-state examples (3,889) to make accurate estimates. Appendix B has more information but we'll highlight it in the main text.