DiscreteTransitionMessagePassingRules

The DiscreteTransition node: a transition between categorical variables through a tensor of probabilities, conditioned on any number of other categorical variables. Use it for the transition and emission factors of hidden Markov models, for a known or a learned transition matrix, and for transitions that switch with a regime or an action.

DiscreteTransitionMessagePassingRules

julia> result = @call_message_update_rule(
           node = DiscreteTransition, target = :in,
           m = (out = Categorical([1.0, 0.0]),), q = (a = PointMass([0.9 0.2; 0.1 0.8]),),
       );

julia> probvec(getresult(result)) ≈ [0.9, 0.2] / 1.1
true

The site has two pages: this one, the node from its definition to its API, and Internals, the tensor algebra its rules are built on.

Overview

A hidden Markov model moves between a finite number of states. The probability of the next state given the current one is a column of a transition matrix A. DiscreteTransition is the factor node of that step: it takes the category of in to the category of out through the column A[:, in]. With conditioning categoricals T1, …, Tn, it takes the fibre A[:, in, T1, …, Tn] of a tensor instead, so the transition can switch with a regime or an action.

A is known, as a point mass, or learned, with a DirichletCollection prior. The node is a tensor node: each of its rules is one contraction of the tensor with its inputs. The same rules therefore serve belief propagation, mean-field and structured variational message passing, for any number of Ts.

Definition

\[p(out \mid in, T_1, \ldots, T_n, A) = A[out, in, T_1, \ldots, T_n], \qquad \textstyle\sum_{i} A[i, j, t_1, \ldots, t_n] = 1\]

The probability that out is category i is the entry of A at i and the categories of the other variables. The constraint makes every fibre A[:, j, t₁, …, tₙ] a probability vector over out. The axes of A are (out, in, T1, …, Tn): out is axis 1, in axis 2 and (:T, k) axis 2 + k. A joint over some of the interfaces covers its members' axes, in order.

Interfaces

using DiscreteTransitionMessagePassingRules, MessagePassingRulesBase, BayesBase, ExponentialFamily
MessagePassingRulesBase.nodespec(DiscreteTransition)

The node has three interfaces and a group T of any length, zero included.

The rules send a Categorical towards out, in and each (:T, k), and a DirichletCollection towards a. As the marginal of a cluster they return a Categorical, a Contingency or a FactorizedCluster.

Algorithm

The node runs under the default algorithm, DefaultAlgorithm, which has no parameters, and a model names none. The node declares no dependencies: each rule takes the inputs of the default scheme.

Supported rules

MessagePassingRulesBase.rule_coverage(DiscreteTransition)

rule_coverage shows one rule per target. Each rule takes whatever inputs the factorisation delivers:

• belief propagation, with every categorical interface in one cluster;
• mean field, q(out) q(in) q(T1) … q(a);
• structured factorisations such as q(out, in) q(a);
• joints of some of the Ts, such as q(out, (T, 1)) q(in) q(a).

The row q(any cluster) is one marginal rule for every cluster. With the average energy, the Bethe free energy is available in every factorisation.

How the rules contract

Every rule is a contraction of E[log A] with its inputs, each along the axes it covers. The marginals are summed out of E[log A], which is then exponentiated, and the messages multiply the result along their axes. With q(a) a point mass and nothing to sum out, the tensor is A itself. So the message towards out is Σⱼ A[:, j] m_in(j) under belief propagation, and exp(Σⱼ E[log A[:, j]] q_in(j)) under mean field.

The message towards a is the expected counts plus one:

\[\alpha[i, j, t_1, \ldots] = 1 + \mathbb{E}_q\big[\mathbb{1}[out = i, in = j, T_1 = t_1, \ldots]\big].\]

The expectation is the outer product of the marginals of the other interfaces, or a joint's tensor. The product of this message with a DirichletCollection(α₀) prior adds the expected counts to α₀.

The marginal of a cluster multiplies the exponentiated tensor by its members' messages. A member observed as a PointMass that is a member of the cluster in its own right is split off, as a FactorizedCluster block. A member inside a whole group T stays in the joint, as a one-hot axis.

Example

A known two-state transition, and the message towards out from a uniform message on in:

A = [0.9 0.2; 0.1 0.8]

@call_message_update_rule(
    node = DiscreteTransition, target = :out,
    m = (in = Categorical([0.5, 0.5]),), q = (a = PointMass(A),),
)

args->discrete_transition_message(args)

@call_message_update_rule runs the rule and draws the inputs it read. The result is the average of the columns of A. getresult extracts the distribution, and the message towards a from mean-field marginals is the expected counts plus one:

julia> using DiscreteTransitionMessagePassingRules, MessagePassingRulesBase, BayesBase, ExponentialFamily

julia> A = [0.9 0.2; 0.1 0.8];

julia> result = @call_message_update_rule(
           node = DiscreteTransition, target = :out,
           m = (in = Categorical([0.5, 0.5]),), q = (a = PointMass(A),),
       );

julia> probvec(getresult(result)) ≈ [0.55, 0.45]
true

julia> result = @call_message_update_rule(
           node = DiscreteTransition, target = :a,
           q = (out = Categorical([0.2, 0.8]), in = Categorical([1.0, 0.0])),
       );

julia> params(getresult(result))[1] ≈ [1.2 1.0; 1.8 1.0]
true

Limitations

• q(a) must be a DirichletCollection or a PointMass of an array.
• E[log A] is clamped away from log 0, so an impossible transition is given a tiny probability rather than zero.
• The contractions are generic, for any number of interfaces: none is specialised to the two-interface case.

API

DiscreteTransition

julia> A = [0.9 0.2; 0.1 0.8];

julia> result = @call_message_update_rule(
           node = DiscreteTransition, target = :out,
           m = (in = Categorical([0.5, 0.5]),), q = (a = PointMass(A),),
       );

julia> probvec(getresult(result)) ≈ [0.55, 0.45]
true

julia> result = @call_message_update_rule(
           node = DiscreteTransition, target = :a,
           q = (out = Categorical([0.2, 0.8]), in = Categorical([1.0, 0.0])),
       );

julia> params(getresult(result))[1] ≈ [1.2 1.0; 1.8 1.0]
true