Reasoning About Regenerate

Gen provides a primitive called regenerate that allows users to ask for certain random choices in a trace to be re-generated from scratch. regenerate is the basis of one variant of the metropolis_hastings operator in Gen’s inference library.

This notebook aims to help you understand the computation that regenerate is performing.

using Gen: bernoulli, @gen

Let’s start by defining a simple generative function:

@gen function foo(prob_a)
    val = true
    if ({:a} ~ bernoulli(prob_a))
        val = ({:b} ~ bernoulli(0.6)) && val
    end
    prob_c = val ? 0.9 : 0.2
    val = ({:c} ~ bernoulli(prob_c)) && val
    return val
end;

Recall the distribution on choice maps for this generative function:

Let’s first obtain an initial trace with ${a \mapsto \mbox{true}, b \mapsto \mbox{false}, c \mapsto \mbox{true}}$, using generate:

using Gen: generate, choicemap

trace, weight = generate(foo, (0.3,), choicemap((:a, true), (:b, false), (:c, true)));

Now, we ask for the value at address :a to be re-generated:

using Gen: regenerate, select, NoChange
(trace, weight, retdiff) = regenerate(trace, (0.3,), (NoChange(),), select(:a));

Note that unlike update, we do not provide the new values for the random choices that we want to change. Instead, we simply pass in a selection indicating the addresses that we want to propose new values for.

Note that select(:a) is equivalent to:

selection = DynamicAddressSet()
push!(selection, :a)

We print the choices in the new trace:

using Gen: get_choices

get_choices(trace)

│
├── :a : false
│
└── :c : true

Re-run the regenerate command until you get a trace where a is false. Note that the address b doesn’t appear in the resulting trace. Then, run the command again until you get a trace where a is true. Note that now there is a value for b. This value of b was sampled along with the new value for a—regenerate will regenerate new values for the selected adddresses, but also any new addresses that may be introduced as a consequence of stochastic control flow.

What distribution is regenerate sampling the selected values from? It turns out that regenerate is using the internal proposal distribution family $q(t; x, u)$, just like like generate. Recall that for @gen functions, the internal proposal distribution is based on ancestral sampling. But whereas generate was given the expicit choice map of constraints ($u$) as an argument, regenerate constructs $u$ by starting with the previous trace $t$ and then removing any selected addresses. In other words, regenerate is like generate, but where the constraints are the choices made in the previous trace less the selected choices.

We can make this concrete. Let us start with a deterministic trace again:

trace, weight = generate(foo, (0.3,), choicemap((:a, true), (:b, false), (:c, true)));

Understanding how regenerate constructs the internal proposal distribution family

We will run regenerate with a selection of just :a. Let’s analyze the internal proposal distribution in this case:

    (trace, weight, retdiff) = regenerate(trace, (0.3,), noargdiff, select(:a));

Since the current trace is $t = {a \mapsto \mbox{true}, b \mapsto \mbox{false}, c \mapsto \mbox{true}}$, the constraints $u$ that will be passed to the internal proposal are $u = {b \mapsto \mbox{false}, c \mapsto \mbox{true}}$ (everything in $t$ but with the mapping for any selected addresses removed).

To compute the internal proposal distribution, we first write down the list of all choice maps $t’$ where $p(t’; x) > 0$, with the probability $p(t’; x)$ listed.

Then, we eliminate any choice maps $t’$ such that $t’(i) \ne u(i)$ for some address $i$ that is contained in both maps $u$ and $t’$.

In particular:

• we eliminate $t’ = {a \mapsto \mbox{true}, b \mapsto \mbox{true}, c \mapsto \mbox{true}}$ because $t’(b) = \mbox{true} \ne u(b) = \mbox{false}$.

• we eliminate $t’ = {a \mapsto \mbox{true}, b \mapsto \mbox{true}, c \mapsto \mbox{false}}$ because $t’(c) = \mbox{false} \ne u(c) = \mbox{true}$.

• we eliminate $t’ = {a \mapsto \mbox{true}, b \mapsto \mbox{false}, c \mapsto \mbox{false}}$ because $t’(c) = \mbox{false} \ne u(c) = \mbox{true}$.

• we eliminate $t’ = {a \mapsto \mbox{false}, c \mapsto \mbox{false}}$ because $t’(c) = \mbox{false} \ne u(c) = \mbox{true}$.

For the remaining choice maps $t’$ we require that $q(t’; x, u) > 0$. The remaining two choice maps are:

The ancestral sampling algorithm has two possible results, depending on whether the new value for a was true or false. It will sample a = true with probability 0.3 and a = false with probability 0.7. If a is sampled to be true, then the existing value of b = false will always be kept. In both cases, the previous value of c = true is kept. Therefore, the internal proposal distribution is:

Reversibility of regenerate

Regenerate has the useful property, that for selected addresses $I$ (e.g. $I = {a}$) and initial choice map $t$ where $p(t; x) > 0$, if regenerate has nonzero probability of producing a new choice map $t’$ from $t$ and $I$, then regenerate also has a nonzero probability of producing choice map $t$ from $t’$ and the same set of selected addresses $I$.

Challenge: convince yourself, or prove, that this is the case.

Understanding the weight returned by regenerate

The weight returned by regenerate, for selected addresses $I$ is:

where $u$ is the restriction of $t$ to the complement of $I$, and where $u’$ is the restriction of $t’$ to the complement of $I$.

We will now manually compute what the weights should be for the two possible transitions from $t = {a \mapsto \mbox{true}, b \mapsto \mbox{false}, c \mapsto \mbox{true}}$, with $I = {a}$.

First, consider $t’ = {a \mapsto \mbox{true}, b \mapsto \mbox{false}, c \mapsto \mbox{true}}$. In this case, $u’ = {b \mapsto \mbox{false}, c \mapsto \mbox{true}}$. The internal proposal distribution in this case is:

Therefore, the weight for the transition from $t = {a \mapsto \mbox{true}, b \mapsto \mbox{false}, c \mapsto \mbox{true}}$ to $t’ = {a \mapsto \mbox{true}, b \mapsto \mbox{false}, c \mapsto \mbox{true}}$ is:

Next, we consider $t’ = {a \mapsto \mbox{false}, c \mapsto \mbox{true}}$. In this case $u’ = {c \mapsto \mbox{true}}$. The internal proposal distribution in this case is:

The weight for the transition from $t = {a \mapsto \mbox{true}, b \mapsto \mbox{false}, c \mapsto \mbox{true}}$ to $t’ = {a \mapsto \mbox{false}, c \mapsto \mbox{true}}$ is:

log((0.7 * 0.9)/(0.3 * 0.4 * 0.2)) + log((0.3 * 0.4)/(0.7))

1.504077396776274

Now that we’ve done all this work, let’s check it against Gen.

Run the cell enough times to sample both of the transitions, and confirm that the weights match with our calculations:

trace, weight = generate(foo, (0.3,), choicemap((:a, true), (:b, false), (:c, true)));
(trace, weight, retdiff) = regenerate(trace, (0.3,), (NoChange(),), select(:a));
display(get_choices(trace))
println("weight: $weight");

│
├── :a : false
│
└── :c : true



weight: 1.504077396776274

Aren’t we glad this is automated by Gen!

Approaching irreducibility

Exercise: Draw a graph in which each random choice map $t$ where $p(t; x) > 0$ is a node, and where there are directed edges from $t$ to $t’$ if applying regenerate to $t$ with selection ${a}$ can produce trace $t’$. Do the same for selections ${b}$ and ${c}$. What about selection ${a,b,c}$?