3. The Memory Model as a World Model

Funding and Authorship

This post was made possible by research funding generously provided by Anki, formerly AnkiHub, and was written in collaboration with Giacomo Randazzo.

---

In this post, we share our updated thoughts on building memory models. We assume you have some familiarity with the problems of spaced repetition memory modeling; if this is new to you, we recommend reading Gwern’s fantastic piece on spaced repetition.

What are we asking the model?

For a given student, card, and point in time, the retrievability of that card is the probability that the student would answer it correctly at that time.

For a given student and card, a forgetting curve is the plot of retrievability over time.

In the traditional “same-card” approach to memory modeling for spaced repetition, we suppose that the retrievability of a card depends only on the review history of that card. Let’s call this assumption the card independence assumption. This approach stands in contrast to building cross-card content-aware models, where we drop the card independence assumption and instead assume that the retrievability of any card is informed by the student’s entire review history across all the cards in their collection.

In the same-card approach, the question that the forgetting curve answers is:

Question 1: What is the probability that the student will answer this card correctly, some specific time after their last review of the same card, conditional on (1) their not reviewing that card in the interim, (2) the assumption that only reviews of the card affect the card’s retrievability, and (3) the student’s entire review history of that card?

In this case, it is reasonable to assume that this curve is a forgetting curve in the classic sense, i.e., it is monotonically decreasing. This is because, given our assumptions, there should be no events between reviews that would cause the memory to become stronger — only forgetting should occur.

Now, let’s consider another question:

Question 2: What is the probability that the student will answer this card correctly, some specific time after a given review of the same card (the ‘reference review’), conditional on the card independence assumption and the student’s entire review history of that card up to the reference review?

In this question, we haven’t specified the absence of intervening reviews, so it’s not reasonable to assume that the curve representing the answer to question 2 is monotonically decreasing; in fact, we expect the student to learn if they study. To answer question 2, we need to draw not only on our understanding of memory dynamics, but also on any prior knowledge we have about how many reviews the student will do, and when, between the reference review and the review for which we wish to predict the retrievability.

This second question is more like student knowledge forecasting than simple memory modeling. The dynamics of memory play a role, but so do our expectations about the student’s study habits.

Another way to see this distinction is that knowledge forecasting implicitly marginalizes over possible future review histories. To predict the student’s knowledge at some future time, we need to consider both what their memory would look like under each possible sequence of intervening reviews and how likely each sequence is to occur. The latter distribution depends on the student’s study behavior and, in the context of spaced repetition, on the scheduling policy. By contrast, next-same-card retrievability prediction (question 1) depends only on our model of memory dynamics.

It’s essential to appreciate that when we drop the card independence assumption and instead try to build cross-card content-aware models, next-same-card retrievability prediction becomes a knowledge forecasting problem, and the closest analogue to question 1 is next-event prediction. We pose this task as question 3:

Question 3: Given some card, what is the probability that the student will answer it correctly, some specific time after the last review of any card, conditional on (1) no other cards being reviewed in the interim, and (2) the student’s entire review history across all cards?

On the other hand, next-same-card prediction in the cross-card content-aware domain would depend on scheduler-mediated expectations about intervening reviews of the other cards. We care about this distinction because we want our models to make predictions by learning the dynamics of human memory, without baking in scheduler-specific base rates that may not generalize.

Of course, there may be (and almost certainly are) unobserved intervening study events that are not captured in the data (e.g. the student taking a practice exam, or doing homework), and unfortunately we expect these to be internalized by the model — the solution here is a more diverse student population. Even so, we want to do the best we can with the data we do have.

An informal point on generalization

An opening question: Why would it be undesirable for the model to internalize scheduler-policy-specific bias?

To address this question, we’re going to take a step back from memory modeling and introduce a thought experiment about statistical modeling in the abstract. If you don’t like this kind of abstraction and would prefer to ground yourself in reality, you can consider $A$ and $B$ to be groups of people, $x$ to be height, and $y$ to be weight.

Consider two sample populations, $A$ and $B$, for whom some covariate $x$ is empirically distributed differently, i.e., $x\sim X_A$ for population $A$ and $x\sim X_B$ for population $B$. Let’s say we’re trying to infer some estimand $y$ given $x$ and some other background knowledge. For the sake of our thought experiment, we have data for population $A$, but not for $B$, though we know about $B$‘s existence and that they probably differ from $A$ with respect to $x$ and $y$.

Now, we have options when training a model $f$ on the data from population $A$ to predict $y$. Let’s consider two:

1. Don’t condition on $x$, in which case the model tries to learn to approximate ${\mathbb{E}}_{x\sim X_A}\left[y\mid x,\text{the other context}\right]$
2. Condition on $x$, in which case the model tries to learn to approximate $y\mid x,\text{the other context}$

In the first case, there is no hope of generalization if $x$ strongly informs $y$ and $X_A$ is very different from $X_B$. The population-specific bias is baked into the model parameters.

In the second case, the model could potentially generalize, even if the empirical supports of $x$ for the two populations are only partially overlapping or even disjoint. Of course, generalization isn’t guaranteed — no model is safe from an arbitrarily hostile data-generating process — but unlike case 1, it’s at least possible, for example if $x$ informs $y$ linearly. Where the empirical supports of $x$ are similar, we may be optimistic about generalization even if the distributions are quite different.

In reality, this is not an all-or-nothing issue: there is more than one observed covariate. The takeaway is to condition on what you can, while trying to avoid incentivizing the model to internalize expectations that likely don’t apply to the population the model will encounter in production.

Relating this back to memory modeling and answering the question at the start of this section: We don’t want the model to learn scheduler-policy-specific bias because we want the model to generalize well under future scheduler policies not represented in the training data. Training on the next-event prediction objective is our attempt to push the model towards learning memory dynamics rather than relying on scheduler-specific base rates that may not generalize to the new schedulers we hope to train. By conditioning on the realized next review event, we avoid having to marginalize over possible intervening recorded review histories. This removes one important way in which scheduler behavior can become entangled with the prediction target.

Unfortunately, the next-event prediction objective isn’t a generalization silver bullet. Conditioning on the next recorded review removes uncertainty about intervening Anki review events, but there may still be unobserved intervening learning events that the model must implicitly marginalize over. Our population is made up of medical students for whom Anki is only one of many study tools. For example, the extent to which students engage with the material being tested by their cards outside of Anki is a population-specific factor we currently can’t condition on. “These are medical students in medical school” will therefore still be a bias in the model weights.

The autoregressive memory model and how it relates to scheduling

We’ve just described why we’re building a next-event prediction model — to capture memory dynamics while minimizing scheduler bias — but we haven’t yet said how we plan to build a useful scheduler.

In short, we plan to use the next-event prediction model as a world model to train scheduling policies (e.g. with reinforcement learning).

Unlike memory modeling, which is a regression problem, scheduling is an optimal control problem. The scheduler asks: given the dynamical system that is the student’s memory, how should I prompt the student over time to optimize their learning relative to some utility function? So scheduling is a question of policy/control optimization. Since we don’t want to train scheduling policies by experimenting on real students, we need a model of the system that the scheduler can interact with during training. Additionally, scheduling is affected by UX concerns in a way that memory modeling is not.

We can’t just train a scheduler as a regression model on the historical data due to the generalization issues already described. The data reflects trajectories under specific existing policies, so learned forecasts will be biased by past scheduler-specific expectations. If the new scheduler diverges from historical policy, its learned expectations will be wrong. A scheduler needs to learn control, and therefore needs to be trained in an environment where it can see how its decisions change its expectations.

The next-event prediction model can be used as such an environment. Given a review history and a proposed next review event, the model gives us a distribution over the outcome of that event. We can sample an outcome, append the realized event to the history, and ask the model again — this is the sense in which the model is autoregressive, and it’s what lets us roll it forward into a whole synthetic study trajectory, one review at a time. During training, the scheduler chooses what is reviewed and when, the model supplies the outcomes, and the policy learns against the trajectories that result. This lets us train policies against an explicit model of memory dynamics rather than directly imitating trajectories produced by historical schedulers.

One important caveat is that conditioning on the next review event removes scheduler dependence from the prediction target, but not from the training distribution. Historical schedulers still determine which review events the model sees, so a new scheduler may take actions outside the model’s training support. A policy trained against the world model may also learn to exploit systematic errors in the model. As we gather data from more diverse scheduling strategies, we can improve the world model and retrain policies against it without changing the overall procedure.

This separation of concerns into world model and policy model has another benefit: the optimal scheduling policy is relative to a user’s goals, and we may want to train different policies for different situations — for example, someone studying for the long term versus someone studying for a specific deadline.

Finally, as we train new policies and deploy them, we can collect more diverse data, which in turn lets us evaluate and improve the memory model. The improvement cycle is: deploy scheduler → new data → better memory model → better policies.