A blog post talking about using differentiable optimization to learn finite state machines
https://google-research.github.io/self-organising-systems/2022/diff-fsm/
Differentiable Finite State Machines
2022-06-07
Alexander Mordvintsev
<div><div>
<p><a href="https://colab.research.google.com/github/google-research/self-organising-systems/blob/master/notebooks/diff_fsm_jax.ipynb#scrollTo=KqJ971iNjenq">Reproduce in a Notebook</a>
<a href="https://github.com/google-research/self-organising-systems/discussions/6"> Discuss on </a>
</p>
<p>In this blog post I'd like to show how differentiable optimization can be used to learn Finite State Machines (FSM) for solving toy string processing tasks. I'll show how simple regularization and initialization techniques can steer continuous optimization towards finding discrete deterministic solutions and increase the training success rate. I think that experiments shown here may have some educational value, e.g. in demonstrating less conventional (and perhaps unexpected) uses of differentiable programming and some elegant JAX tricks.</p>
<p>There exists a vast literature on various approaches to FSM synthesis from data, so it's highly likely that the particular method described here is already well known. A quick review of related literature will be given in the <a href="#related-work">Related Work</a> section.</p>
<p>Here is an example of a toy problem we are going to solve. Consider a pair of input-output strings: </p>
<figure><pre> In: 01010100100111111
Out: 01000100000101010
</pre></figure>
<p>We'd like to design an algorithm that would produce the expected output string if we feed it with the input string character by character. There are many different possibilities to perform this task. The brute force way would be just to store the expected output string and emit it character by character without paying attention to the input at all. This is actually a pretty complex solution, for example implementing such an algorithm as a state machine would require the number of states that is equal to the string length. Intuitively, ideas behind <a href="https://en.wikipedia.org/wiki/Occam%27s_razor">Occam's razor</a> and <a href="https://en.wikipedia.org/wiki/Kolmogorov_complexity">Kolmogorov's complexity</a> suggest that simpler solutions are preferable. The simplicity gives us hope that the solution would be interpretable and generalize to input sequences beyond the training set.</p>
<p>The simplest algorithm for the string pair above scans the input and replaces every second '1' character it encounters with '0'. A type of FSM called <a href="https://en.wikipedia.org/wiki/Mealy_machine">Mealy Machine</a> is natural way of representing such algorithms:</p>
<figure>
<div>
<table>
 <thead>
 <tr>
 <th> </th>
 <th> </th>
 <th>T: state'</th>
 <th>R: output</th>
 <th>s0</th>
 </tr>
 <tr>
 <th> </th>
 <th> </th>
 <th>A</th>
 <th>B</th>
 <th>0</th>
 <th>1</th>
 <th></th>
 </tr>
 <tr>
 <th>input</th>
 <th>state</th>
 <th> </th>
 <th> </th>
 <th> </th>
 <th> </th>
 <th> </th>
 </tr>
 </thead>
 <tbody>
 <tr>
 <th>0</th>
 <th>A</th>
 <td>1.00</td>
 <td>0.00</td>
 <td>1.00</td>
 <td>0.00</td>
 <td>1.00</td>
 </tr>
 <tr>
 <th>B</th>
 <td>0.00</td>
 <td>1.00</td>
 <td>1.00</td>
 <td>0.00</td>
 <td>0.00</td>
 </tr>
 <tr>
 <th>1</th>
 <th>A</th>
 <td>0.00</td>
 <td>1.00</td>
 <td>0.00</td>
 <td>1.00</td>
 <td>1.00</td>
 </tr>
 <tr>
 <th>B</th>
 <td>1.00</td>
 <td>0.00</td>
 <td>1.00</td>
 <td>0.00</td>
 <td>0.00</td>
 </tr>
 </tbody>
</table>
</div>
<figcaption>Mealy machine and its transition table</figcaption>
</figure>
<p>This machine starts at the state "A". The label at each edge has a form "&lt;input&gt;/&lt;output&gt;": &lt;input&gt; character activates the transition along the given edge, and &lt;output&gt; character is emitted to the output string during the transitions. Transition tables are a common way to represent FSMs. Each row of the transition table above corresponds to a different combination of <strong>input</strong> and <strong>state</strong>, and represents the new <strong>state'</strong> and the <strong>output</strong> value, encoded as two one-hot vectors. This table also includes the column <strong>s0</strong>, a one-hot vector, encoding the starting FSM state (which is replicated two times for completeness). Formally, this table can be represented by three tensors with following axes:</p>
<p>$$ \begin{aligned}
T&amp;:[\mathsf{input}, \mathsf{state}, \mathsf{state'}] \\
R&amp;:[\mathsf{input}, \mathsf{state}, \mathsf{output}] \\
s^0&amp;:[\mathsf{state}]
\end{aligned} $$</p>
<p>Now we can express the FSM output and state transition computation as the following sums</p>
<p>$$ \begin{aligned}
y^t_{\mathsf{output}} &amp;= \sum_{\mathsf{input},\mathsf{state}} x^t_\mathsf{input} s^t_\mathsf{state} R_{\mathsf{input}, \mathsf{state}, \mathsf{output}} \\
s^{t+1}_{\mathsf{state'}} &amp;= \sum_{\mathsf{input},\mathsf{state}} x^t_\mathsf{input} s^t_\mathsf{state} T_{\mathsf{input}, \mathsf{state}, \mathsf{state'}}
\end{aligned}$$</p>
<p>where x, y, and s are sequences of one-hot vectors, representing inputs, outputs and FSM states. Luckily, modern differentiable programming frameworks (I'll be using JAX in this tutorial) provide the powerful <code>'</code><a href="https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.einsum.html"><code>einsum</code></a><code>'</code> function to compute these kinds of expressions. It's easy to combine it with the <code>'</code><a href="https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html"><code>scan</code></a><code>'</code> primitive to compute FSM outputs and intermediate states for a given input:</p>
FSM = namedtuple('FSM', 'T R s0')
def run_fsm(fsm: FSM, inputs):
 def f(s, x):
 y = jp.einsum('x,s,xsy-&gt;y', x, s, fsm.R)
 s1 = jp.einsum('x,s,xst-&gt;t', x, s, fsm.T)
 return s1, (y, s1)
 _, (outputs, states) = jax.lax.scan(f, fsm.s0, inputs)
 return outputs, jp.vstack([fsm.s0, states])
<p>For now we were presenting states and characters as one-hot vectors, but we can also think of them as <i>probability distributions</i>. This relaxes the discrete set of one-hot vectors to a continuous space of vectors of non-negative values that sum up to one. The resulting model is called a <a href="https://en.wikipedia.org/wiki/Probabilistic_automaton">probabilistic automaton</a> and the above math still holds.</p>
<p>Given that we represented a state machine in terms of differentiable operations, why don't we try to use gradient descent to find a machine given the input-output pair? We need to make sure that we use the right parameterisation that enforces rows of $T$, $R$ and $s^0$ to be probability distributions. Usually this is achieved by using the 'softmax' function. $(T, R, s^0)$ are parametrized with real-valued tensors and softmax is applied along the last dimension to turn them into probabilities. Our toy examples are going to use the alphabet that consists of just two characters: '0' and '1'. We assume that we don't know the required number of FSM states in advance, so in the spirit of deep learning we are going to use an overparameterized FSM representation that has a maximum of 8 states.</p>
<p>CHAR_N = 2</p>
<p>STATE_N = 8</p>
<p>Params = namedtuple('Params', 'T R s0')</p>
<p>def init_fsm(key, noise=1e-3) -&gt; Params:</p>
<p> k1, k2, k3 = jax.random.split(key, 3)</p>
<p> T = jax.random.normal(k1, [CHAR_N, STATE_N, STATE_N]) * noise</p>
<p> R = jax.random.normal(k2, [CHAR_N, STATE_N, CHAR_N]) * noise</p>
<p> s0 = jax.random.normal(k3, [STATE_N]) * noise</p>
<p> return Params(T, R, s0)</p>
<p>def hardmax(x):</p>
<p> return nn.one_hot(x.argmax(-1), x.shape[-1])</p>
<p>def decode_fsm(params: Params, hard=False) -&gt; FSM:</p>
<p> T, R, s0 = params</p>
<p> f = hardmax if hard else nn.softmax</p>
<p> return FSM(f(T), f(R), f(s0))</p>
<p><code>'init_fsm'</code> function creates a randomly initialized FSM. Parameters are set to very small values, so all distributions are close to uniform. 'decode_fsm' serves to convert parameters into FSM matrices to consume by <code>'run_fsm'</code>. It has two operation modes: 'soft', when softmax is used, and 'hard', when all distributions get collapsed into one-hot vectors, and stochastic FSM becomes deterministic.</p>
<p>Let's try to simply minimize the sum of squared-differences between the FSM output and the reference output sequence:</p>
<p>fsm = decode_fsm(params)</p>
<p>y, s = run_fsm(fsm, x)</p>
<p>error = jp.square(y-y0).sum()</p>
<p>Throwing ADAM optimizer into minimizing the error indeed makes it go down. Getting there took a bit of fiddling with the optimizer parameters. Motivated by the lack of loss stochasticity and small number of training steps (400) I use a pretty high learning rate (0.25) and set <code>beta1=beta2=0.5</code>.</p>
<figure>
</figure>
<p>Here is the diagram showing input, output and state probability distributions produced by the resulting FSM at different time steps (circle size is proportional to the probability):</p>
<figure>
<figcaption></figcaption>
</figure>
<p>We see that although the model manages to perfectly reproduce the desired output for the given input string, it's using many more states than necessary, and at some timesteps (e.g. 0 or 3) the probability is distributed among a number of states, so the process doesn't look deterministic. In machine learning, the usual strategy to steer the optimization into a more desirable solutions region is to introduce additional regularizer objectives. In our case "desirable solutions" are the ones that have deterministic behavior and use the least number of states. I experimented with a few ways to achieve these goals. The simplest and one of the most efficient method I found was to <strong>penalize the entropy ($H$) of the average state probability across time steps</strong>: $$ p_i = {1 \over N_\mathsf{step}}\sum_t s^t_i $$ $$ H = -\sum_i p_i log(p_i) $$</p>
<p>The total loss is therefore $$ L_\mathsf{total} = L_\mathsf{error} + w_H H $$</p>
<p>where $w_H$ is a regularization coefficient that is set by default to 0.01. This regularization happened to be sufficient to steer optimization to the deterministic solution that only uses two states:</p>
<figure>
</figure>
<p>We see that states 'A' and 'H' of the learned machine are equivalent to the states 'A' and 'B' of the expected minimal FSM.</p>
<h2>Robustness and initialization</h2>
<p>Let's check how robust and dependent on the initialization the solution is. <code>init_fsm</code> function uses the <code>key</code> argument to generate the initial FSM parameters for the optimization. I wrote a <a href="https://colab.research.google.com/github/znah/post-diff-fsm/blob/master/notebooks/diff_fsm_jax.ipynb#scrollTo=BVR4SIC3cbS2"><code>Trainer</code></a> class that encapsulates the task training and evaluation, and can be used like this:</p>
<p>x = '01010100100111111'</p>
<p>y0 = '01000100000101010'</p>
<p>trainer = Trainer(x, y0)</p>
<p>key = jax.random.PRNGKey(1)</p>
<p>result = trainer.run(key)</p>
<p>Returned data contains the final learned parameters, training log and evaluation result (number of errors and used states) by the deterministic (decoded with <code>hard=True</code>) version of the learned FSM. Here I'd like to give a great shout-out to <a href="https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html#jax.vmap"><code>jax.vmap</code></a> function, because with its help running 100 parallel experiments with different random keys boils down to the following lines:</p>
<p>keys = jax.random.split(key, 100)</p>
<p>results = jax.vmap(trainer.run)(keys)</p>
<p>By default <a href="https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html#jax.vmap"><code>vmap</code></a> replicates the nested function by introducing an extra dimension to all input and output arrays. Single line change creates many Adam optimizers and runs many training loops in parallel. The diagram below shows results from running 100 randomly initialized experiments with four different values of $w_H$. We see that strong regularization increases the number of runs that have converged to the minimal two-state FSM, but also increases the number of erroneous runs, i.e. those that produced the machine that couldn't correctly process the training sequence.</p>
<figure></figure>
<p>Entropy regularization is just one of many possible methods to steer the optimization towards desired solutions. Intuitively, machines that make fewer state transitions during processing of the training sequence are executing simpler, and thus, preferable behavior. Of course we may try to introduce another loss term that would penalize state transitions, but I'd like to demonstrate another approach here. What I tried is to initialize the transition tensor $T$ with a bias towards preserving the current state and disregarding the input. I do so by adding the identity matrix, that gets broadcasted for all input characters.</p>
<p> T = jax.random.normal(k1, [CHAR_N, STATE_N, STATE_N]) * noise</p>
<p> T += jp.eye(STATE_N) * lazy_bias</p>
<p>Setting <code>lazy_bias=1.0</code> happened to dramatically change the distribution of training outcomes, leading to higher chances of finding the minimal solution for lower values of $w_H$.</p>
<figure></figure>
<p>My observations showed that for simple problems the optimization is not particularly sensitive to the choice of $w_H$, so I picked the value 0.01</p>
<h2>More toy problems</h2>
<p>To see how the procedure described above handles different problems, I've created a "dataset" of 10 toy tasks and ran 1000 training attempts for each of them (thanks to JAX it takes a few seconds per task even on CPU). Almost all training runs found solutions that correctly process given training sequences. Only a few erroneous state machines were produced for tasks 1 and 9. The table below shows distribution of the number of states in discovered state machines, and minimal FSM diagram(s). </p>
<p>I rename states, assigning letters in the order of the first visit during processing of the training input. This way we can equalize solutions that only differ by the states' order. Note that these diagrams only show nodes and edges that were traversed during processing of the training sequence, missing edges are considered undefined (e.g. task 3 FSM doesn't have an edge corresponding to the input "1" from the state B'").</p>
<div>
<p><a href="https://colab.research.google.com/github/google-research/self-organising-systems/blob/master/notebooks/diff_fsm_jax.ipynb#scrollTo=X7vKvta5-Nle">Reproduce in a Notebook</a></p><div>
 <div><pre><b>Task 0: skip every second '1'</b>
Training input : 01010100100111111
Expected outout : 01000100000101010
</pre>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
</div>
<div>
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/1
B-&gt;A
1/0
B-&gt;B
0/0
<p>
 x 978</p></div></div>
<div>
 <div><pre><b>Task 1: skip every third '1'</b>
Training input : 0101010000111111
Expected outout : 0101000000110110
</pre>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
</div>
<div>
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/1
B-&gt;B
0/0
C
C
B-&gt;C
1/1
C-&gt;A
1/0
C-&gt;C
0/0
<p>
 x 676</p></div></div>
<div>
 <div><pre><b>Task 2: emit every third '1'</b>
Training input : 01010100111111
Expected outout : 00000100001001
</pre>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
</div>
<div>
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/0
B-&gt;B
0/0
C
C
B-&gt;C
1/0
C-&gt;A
1/1
C-&gt;C
0/0
<p>
 x 469</p></div></div>
<div>
 <div><pre><b>Task 3: emit 0s, when meet 1, start emitting 1s</b>
Training input : 00010000
Expected outout : 00011111
</pre>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
</div>
<div>
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/1
B-&gt;B
0/1
<p>
 x 999</p></div></div>
<div>
 <div><pre><b>Task 4: invert bits</b>
Training input : 01001110
Expected outout : 10110001
</pre>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
</div>
<div>
%3
A
A
A-&gt;A
1/0
A-&gt;A
0/1
<p>
 x 1000</p></div></div>
<div>
 <div><pre><b>Task 5: generate repeating '01' sequence</b>
Training input : 00000000
Expected outout : 01010101
</pre>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
</div>
<div>
%3
A
A
B
B
A-&gt;B
0/0
B-&gt;A
0/1
<p>
 x 1000</p></div></div>
<div>
 <div><pre><b>Task 6: replace '1'-&gt;'111'</b>
Training input : 00100001000
Expected outout : 00111001110
</pre>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
</div>
<div>
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/1
C
C
B-&gt;C
0/1
C-&gt;A
0/1
<p>
 x 1000</p></div></div>
<div>
 <div><pre><b>Task 7: shift right by 1 bit</b>
Training input : 0110010011001
Expected outout : 0011001001100
</pre>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
</div>
<div>
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/0
B-&gt;A
0/1
B-&gt;B
1/1
<p>
 x 1000</p></div></div>
<div>
 <div><pre><b>Task 8: shift right by 2 bits</b>
Training input : 000111010001100100
Expected outout : 000001110100011001
</pre>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
</div>
<div>
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/0
C
C
B-&gt;C
1/0
D
D
B-&gt;D
0/0
C-&gt;C
1/1
C-&gt;D
0/1
D-&gt;A
0/1
D-&gt;B
1/1
<p>
 x 919</p></div></div>
<div>
 <div><pre><b>Task 9: skip every fourth '1'</b>
Training input : 01010100100111111
Expected outout : 01010100000111011
</pre>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
</div>
<div>
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/1
B-&gt;B
0/0
C
C
B-&gt;C
1/1
C-&gt;C
0/0
D
D
C-&gt;D
1/1
D-&gt;A
1/0
D-&gt;D
0/0
<p>
 x 84</p></div><b>Click to see all FSMs (6 more)</b>
 <div>
%3
A
A
C
C
A-&gt;C
1/1
B
B
A-&gt;B
0/0
D
D
C-&gt;D
1/0
B-&gt;A
1/1
B-&gt;C
0/0
D-&gt;B
1/1
D-&gt;D
0/0
<p>
 x 83</p></div><div>
%3
A
A
C
C
A-&gt;C
1/1
B
B
A-&gt;B
0/0
C-&gt;B
1/0
D
D
C-&gt;D
0/0
B-&gt;A
1/1
B-&gt;C
0/0
D-&gt;B
1/1
<p>
 x 37</p></div><div>
%3
A
A
B
B
A-&gt;B
0/0
D
D
A-&gt;D
1/1
C
C
C-&gt;B
0/0
C-&gt;B
1/0
B-&gt;A
1/1
B-&gt;C
0/0
D-&gt;C
1/1
<p>
 x 10</p></div><div>
%3
A
A
B
B
A-&gt;B
0/0
A-&gt;B
1/1
C
C
B-&gt;C
1/1
D
D
B-&gt;D
0/0
C-&gt;A
1/0
C-&gt;D
0/0
D-&gt;A
1/1
D-&gt;C
0/0
<p>
 x 1</p></div><div>
%3
A
A
B
B
A-&gt;B
0/0
A-&gt;B
1/1
C
C
C-&gt;B
0/0
D
D
C-&gt;D
1/1
B-&gt;C
1/1
B-&gt;D
0/0
D-&gt;A
0/0
D-&gt;B
1/0
<p>
 x 1</p></div><div>
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/1
C
C
D
D
C-&gt;D
1/1
C-&gt;D
0/0
D-&gt;A
1/0
D-&gt;C
0/0
B-&gt;C
1/1
B-&gt;C
0/0
<p>
 x 1</p></div></div>
</div>
<p>For <i>most</i> problems it's easy to validate that the intended solution was indeed found. The task 8 ("shift the string right by two characters") FSM may look a bit cumbersome, but it's easy to understand what the state machine is doing by renaming the states: <code>A-&gt;00, B-&gt;01, C-&gt;11, D-&gt;10</code>. Now the label of each state reflects two most recent input bits, stored by the FSM and ready to be emitted.</p>
<figure></figure>
<p>Finally, task 9 (replace every 4th '1' with 0') happened to be an interesting example of an insufficiently defined problem. I expected the optimization to find the solution similar to problems 0 and 1: a loop with 4 states, that counts a number of encountered '1'-s and replaces the last with '0'. Turns out that there is more than one 4-state solution for the training sequence <code>'01010100100111111'</code>. Optimization found a 4-state FSM 217 times, only 84 out of which were the machine I expected. If you expand the last cell of the table, which shows all discovered FSMs and their counts, you will see that we happened to be quite fortunate, because the second most popular machine count is smaller by just one. The other state machines have overfitted to the particular training sequence and are not going to perform correctly on other inputs. The usual rule of machine learning, when low training loss doesn't guarantee generalization, applies here as well. Additional initialization and regularization strategies may increase the chance of finding the well performing solutions.</p>
<p>The amount of literature published on almost any topic is overwhelming these days. Finite State Machine learning is not an exception, especially given that the concept of state machines lies at the very foundation of computer science. Here are just some of the vast number of relevant articles, some of which are not directly related to FSM learning, but still conceptually close. Please feel free to leave additional pointers and comments in the <a href="https://github.com/google-research/self-organising-systems/discussions/6">article discussion</a>.</p>
<p><a href="https://learnlib.de/">LearnLib</a> is an extensive library of FSM learning algorithms that are used in practice of e.g. communication systems modeling. An extensive discussion of probabilistic FSM learning can also be found here .</p>
<p>State machines were also shown to be useful in more complex differentiable machine learning systems . For example, it was shown that it's possible to extract compact FSMs from RNN-based game policy controllers and string matchers, using post-training quantization.</p>
<p>On another end Neural Turing Machines take inspiration from classical discrete computer science concepts and implement trainable differentiable analogies.</p>
<p>An interesting recent work tries to solve toy sequential string processing problems using very small recurrent neural networks instead of discrete state machines, and uses genetic algorithms instead of differentiable optimization. Authors make emphasis on interpretability and generalization of very compact models.</p>
<h2>Conclusion</h2>
<p>In this tutorial we've seen that simple deterministic string processing finite state machines can be discovered as a limit case of differentiable stochastic state machines. I've only covered simple toy problems in this tutorial, but differentiability allows incorporating FSMs in more complex end-to-end pipelines that I'm going to talk about in future publications.</p>
</div>||||I||||Differentiable Finite State Machines

Continuous optimization for discrete control

Reproduce in a Notebook Discuss on

In this blog post I'd like to show how differentiable optimization can be used to learn Finite State Machines (FSM) for solving toy string processing tasks. I'll show how simple regularization and initialization techniques can steer continuous optimization towards finding discrete deterministic solutions and increase the training success rate. I think that experiments shown here may have some educational value, e.g. in demonstrating less conventional (and perhaps unexpected) uses of differentiable programming and some elegant JAX tricks.

There exists a vast literature on various approaches to FSM synthesis from data, so it's highly likely that the particular method described here is already well known. A quick review of related literature will be given in the Related Work section.

Here is an example of a toy problem we are going to solve. Consider a pair of input-output strings:


 In: 01010100100111111
Out: 01000100000101010


We'd like to design an algorithm that would produce the expected output string if we feed it with the input string character by character. There are many different possibilities to perform this task. The brute force way would be just to store the expected output string and emit it character by character without paying attention to the input at all. This is actually a pretty complex solution, for example implementing such an algorithm as a state machine would require the number of states that is equal to the string length. Intuitively, ideas behind Occam's razor and Kolmogorov's complexity suggest that simpler solutions are preferable. The simplicity gives us hope that the solution would be interpretable and generalize to input sequences beyond the training set.

The simplest algorithm for the string pair above scans the input and replaces every second '1' character it encounters with '0'. A type of FSM called Mealy Machine is natural way of representing such algorithms:

 T: state' R: output s0 
 A B 0 1 
input state 
0 A 1.00 0.00 1.00 0.00 1.00
B 0.00 1.00 1.00 0.00 0.00
1 A 0.00 1.00 0.00 1.00 1.00
B 1.00 0.00 1.00 0.00 0.00


Mealy machine and its transition table

This machine starts at the state "A". The label at each edge has a form "<input>/<output>": <input> character activates the transition along the given edge, and <output> character is emitted to the output string during the transitions. Transition tables are a common way to represent FSMs. Each row of the transition table above corresponds to a different combination of input and state, and represents the new state' and the output value, encoded as two one-hot vectors. This table also includes the column s0, a one-hot vector, encoding the starting FSM state (which is replicated two times for completeness). Formally, this table can be represented by three tensors with following axes:

$$ \begin{aligned} T&:[\mathsf{input}, \mathsf{state}, \mathsf{state'}] \\ R&:[\mathsf{input}, \mathsf{state}, \mathsf{output}] \\ s^0&:[\mathsf{state}] \end{aligned} $$

Now we can express the FSM output and state transition computation as the following sums

$$ \begin{aligned} y^t_{\mathsf{output}} &= \sum_{\mathsf{input},\mathsf{state}} x^t_\mathsf{input} s^t_\mathsf{state} R_{\mathsf{input}, \mathsf{state}, \mathsf{output}} \\ s^{t+1}_{\mathsf{state'}} &= \sum_{\mathsf{input},\mathsf{state}} x^t_\mathsf{input} s^t_\mathsf{state} T_{\mathsf{input}, \mathsf{state}, \mathsf{state'}} \end{aligned}$$

where x, y, and s are sequences of one-hot vectors, representing inputs, outputs and FSM states. Luckily, modern differentiable programming frameworks (I'll be using JAX in this tutorial) provide the powerful 'einsum' function to compute these kinds of expressions. It's easy to combine it with the 'scan' primitive to compute FSM outputs and intermediate states for a given input:

FSM = namedtuple('FSM', 'T R s0') def run_fsm(fsm: FSM, inputs): def f(s, x): y = jp.einsum('x,s,xsy->y', x, s, fsm.R) s1 = jp.einsum('x,s,xst->t', x, s, fsm.T) return s1, (y, s1) _, (outputs, states) = jax.lax.scan(f, fsm.s0, inputs) return outputs, jp.vstack([fsm.s0, states])

For now we were presenting states and characters as one-hot vectors, but we can also think of them as probability distributions. This relaxes the discrete set of one-hot vectors to a continuous space of vectors of non-negative values that sum up to one. The resulting model is called a probabilistic automaton and the above math still holds.

Given that we represented a state machine in terms of differentiable operations, why don't we try to use gradient descent to find a machine given the input-output pair? We need to make sure that we use the right parameterisation that enforces rows of $T$, $R$ and $s^0$ to be probability distributions. Usually this is achieved by using the 'softmax' function. $(T, R, s^0)$ are parametrized with real-valued tensors and softmax is applied along the last dimension to turn them into probabilities. Our toy examples are going to use the alphabet that consists of just two characters: '0' and '1'. We assume that we don't know the required number of FSM states in advance, so in the spirit of deep learning we are going to use an overparameterized FSM representation that has a maximum of 8 states.

CHAR_N = 2

STATE_N = 8

Params = namedtuple('Params', 'T R s0')

def init_fsm(key, noise=1e-3) -> Params:

k1, k2, k3 = jax.random.split(key, 3)

T = jax.random.normal(k1, [CHAR_N, STATE_N, STATE_N]) * noise

R = jax.random.normal(k2, [CHAR_N, STATE_N, CHAR_N]) * noise

s0 = jax.random.normal(k3, [STATE_N]) * noise

return Params(T, R, s0)

def hardmax(x):

return nn.one_hot(x.argmax(-1), x.shape[-1])

def decode_fsm(params: Params, hard=False) -> FSM:

T, R, s0 = params

f = hardmax if hard else nn.softmax

return FSM(f(T), f(R), f(s0))

'init_fsm' function creates a randomly initialized FSM. Parameters are set to very small values, so all distributions are close to uniform. 'decode_fsm' serves to convert parameters into FSM matrices to consume by 'run_fsm'. It has two operation modes: 'soft', when softmax is used, and 'hard', when all distributions get collapsed into one-hot vectors, and stochastic FSM becomes deterministic.

Let's try to simply minimize the sum of squared-differences between the FSM output and the reference output sequence:

fsm = decode_fsm(params)

y, s = run_fsm(fsm, x)

error = jp.square(y-y0).sum()

Throwing ADAM optimizer into minimizing the error indeed makes it go down. Getting there took a bit of fiddling with the optimizer parameters. Motivated by the lack of loss stochasticity and small number of training steps (400) I use a pretty high learning rate (0.25) and set beta1=beta2=0.5.

Here is the diagram showing input, output and state probability distributions produced by the resulting FSM at different time steps (circle size is proportional to the probability):

We see that although the model manages to perfectly reproduce the desired output for the given input string, it's using many more states than necessary, and at some timesteps (e.g. 0 or 3) the probability is distributed among a number of states, so the process doesn't look deterministic. In machine learning, the usual strategy to steer the optimization into a more desirable solutions region is to introduce additional regularizer objectives. In our case "desirable solutions" are the ones that have deterministic behavior and use the least number of states. I experimented with a few ways to achieve these goals. The simplest and one of the most efficient method I found was to penalize the entropy ($H$) of the average state probability across time steps: $$ p_i = {1 \over N_\mathsf{step}}\sum_t s^t_i $$ $$ H = -\sum_i p_i log(p_i) $$

The total loss is therefore $$ L_\mathsf{total} = L_\mathsf{error} + w_H H $$

where $w_H$ is a regularization coefficient that is set by default to 0.01. This regularization happened to be sufficient to steer optimization to the deterministic solution that only uses two states:

We see that states 'A' and 'H' of the learned machine are equivalent to the states 'A' and 'B' of the expected minimal FSM.

Robustness and initialization

Let's check how robust and dependent on the initialization the solution is. init_fsm function uses the key argument to generate the initial FSM parameters for the optimization. I wrote a Trainer class that encapsulates the task training and evaluation, and can be used like this:

x = '01010100100111111'

y0 = '01000100000101010'

trainer = Trainer(x, y0)

key = jax.random.PRNGKey(1)

result = trainer.run(key)

Returned data contains the final learned parameters, training log and evaluation result (number of errors and used states) by the deterministic (decoded with hard=True) version of the learned FSM. Here I'd like to give a great shout-out to jax.vmap function, because with its help running 100 parallel experiments with different random keys boils down to the following lines:

keys = jax.random.split(key, 100)

results = jax.vmap(trainer.run)(keys)

By default vmap replicates the nested function by introducing an extra dimension to all input and output arrays. Single line change creates many Adam optimizers and runs many training loops in parallel. The diagram below shows results from running 100 randomly initialized experiments with four different values of $w_H$. We see that strong regularization increases the number of runs that have converged to the minimal two-state FSM, but also increases the number of erroneous runs, i.e. those that produced the machine that couldn't correctly process the training sequence.

Entropy regularization is just one of many possible methods to steer the optimization towards desired solutions. Intuitively, machines that make fewer state transitions during processing of the training sequence are executing simpler, and thus, preferable behavior. Of course we may try to introduce another loss term that would penalize state transitions, but I'd like to demonstrate another approach here. What I tried is to initialize the transition tensor $T$ with a bias towards preserving the current state and disregarding the input. I do so by adding the identity matrix, that gets broadcasted for all input characters.

T = jax.random.normal(k1, [CHAR_N, STATE_N, STATE_N]) * noise

T += jp.eye(STATE_N) * lazy_bias

Setting lazy_bias=1.0 happened to dramatically change the distribution of training outcomes, leading to higher chances of finding the minimal solution for lower values of $w_H$.

My observations showed that for simple problems the optimization is not particularly sensitive to the choice of $w_H$, so I picked the value 0.01

More toy problems

To see how the procedure described above handles different problems, I've created a "dataset" of 10 toy tasks and ran 1000 training attempts for each of them (thanks to JAX it takes a few seconds per task even on CPU). Almost all training runs found solutions that correctly process given training sequences. Only a few erroneous state machines were produced for tasks 1 and 9. The table below shows distribution of the number of states in discovered state machines, and minimal FSM diagram(s).

I rename states, assigning letters in the order of the first visit during processing of the training input. This way we can equalize solutions that only differ by the states' order. Note that these diagrams only show nodes and edges that were traversed during processing of the training sequence, missing edges are considered undefined (e.g. task 3 FSM doesn't have an edge corresponding to the input "1" from the state B'").

 Reproduce in a Notebook
 Task 0: skip every second '1'
 Training input : 01010100100111111
 Expected outout : 01000100000101010
 
 A 0/0 B 1/1 1/0 0/0 x 978
 Task 1: skip every third '1'
 Training input : 0101010000111111
 Expected outout : 0101000000110110
 
 A 0/0 B 1/1 0/0 C 1/1 1/0 0/0 x 676
 Task 2: emit every third '1'
 Training input : 01010100111111
 Expected outout : 00000100001001
 
 A 0/0 B 1/0 0/0 C 1/0 1/1 0/0 x 469
 Task 3: emit 0s, when meet 1, start emitting 1s
 Training input : 00010000
 Expected outout : 00011111
 
 A 0/0 B 1/1 0/1 x 999
 Task 4: invert bits
 Training input : 01001110
 Expected outout : 10110001
 
 A 1/0 0/1 x 1000
 Task 5: generate repeating '01' sequence
 Training input : 00000000
 Expected outout : 01010101
 
 A B 0/0 0/1 x 1000
 Task 6: replace '1'->'111'
 Training input : 00100001000
 Expected outout : 00111001110
 
 A 0/0 B 1/1 C 0/1 0/1 x 1000
 Task 7: shift right by 1 bit
 Training input : 0110010011001
 Expected outout : 0011001001100
 
 A 0/0 B 1/0 0/1 1/1 x 1000
 Task 8: shift right by 2 bits
 Training input : 000111010001100100
 Expected outout : 000001110100011001
 
 A 0/0 B 1/0 C 1/0 D 0/0 1/1 0/1 0/1 1/1 x 919
 Task 9: skip every fourth '1'
 Training input : 01010100100111111
 Expected outout : 01010100000111011
 
 A 0/0 B 1/1 0/0 C 1/1 0/0 D 1/1 1/0 0/0 x 84
 Click to see all FSMs (6 more)
 A C 1/1 B 0/0 D 1/0 1/1 0/0 1/1 0/0 x 83
 A C 1/1 B 0/0 1/0 D 0/0 1/1 0/0 1/1 x 37
 A B 0/0 D 1/1 C 0/0 1/0 1/1 0/0 1/1 x 10
 A B 0/0 1/1 C 1/1 D 0/0 1/0 0/0 1/1 0/0 x 1
 A B 0/0 1/1 C 0/0 D 1/1 1/1 0/0 0/0 1/0 x 1
 A 0/0 B 1/1 C D 1/1 0/0 1/0 0/0 1/1 0/0 x 1

For most problems it's easy to validate that the intended solution was indeed found. The task 8 ("shift the string right by two characters") FSM may look a bit cumbersome, but it's easy to understand what the state machine is doing by renaming the states: A->00, B->01, C->11, D->10. Now the label of each state reflects two most recent input bits, stored by the FSM and ready to be emitted.

Finally, task 9 (replace every 4th '1' with 0') happened to be an interesting example of an insufficiently defined problem. I expected the optimization to find the solution similar to problems 0 and 1: a loop with 4 states, that counts a number of encountered '1'-s and replaces the last with '0'. Turns out that there is more than one 4-state solution for the training sequence '01010100100111111'. Optimization found a 4-state FSM 217 times, only 84 out of which were the machine I expected. If you expand the last cell of the table, which shows all discovered FSMs and their counts, you will see that we happened to be quite fortunate, because the second most popular machine count is smaller by just one. The other state machines have overfitted to the particular training sequence and are not going to perform correctly on other inputs. The usual rule of machine learning, when low training loss doesn't guarantee generalization, applies here as well. Additional initialization and regularization strategies may increase the chance of finding the well performing solutions.

Related work

The amount of literature published on almost any topic is overwhelming these days. Finite State Machine learning is not an exception, especially given that the concept of state machines lies at the very foundation of computer science. Here are just some of the vast number of relevant articles, some of which are not directly related to FSM learning, but still conceptually close. Please feel free to leave additional pointers and comments in the article discussion.

LearnLib is an extensive library of FSM learning algorithms that are used in practice of e.g. communication systems modeling. An extensive discussion of probabilistic FSM learning can also be found here .

State machines were also shown to be useful in more complex differentiable machine learning systems . For example, it was shown that it's possible to extract compact FSMs from RNN-based game policy controllers and string matchers, using post-training quantization.

On another end Neural Turing Machines take inspiration from classical discrete computer science concepts and implement trainable differentiable analogies.

An interesting recent work tries to solve toy sequential string processing problems using very small recurrent neural networks instead of discrete state machines, and uses genetic algorithms instead of differentiable optimization. Authors make emphasis on interpretability and generalization of very compact models.

Conclusion

In this tutorial we've seen that simple deterministic string processing finite state machines can be discovered as a limit case of differentiable stochastic state machines. I've only covered simple toy problems in this tutorial, but differentiability allows incorporating FSMs in more complex end-to-end pipelines that I'm going to talk about in future publications.

Acknowledgements

I'd like to thank my colleagues Ettore Randazzo, Eyvind Niklasson, Dominik Roblek and David Ha for their feedback, suggestions and relevant work pointers. This article was prepared using the Distill template and customized Google Docs->HTML workflow.

Reproduce in a Notebook
 Discuss on 
In this blog post I'd like to show how differentiable optimization can be used to learn Finite State Machines (FSM) for solving toy string processing tasks. I'll show how simple regularization and initialization techniques can steer continuous optimization towards finding discrete deterministic solutions and increase the training success rate. I think that experiments shown here may have some educational value, e.g. in demonstrating less conventional (and perhaps unexpected) uses of differentiable programming and some elegant JAX tricks.
There exists a vast literature on various approaches to FSM synthesis from data, so it's highly likely that the particular method described here is already well known. A quick review of related literature will be given in the Related Work section.
Here is an example of a toy problem we are going to solve. Consider a pair of input-output strings: 
  In: 01010100100111111
Out: 01000100000101010
 
We'd like to design an algorithm that would produce the expected output string if we feed it with the input string character by character. There are many different possibilities to perform this task. The brute force way would be just to store the expected output string and emit it character by character without paying attention to the input at all. This is actually a pretty complex solution, for example implementing such an algorithm as a state machine would require the number of states that is equal to the string length. Intuitively, ideas behind Occam's razor and Kolmogorov's complexity suggest that simpler solutions are preferable. The simplicity gives us hope that the solution would be interpretable and generalize to input sequences beyond the training set.
The simplest algorithm for the string pair above scans the input and replaces every second '1' character it encounters with '0'. A type of FSM called Mealy Machine is natural way of representing such algorithms:
 
 
  
  
 T: state'
 R: output
 s0
 
 
  
  
 A
 B
 0
 1
  
 
 
 input
 state
  
  
  
  
  
 
 
 
 
 0
 A
 1.00
 0.00
 1.00
 0.00
 1.00
 
 
 B
 0.00
 1.00
 1.00
 0.00
 0.00
 
 
 1
 A
 0.00
 1.00
 0.00
 1.00
 1.00
 
 
 B
 1.00
 0.00
 1.00
 0.00
 0.00
 
 
Mealy machine and its transition table
This machine starts at the state "A". The label at each edge has a form "&lt;input&gt;/&lt;output&gt;": &lt;input&gt; character activates the transition along the given edge, and &lt;output&gt; character is emitted to the output string during the transitions. Transition tables are a common way to represent FSMs. Each row of the transition table above corresponds to a different combination of input and state, and represents the new state' and the output value, encoded as two one-hot vectors. This table also includes the column s0, a one-hot vector, encoding the starting FSM state (which is replicated two times for completeness). Formally, this table can be represented by three tensors with following axes:
$$ \begin{aligned}
T&amp;:[\mathsf{input}, \mathsf{state}, \mathsf{state'}] \\
R&amp;:[\mathsf{input}, \mathsf{state}, \mathsf{output}] \\
s^0&amp;:[\mathsf{state}]
\end{aligned} $$
Now we can express the FSM output and state transition computation as the following sums
$$ \begin{aligned}
y^t_{\mathsf{output}} &amp;= \sum_{\mathsf{input},\mathsf{state}} x^t_\mathsf{input} s^t_\mathsf{state} R_{\mathsf{input}, \mathsf{state}, \mathsf{output}} \\
s^{t+1}_{\mathsf{state'}} &amp;= \sum_{\mathsf{input},\mathsf{state}} x^t_\mathsf{input} s^t_\mathsf{state} T_{\mathsf{input}, \mathsf{state}, \mathsf{state'}}
\end{aligned}$$
where x, y, and s are sequences of one-hot vectors, representing inputs, outputs and FSM states. Luckily, modern differentiable programming frameworks (I'll be using JAX in this tutorial) provide the powerful '  einsum  ' function to compute these kinds of expressions. It's easy to combine it with the '  scan  ' primitive to compute FSM outputs and intermediate states for a given input:
FSM = namedtuple('FSM', 'T R s0')
def run_fsm(fsm: FSM, inputs):
 def f(s, x):
 y = jp.einsum('x,s,xsy-&gt;y', x, s, fsm.R)
 s1 = jp.einsum('x,s,xst-&gt;t', x, s, fsm.T)
 return s1, (y, s1)
 _, (outputs, states) = jax.lax.scan(f, fsm.s0, inputs)
 return outputs, jp.vstack([fsm.s0, states])
For now we were presenting states and characters as one-hot vectors, but we can also think of them as probability distributions. This relaxes the discrete set of one-hot vectors to a continuous space of vectors of non-negative values that sum up to one. The resulting model is called a probabilistic automaton and the above math still holds.
Given that we represented a state machine in terms of differentiable operations, why don't we try to use gradient descent to find a machine given the input-output pair? We need to make sure that we use the right parameterisation that enforces rows of $T$, $R$ and $s^0$ to be probability distributions. Usually this is achieved by using the 'softmax' function. $(T, R, s^0)$ are parametrized with real-valued tensors and softmax is applied along the last dimension to turn them into probabilities. Our toy examples are going to use the alphabet that consists of just two characters: '0' and '1'. We assume that we don't know the required number of FSM states in advance, so in the spirit of deep learning we are going to use an overparameterized FSM representation that has a maximum of 8 states.
CHAR_N = 2
STATE_N = 8
Params = namedtuple('Params', 'T R s0')
def init_fsm(key, noise=1e-3) -&gt; Params:
 k1, k2, k3 = jax.random.split(key, 3)
 T = jax.random.normal(k1, [CHAR_N, STATE_N, STATE_N]) * noise
 R = jax.random.normal(k2, [CHAR_N, STATE_N, CHAR_N]) * noise
 s0 = jax.random.normal(k3, [STATE_N]) * noise
 return Params(T, R, s0)
def hardmax(x):
 return nn.one_hot(x.argmax(-1), x.shape[-1])
def decode_fsm(params: Params, hard=False) -&gt; FSM:
 T, R, s0 = params
 f = hardmax if hard else nn.softmax
 return FSM(f(T), f(R), f(s0))
 'init_fsm' function creates a randomly initialized FSM. Parameters are set to very small values, so all distributions are close to uniform. 'decode_fsm' serves to convert parameters into FSM matrices to consume by 'run_fsm'. It has two operation modes: 'soft', when softmax is used, and 'hard', when all distributions get collapsed into one-hot vectors, and stochastic FSM becomes deterministic.
Let's try to simply minimize the sum of squared-differences between the FSM output and the reference output sequence:
fsm = decode_fsm(params)
y, s = run_fsm(fsm, x)
error = jp.square(y-y0).sum()
Throwing ADAM optimizer into minimizing the error indeed makes it go down. Getting there took a bit of fiddling with the optimizer parameters. Motivated by the lack of loss stochasticity and small number of training steps (400) I use a pretty high learning rate (0.25) and set beta1=beta2=0.5.
Here is the diagram showing input, output and state probability distributions produced by the resulting FSM at different time steps (circle size is proportional to the probability):
 
We see that although the model manages to perfectly reproduce the desired output for the given input string, it's using many more states than necessary, and at some timesteps (e.g. 0 or 3) the probability is distributed among a number of states, so the process doesn't look deterministic. In machine learning, the usual strategy to steer the optimization into a more desirable solutions region is to introduce additional regularizer objectives. In our case "desirable solutions" are the ones that have deterministic behavior and use the least number of states. I experimented with a few ways to achieve these goals. The simplest and one of the most efficient method I found was to penalize the entropy ($H$) of the average state probability across time steps: $$ p_i = {1 \over N_\mathsf{step}}\sum_t s^t_i $$ $$ H = -\sum_i p_i log(p_i) $$
The total loss is therefore $$ L_\mathsf{total} = L_\mathsf{error} + w_H H $$
where $w_H$ is a regularization coefficient that is set by default to 0.01. This regularization happened to be sufficient to steer optimization to the deterministic solution that only uses two states:
We see that states 'A' and 'H' of the learned machine are equivalent to the states 'A' and 'B' of the expected minimal FSM.
Robustness and initialization
Let's check how robust and dependent on the initialization the solution is. init_fsm function uses the key argument to generate the initial FSM parameters for the optimization. I wrote a  Trainer  class that encapsulates the task training and evaluation, and can be used like this:
x = '01010100100111111'
y0 = '01000100000101010'
trainer = Trainer(x, y0)
key = jax.random.PRNGKey(1)
result = trainer.run(key)
Returned data contains the final learned parameters, training log and evaluation result (number of errors and used states) by the deterministic (decoded with hard=True) version of the learned FSM. Here I'd like to give a great shout-out to  jax.vmap  function, because with its help running 100 parallel experiments with different random keys boils down to the following lines:
keys = jax.random.split(key, 100)
results = jax.vmap(trainer.run)(keys)
By default  vmap  replicates the nested function by introducing an extra dimension to all input and output arrays. Single line change creates many Adam optimizers and runs many training loops in parallel. The diagram below shows results from running 100 randomly initialized experiments with four different values of $w_H$. We see that strong regularization increases the number of runs that have converged to the minimal two-state FSM, but also increases the number of erroneous runs, i.e. those that produced the machine that couldn't correctly process the training sequence.
 
Entropy regularization is just one of many possible methods to steer the optimization towards desired solutions. Intuitively, machines that make fewer state transitions during processing of the training sequence are executing simpler, and thus, preferable behavior. Of course we may try to introduce another loss term that would penalize state transitions, but I'd like to demonstrate another approach here. What I tried is to initialize the transition tensor $T$ with a bias towards preserving the current state and disregarding the input. I do so by adding the identity matrix, that gets broadcasted for all input characters.
 T = jax.random.normal(k1, [CHAR_N, STATE_N, STATE_N]) * noise
 T += jp.eye(STATE_N) * lazy_bias
Setting lazy_bias=1.0 happened to dramatically change the distribution of training outcomes, leading to higher chances of finding the minimal solution for lower values of $w_H$.
 
My observations showed that for simple problems the optimization is not particularly sensitive to the choice of $w_H$, so I picked the value 0.01
More toy problems
To see how the procedure described above handles different problems, I've created a "dataset" of 10 toy tasks and ran 1000 training attempts for each of them (thanks to JAX it takes a few seconds per task even on CPU). Almost all training runs found solutions that correctly process given training sequences. Only a few erroneous state machines were produced for tasks 1 and 9. The table below shows distribution of the number of states in discovered state machines, and minimal FSM diagram(s). 
I rename states, assigning letters in the order of the first visit during processing of the training input. This way we can equalize solutions that only differ by the states' order. Note that these diagrams only show nodes and edges that were traversed during processing of the training sequence, missing edges are considered undefined (e.g. task 3 FSM doesn't have an edge corresponding to the input "1" from the state B'").
 Reproduce in a Notebook  
   Task 0: skip every second '1'
Training input : 01010100100111111
Expected outout : 01000100000101010
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/1
B-&gt;A
1/0
B-&gt;B
0/0
 x 978  
   Task 1: skip every third '1'
Training input : 0101010000111111
Expected outout : 0101000000110110
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/1
B-&gt;B
0/0
C
C
B-&gt;C
1/1
C-&gt;A
1/0
C-&gt;C
0/0
 x 676  
   Task 2: emit every third '1'
Training input : 01010100111111
Expected outout : 00000100001001
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/0
B-&gt;B
0/0
C
C
B-&gt;C
1/0
C-&gt;A
1/1
C-&gt;C
0/0
 x 469  
   Task 3: emit 0s, when meet 1, start emitting 1s
Training input : 00010000
Expected outout : 00011111
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/1
B-&gt;B
0/1
 x 999  
   Task 4: invert bits
Training input : 01001110
Expected outout : 10110001
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
%3
A
A
A-&gt;A
1/0
A-&gt;A
0/1
 x 1000  
   Task 5: generate repeating '01' sequence
Training input : 00000000
Expected outout : 01010101
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
%3
A
A
B
B
A-&gt;B
0/0
B-&gt;A
0/1
 x 1000  
   Task 6: replace '1'-&gt;'111'
Training input : 00100001000
Expected outout : 00111001110
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/1
C
C
B-&gt;C
0/1
C-&gt;A
0/1
 x 1000  
   Task 7: shift right by 1 bit
Training input : 0110010011001
Expected outout : 0011001001100
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/0
B-&gt;A
0/1
B-&gt;B
1/1
 x 1000  
   Task 8: shift right by 2 bits
Training input : 000111010001100100
Expected outout : 000001110100011001
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/0
C
C
B-&gt;C
1/0
D
D
B-&gt;D
0/0
C-&gt;C
1/1
C-&gt;D
0/1
D-&gt;A
0/1
D-&gt;B
1/1
 x 919  
   Task 9: skip every fourth '1'
Training input : 01010100100111111
Expected outout : 01010100000111011
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/1
B-&gt;B
0/0
C
C
B-&gt;C
1/1
C-&gt;C
0/0
D
D
C-&gt;D
1/1
D-&gt;A
1/0
D-&gt;D
0/0
 x 84  Click to see all FSMs (6 more)
 
%3
A
A
C
C
A-&gt;C
1/1
B
B
A-&gt;B
0/0
D
D
C-&gt;D
1/0
B-&gt;A
1/1
B-&gt;C
0/0
D-&gt;B
1/1
D-&gt;D
0/0
 x 83  
%3
A
A
C
C
A-&gt;C
1/1
B
B
A-&gt;B
0/0
C-&gt;B
1/0
D
D
C-&gt;D
0/0
B-&gt;A
1/1
B-&gt;C
0/0
D-&gt;B
1/1
 x 37  
%3
A
A
B
B
A-&gt;B
0/0
D
D
A-&gt;D
1/1
C
C
C-&gt;B
0/0
C-&gt;B
1/0
B-&gt;A
1/1
B-&gt;C
0/0
D-&gt;C
1/1
 x 10  
%3
A
A
B
B
A-&gt;B
0/0
A-&gt;B
1/1
C
C
B-&gt;C
1/1
D
D
B-&gt;D
0/0
C-&gt;A
1/0
C-&gt;D
0/0
D-&gt;A
1/1
D-&gt;C
0/0
 x 1  
%3
A
A
B
B
A-&gt;B
0/0
A-&gt;B
1/1
C
C
C-&gt;B
0/0
D
D
C-&gt;D
1/1
B-&gt;C
1/1
B-&gt;D
0/0
D-&gt;A
0/0
D-&gt;B
1/0
 x 1  
%3
A
A
A-&gt;A
0/0
B
B
A-&gt;B
1/1
C
C
D
D
C-&gt;D
1/1
C-&gt;D
0/0
D-&gt;A
1/0
D-&gt;C
0/0
B-&gt;C
1/1
B-&gt;C
0/0
 x 1  
For most problems it's easy to validate that the intended solution was indeed found. The task 8 ("shift the string right by two characters") FSM may look a bit cumbersome, but it's easy to understand what the state machine is doing by renaming the states: A-&gt;00, B-&gt;01, C-&gt;11, D-&gt;10. Now the label of each state reflects two most recent input bits, stored by the FSM and ready to be emitted.
 
Finally, task 9 (replace every 4th '1' with 0') happened to be an interesting example of an insufficiently defined problem. I expected the optimization to find the solution similar to problems 0 and 1: a loop with 4 states, that counts a number of encountered '1'-s and replaces the last with '0'. Turns out that there is more than one 4-state solution for the training sequence '01010100100111111'. Optimization found a 4-state FSM 217 times, only 84 out of which were the machine I expected. If you expand the last cell of the table, which shows all discovered FSMs and their counts, you will see that we happened to be quite fortunate, because the second most popular machine count is smaller by just one. The other state machines have overfitted to the particular training sequence and are not going to perform correctly on other inputs. The usual rule of machine learning, when low training loss doesn't guarantee generalization, applies here as well. Additional initialization and regularization strategies may increase the chance of finding the well performing solutions.
The amount of literature published on almost any topic is overwhelming these days. Finite State Machine learning is not an exception, especially given that the concept of state machines lies at the very foundation of computer science. Here are just some of the vast number of relevant articles, some of which are not directly related to FSM learning, but still conceptually close. Please feel free to leave additional pointers and comments in the article discussion.
 LearnLib is an extensive library of FSM learning algorithms that are used in practice of e.g. communication systems modeling. An extensive discussion of probabilistic FSM learning can also be found here .
State machines were also shown to be useful in more complex differentiable machine learning systems . For example, it was shown that it's possible to extract compact FSMs from RNN-based game policy controllers and string matchers, using post-training quantization.
On another end Neural Turing Machines take inspiration from classical discrete computer science concepts and implement trainable differentiable analogies.
An interesting recent work tries to solve toy sequential string processing problems using very small recurrent neural networks instead of discrete state machines, and uses genetic algorithms instead of differentiable optimization. Authors make emphasis on interpretability and generalization of very compact models.
Conclusion
In this tutorial we've seen that simple deterministic string processing finite state machines can be discovered as a limit case of differentiable stochastic state machines. I've only covered simple toy problems in this tutorial, but differentiability allows incorporating FSMs in more complex end-to-end pipelines that I'm going to talk about in future publications.
Differentiable Finite State Machines
https://google-research.github.io/self-organising-systems/2022/diff-fsm/
https://arxiv.org/abs/1802.03676v1
Differentiable Dynamic Programming for Structured Prediction and Attention
2018-02-11
Mensch; Arthur; Blondel; Mathieu
<div><div>
 
 
 
 
 
 
 <p><a href="/pdf/1802.03676v1">Download PDF</a></p><blockquote>
 Abstract: Dynamic programming (DP) solves a variety of structured combinatorial
problems by iteratively breaking them down into smaller subproblems. In spite
of their versatility, DP algorithms are usually non-differentiable, which
hampers their use as a layer in neural networks trained by backpropagation. To
address this issue, we propose to smooth the max operator in the dynamic
programming recursion, using a strongly convex regularizer. This allows to
relax both the optimal value and solution of the original combinatorial
problem, and turns a broad class of DP algorithms into differentiable
operators. Theoretically, we provide a new probabilistic perspective on
backpropagating through these DP operators, and relate them to inference in
graphical models. We derive two particular instantiations of our framework, a
smoothed Viterbi algorithm for sequence prediction and a smoothed DTW algorithm
for time-series alignment. We showcase these instantiations on two structured
prediction tasks and on structured and sparse attention for neural machine
translation.
 </blockquote>
 
 
 </div><div>
 <h2>Submission history</h2><p> From: Mathieu Blondel [<a href="/show-email/d0caec15/1802.03676">view email</a>]
 <br /><strong>[v1]</strong>
Sun, 11 Feb 2018 01:36:06 UTC (1,168 KB)<br />
 <strong><a href="/abs/1802.03676v2">[v2]</a></strong>
 Tue, 20 Feb 2018 06:01:45 UTC (1,168 KB)<br /></p></div>||||I|||| Skip to main content
 We gratefully acknowledge support from
 the Simons Foundation and member institutions.
 > stat > arXiv:1802.03676v1

 Help | Advanced Search

 All fields Title Author Abstract Comments Journal reference ACM classification MSC classification Report number arXiv identifier DOI ORCID arXiv author ID Help pages Full text
 Search
 GO

 quick links

 * Login
 * Help Pages
 * About

 Statistics > Machine Learning

 arXiv:1802.03676v1 (stat)
 [Submitted on 11 Feb 2018 (this version), latest version 20 Feb 2018 (v2)]

 Title: Differentiable Dynamic Programming for Structured Prediction and Attention

 Authors: Arthur Mensch, Mathieu Blondel
 Download PDF
 Abstract: Dynamic programming (DP) solves a variety of structured combinatorial problems by iteratively breaking them down into smaller subproblems. In spite of their versatility, DP algorithms are usually non-differentiable, which hampers their use as a layer in neural networks trained by backpropagation. To address this issue, we propose to smooth the max operator in the dynamic programming recursion, using a strongly convex regularizer. This allows to relax both the optimal value and solution of the original combinatorial problem, and turns a broad class of DP algorithms into differentiable operators. Theoretically, we provide a new probabilistic perspective on backpropagating through these DP operators, and relate them to inference in graphical models. We derive two particular instantiations of our framework, a smoothed Viterbi algorithm for sequence prediction and a smoothed DTW algorithm for time-series alignment. We showcase these instantiations on two structured prediction tasks and on structured and sparse attention for neural machine translation.
 Subjects: Machine Learning (stat.ML) ; Machine Learning (cs.LG)
 Cite as: arXiv:1802.03676 [stat.ML] 
 (or arXiv:1802.03676v1 [stat.ML] for this version) 
 https://doi.org/10.48550/arXiv.1802.03676 
 Focus to learn more 
 arXiv-issued DOI via DataCite 
 

 Submission history

 From: Mathieu Blondel [view email]
 [v1] Sun, 11 Feb 2018 01:36:06 UTC (1,168 KB)
 [v2] Tue, 20 Feb 2018 06:01:45 UTC (1,168 KB)
 Full-text links:

 Download:

 * PDF
 * Other formats
 (license)
 Current browse context:
 stat.ML
 < prev | next >
 new | recent | 1802
 Change to browse by:
 cs
 cs.LG
 stat

 References & Citations

 * NASA ADS
 * Google Scholar
 * Semantic Scholar
 a export bibtex citation Loading...

 Bibtex formatted citation

 ×
 loading...
 Data provided by:

 Bookmark

 Bibliographic Tools

 Bibliographic and Citation Tools

 Bibliographic Explorer Toggle
 Bibliographic Explorer (What is the Explorer?)
 Litmaps Toggle
 Litmaps (What is Litmaps?)
 scite.ai Toggle
 scite Smart Citations (What are Smart Citations?)
 Code, Data, Media

 Code, Data and Media Associated with this Article

 Links to Code Toggle
 Papers with Code (What is Papers with Code?)
 ScienceCast Toggle
 ScienceCast (What is ScienceCast?)
 Demos

 Demos

 Replicate Toggle
 Replicate (What is Replicate?)
 Spaces Toggle
 Hugging Face Spaces (What is Spaces?)
 Related Papers

 Recommenders and Search Tools

 Connected Papers Toggle
 Connected Papers (What is Connected Papers?)
 Core recommender toggle
 CORE Recommender (What is CORE?)
 About arXivLabs

 arXivLabs: experimental projects with community collaborators

 arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

 Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

 Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.

 Which authors of this paper are endorsers? | Disable MathJax (What is MathJax?)
 * About
 * Help
 * Click here to contact arXiv Contact
 * Click here to subscribe Subscribe
 * Copyright
 * Privacy Policy
 * Web Accessibility Assistance

 * arXiv Operational Status
 Get status notifications via email or slack

Differentiable Dynamic Programming for Structured Prediction and Attention
February 21, 2018
Arthur Mensch arthur.mensch@m4x.org 
NTT Communication Science Laboratories Kyoto
Inria Parietal Saclay
France, Japan
Mathieu Blondel mathieu@mblondel.org 
NTT Communication Science Laboratories Kyoto
Inria Parietal Saclay
France, Japan
Differentiable Dynamic Programming for Structured Prediction and Attention
February 21, 2018Index words Dynamic programmingsmoothingstructured predictionattentionViterbiDTW
Dynamic programming (DP) solves a variety of structured combinatorial problems by iteratively breaking them down into smaller subproblems. In spite of their versatility, DP algorithms are usually non-differentiable, which hampers their use as a layer in neural networks trained by backpropagation. To address this issue, we propose to smooth the max operator in the dynamic programming recursion, using a strongly convex regularizer. This allows to relax both the optimal value and solution of the original combinatorial problem, and turns a broad class of DP algorithms into differentiable operators. Theoretically, we provide a new probabilistic perspective on backpropagating through these DP operators, and relate them to inference in graphical models. We derive two particular instantiations of our framework, a smoothed Viterbi algorithm for sequence prediction and a smoothed DTW algorithm for time-series alignment. We showcase these instantiations on two structured prediction tasks and on structured and sparse attention for neural machine translation.
Introduction
Modern neural networks are composed of multiple layers of nested functions. Although layers are usually constituted of elementary linear algebraic operations and simple non-linearities, there is a growing need for layers that output the value or the solution of an optimization problem. This can be used to design loss functions that capture relevant regularities in the input (Lample et al., 2016;Cuturi & Blondel, 2017) or to create layers that impose prior structure on the output (Kim et al., 2017;Amos & Kolter, 2017;Niculae & Blondel, 2017;Djolonga & Krause, 2017).
Among these works, several involve a convex optimization problem (Amos & Kolter, 2017;Niculae & Blondel, 2017;Djolonga & Krause, 2017); others solve certain combinatorial optimization problems by dynamic programming (Lample et al., 2016;Kim et al., 2017;Cuturi & Blondel, 2017). However, because dynamic programs (Bellman, 1952) are usually non-differentiable, virtually all these works resort to the formalism Figure 1. ToDo tions of our framework, a smoothed Viterbi algorithm for sequence prediction and a smoothed DTW algorithm for supervised time-series alignment ( §4). We showcase these two instantiations on structured prediction tasks ( §5) and on structured attention for neural machine translation (6).
DTW H (✓) = 7.49
Notation. We denote scalars, vectors and matrices using lower-case, bold lower-case and bold upper-case letters, e.g., y, y and Y , respectively. Given a matrix Y , we denote its elements by y i,j and its rows by y i . We denote the Frobenius inner product between two matrices A and B by hA, Bi , P i,j a i,j b i,j . We write the (D 1)-probability simplex by 4 D , { 2 R D + : k k 1 = 1}. We denote conv(Y) , { P Y 2Y Y Y : 2 4 |Y| } the convex hull of Y and [N ] the set {1, . . . , N}. We denote the Shannon entropy by H(q) , P i q i log q i .
Smoothed max operators
Smoothed max operators (Nesterov, 2005;Niculae & Blondel, 2017), as their name indicates, are smooth approximations of max operators. These operators will serve as a powerful and generic abstraction to define differentiable dynamic programming layers in §3. Formally, let ⌦ : R D ! R be a strongly convex regularizer and x 2 R D . Then, we define the max operator smoothed by ⌦ as:
max ⌦ (x) , max q24 D hq, xi ⌦(q).(1)
In other words, max ⌦ is the convex conjugate of ⌦, restricted to the simplex. By the duality between strong convexity and smoothness, max ⌦ is smooth (differentiable everywhere and with Lipschitz continuous gradient). Since the argument that achieves the maximum in (1) is unique, as per Danskin's theorem, it equals the gradient of max ⌦ : rmax ⌦ (x) = argmax q24 D hq, xi ⌦(q).
By Rademacher's theorem, rmax ⌦ (x) is differentiable almost everywhere. We will denote Clarke's generalized Jacobian of rmax ⌦ (x), or equivalently the generalized Hessian of max ⌦ (x), by r 2 max ⌦ (x). Next, we state several properties that will be useful throughout this paper.
Lemma 1. Properties of max ⌦ o Let x = (x 1 , . . . , x D ) > 2 R D . 1. Boundedness: If ⌦ is lowe upper-bounded by U ⌦,D on max(x) U ⌦,D  max ⌦ (x 2. Distributivity of + over max max ⌦ (x + c1) = max ⌦ (x 3. Commutativity: If ⌦(P q) = mutation matrix, then max ⌦ 4. Non-decreasingness in each max ⌦ (x)  max ⌦ (y) 8x 5. Insensitivity to 1: x j = Proofs are given in §A.1. In pa whenever ⌦(q) = P D i=1 !(q i ), focus in this paper on two speci tropy and squared`2 norm. For th above are satisfied and we can d sions for max ⌦ , its gradient and When using negentropy, max ⌦ b This operator satisfies associativ makes it natural to use in dynam squared`2 regularization, as ob tudillo (2016), the gradient rm prove useful to impose sparsity in
Differentiable DP layer
Dynamic programing (DP) is a combinatorial optimization prob ing problems on smaller sets. W egory of algorithms in a broad s max operators to define different
Dynamic programming on
The archetype case of dynamic ery problem that it solves can be tion of the highest-scoring path b end node, on a weighted directed therefore introduce our formalism and give concrete examples in §4 Formally, let G = (V, E) be a edges E. We write N = |V| Without loss of generality, we n logical order, from 1 (start) to N Node 1 is the only node without only node without children. Ever a parent node j to a child node i gather the edge weights in a mat ting ✓ i,j = 1 if (i, j) / 2 E an Figure 1. ToDo tions of our framework, a smoothed Viterbi algorith sequence prediction and a smoothed DTW algorith supervised time-series alignment ( §4). We showcase two instantiations on structured prediction tasks ( §5 on structured attention for neural machine translation DTW k·k 2 (✓) = 9.61 Notation. We denote scalars, vectors and matrices lower-case, bold lower-case and bold upper-case l e.g., y, y and Y , respectively. Given a matrix Y , w note its elements by y i,j and its rows by y i . We deno Frobenius inner product between two matrices A and hA, Bi , P i,j a i,j b i,j . We write the (D 1)-proba simplex by 4 D , { 2 R D + : k k 1 = 1}. We d conv(Y) , { P Y 2Y Y Y : 2 4 |Y| } the conve of Y and [N ] the set {1, . . . , N}. We denote the Sh entropy by H(q) , P i q i log q i .
Smoothed max operators
Smoothed max operators (Nesterov, 2005;Niculae & del, 2017), as their name indicates, are smooth ap imations of max operators. These operators will as a powerful and generic abstraction to define dif tiable dynamic programming layers in §3. Formal ⌦ : R D ! R be a strongly convex regularizer and x 2 Then, we define the max operator smoothed by ⌦ as:
max ⌦ (x) , max q24 D hq, xi ⌦(q).
In other words, max ⌦ is the convex conjugate of ⌦ stricted to the simplex. By the duality between strong vexity and smoothness, max ⌦ is smooth (differentiab erywhere and with Lipschitz continuous gradient). the argument that achieves the maximum in (1) is u as per Danskin's theorem, it equals the gradient of m rmax ⌦ (x) = argmax q24 D hq, xi ⌦(q).
By Rademacher's theorem, rmax ⌦ (x) is differentia most everywhere. We will denote Clarke's generaliz cobian of rmax ⌦ (x), or equivalently the generalized sian of max ⌦ (x), by r 2 max ⌦ (x). Next, we state s properties that will be useful throughout this paper. Figure 1: DTW Ω (θ) is an instantiation of the proposed smoothed dynamic programming operator, DP Ω (θ), to the dynamic time warping (DTW) computational graph. In this picture, θ is the squared Euclidean distance matrix between the observations of two time-series. The gradient ∇DTW Ω (θ) is equal to the expected alignment under a certain random walk characterized in §3.3 and is a sound continuous relaxation to the hard DTW alignment between the two time-series (here depicted with a yellow path). Unlike negentropy regularization (left), 2 2 regularization leads to exactly sparse alignments (right). Our framework allows to backpropagate through both DTW Ω (θ) and ∇DTW Ω (θ), which makes it possible to learn the distance matrix θ end-to-end. of conditional random fields (CRFs) (Lafferty et al., 2001), which can be seen as changing the semiring used by the dynamic program -replacing all values by their exponentials and all (max, +) operations with (+, ×) operations (Verdu & Poor, 1987). While this modification smoothes the dynamic program, it looses the sparsity of solutions, since hard assignments become soft ones. Moreover, a general understanding of how to relax and differentiate dynamic programs is lacking. In this work, we propose to do so by leveraging smoothing (Moreau, 1965;Nesterov, 2005) and backpropagation (Linnainmaa, 1970). We make the following contributions.
1) We present a unified framework for turning a broad class of dynamic programs (DP) into differentiable operators. Unlike existing works, we propose to change the semiring to use (max Ω , +) operations, where max Ω is a max operator smoothed with a strongly convex regularizer Ω ( §2).
2) We show that the resulting DP operators, that we call DP Ω , are smoothed relaxations of the original DP algorithm and satisfy several key properties, chief among them convexity. In addition, we show that their gradient, ∇DP Ω , is equal to the expected trajectory of a certain random walk and can be used as a sound relaxation to the original dynamic program's solution. Using negative entropy for Ω recovers existing CRFbased works from a different perspective -we provide new arguments as to why this Ω is a good choice. On the other hand, using squared 2 norm for Ω leads to new algorithms whose expected solution is sparse. We derive a clean and efficient method to backpropagate gradients, both through DP Ω and ∇DP Ω . This allows us to define differentiable DP layers that can be incorporated in neural networks trained end-to-end ( §3).
3) We illustrate how to to derive two particular instantiations of our framework, a smoothed Viterbi algorithm for sequence prediction and a smoothed DTW algorithm for supervised time-series alignment ( §4). The latter is illustrated in Figure 1. Finally, we showcase these two instantiations on structured prediction tasks ( §5) and on structured attention for neural machine translation ( §6).
Notation. We denote scalars, vectors and matrices using lower-case, bold lower-case and bold upper-case letters, e.g., y, y and Y . We denote the elements of Y by y i,j and its rows by y i . We denote the Frobenius inner product between A and B by A, B i,j a i,j b i,j . We denote the (D − 1)-probability simplex by . . . , N } and supp(x) {j ∈ [D] : x j = 0} the support of x ∈ R D . We denote the Shannon entropy by H(q) i q i log q i .
D {λ ∈ R D + : λ 1 = 1}. We write conv(Y) { Y ∈Y λ Y Y : λ ∈ |Y| } the convex hull of Y, [N ] the set {1,
We will release an optimized modular PyTorch implementation for reproduction and reuse. 
Smoothed max operators
In this section, we introduce smoothed max operators (Nesterov, 2005;Beck & Teboulle, 2012;Niculae & Blondel, 2017), that will serve as a powerful and generic abstraction to define differentiable dynamic programs in §3. Formally, let Ω : R D → R be a strongly convex regularizer on D and let x ∈ R D . We define the max operator smoothed by Ω as:
max Ω (x) max q∈ D q, x − Ω(q).(1)
In other words, max Ω is the convex conjugate of Ω, restricted to the simplex. From the duality between strong convexity and smoothness, max Ω is smooth: differentiable everywhere and with Lipschitz continuous gradient. Since the argument that achieves the maximum in (1) is unique, from Danskin's theorem (1966), it is equal to the gradient of max Ω :
∇max
Ω (x) = argmax q∈ D q, x − Ω(q).
The gradient is differentiable almost everywhere for any strongly-convex Ω (everywhere for negentropy). Next, we state properties that will be useful throughout this paper.
Lemma 1. Properties of max Ω operators Let x = (x 1 , . . . , x D ) ∈ R D .
Boundedness:
If Ω is lower-bounded by L Ω,D and upper-bounded by U Ω,D on the simplex D , then
max(x) − U Ω,D ≤ max Ω (x) ≤ max(x) − L Ω,D . 2. Distributivity of + over max Ω : max Ω (x + c1) = max Ω (x) + c ∀c ∈ R.
3. Commutativity: If Ω(P q) = Ω(q), where P is a permutation matrix, then max Ω (P x) = max Ω (x).
4. Non-decreasingness in each coordinate: max Ω (x) ≤ max Ω (y) ∀x ≤ y.
Insensitivity to
−∞: x j = −∞ ⇒ ∇max Ω (x) j = 0.
Proofs are given in §A.1. In particular, property 3 holds whenever Ω(q) = D i=1 ω(q i ), for some function ω. We focus in this paper on two specific regularizers Ω: the negentropy −H and the squared 2 norm. For these choices, all properties above are satisfied and we can derive closed-form expressions for max Ω , its gradient and its Hessian -see §B.1. When using negentropy, max Ω becomes the log-sum-exp and ∇max Ω the softmax. The former satisfies associativity, which as we shall see, makes it natural to use in dynamic programming. With the squared 2 regularization, as observed by Martins & Astudillo (2016); Niculae & Blondel (2017), the gradient ∇max Ω is sparse. This will prove useful to enforce sparsity in the models we study.
Differentiable DP layers
Dynamic programming (DP) is a generic way of solving combinatorial optimization problems by recursively solving problems on smaller sets. We first introduce this category of algorithms in a broad setting, then use smoothed max operators to define differentiable DP layers.
Dynamic programming on a DAG
Every problem solved by dynamic programming reduces to finding the highest-scoring path between a start node and an end node, on a weighted directed acyclic graph (DAG). We therefore introduce our formalism on this generic problem, and give concrete examples in §4.
Formally, let G = (V, E) be a DAG, with nodes V and edges E. We write N = |V| ≥ 2 the number of nodes. Without loss of generality, we number the nodes in topological order, from 1 (start) to N (end), and thus Differentiable Dynamic Programming for Structured Prediction and Attention 4 V = [N ]. Node 1 is the only node without parents, and node N the only node without children. Every directed edge (i, j) from a parent node j to a child node i has a weight θ i,j ∈ R. We gather the edge weights in a matrix θ ∈ Θ ⊆ R N ×N , setting θ i,j = −∞ if (i, j) / ∈ E and θ 1,1 = 1. We consider the set Y of all paths in G from node 1 to node N . Any path Y ∈ Y can be represented as a N × N binary matrix, with y i,j = 1 if the path goes through the edge (i, j) and y i,j = 0 otherwise. In the sequel, paths will have a one-to-one correspondence with discrete structures such as sequences or alignments. Using this representation, Y , θ corresponds to the cumulated sum of edge weights, along the path Y . The computation of the highest score among all paths amounts to solving the combinatorial problem
LP(θ) max Y ∈Y Y , θ ∈ R.
(2)
Although the size of Y is in general exponential in N , LP(θ) can be computed in one topologically-ordered pass over G using dynamic programming. We let P i be the set of parent nodes of node i in graph G and define recursively
v 1 (θ) 0 ∀ i ∈ [2, . . . , N ] : v i (θ) max j∈Pi θ i,j + v j (θ).(3)
This algorithm outputs DP(θ) v N (θ). We now show that this is precisely the highest score among all paths.
Proposition 1. Optimality of dynamic programming ∀θ ∈ Θ : DP(θ) = LP(θ)
The optimality of the recursion (3) is a well-known result (Bellman, 1952). We prove it again with our formalism in §A.2, since it exhibits the two key properties that the max operator must satisfy to guarantee optimality: distributivity of + over it and associativity. The cost of computing DP Ω (θ) is O(|E|), which is exponentially better than O(|Y|).
In many applications, we will often rather be interested in the argument that achieves the maximum, i.e., one of the highest-scoring paths
Y (θ) ∈ argmax Y ∈Y Y , θ .(4)
This argument can be computed by backtracking, that we now relate to computing subgradients of LP(θ).
Linear program, lack of differentiality. Unfortunately, LP(θ) is not differentiable everywhere. To see why this is the case, notice that (2) can be rewritten as a linear program over the convex polytope conv(Y):
LP(θ) = max Y ∈conv(Y) Y , θ .
From the generalized Danskin theorem (Bertsekas, 1971),
Y (θ) ∈ argmax Y ∈conv(Y) Y , θ ⊆ ∂LP(θ),
where ∂ denotes the subdifferential of LP(θ), i.e., the set of subgradients. When Y (θ) is unique, ∂LP(θ) is a singleton and Y is equal to the gradient of LP(θ), that we write ∇LP(θ). Unfortunately, Y (θ) is not always unique, meaning that LP(θ) is not differentiable everywhere. This hinders optimization as we can only train models involving LP(θ) with subgradient methods. Worse, Y (θ), a function from Θ to Y, is discontinuous and has null or undefined derivatives. It is thus impossible to use it in a model trained by gradient descent.
Smoothed max layers
To address the lack of differentiability of dynamic programming, we introduce the operator max Ω , presented in §2, and consider two approaches.
Smoothing the linear program. Let us define the Ω-smoothed maximum of a function f : Y → R over a finite set Y using the following shorthand notation:
max Ω Y ∈Y f (Y ) max Ω ((f (Y )) Y ∈Y ).
A natural way to circumvent the lack of differentiability of LP(θ) is then to replace the global max operator by max Ω :
LP Ω (θ) max Ω Y ∈Y Y , θ ∈ R.(5)
From §2, LP Ω (θ) is convex and, as long as Ω is strongly convex, differentiable everywhere. In addition, ∇LP Ω (θ) is Lipschitz continuous and thus differentiable almost everywhere. Unfortunately, solving (5) for general strongly convex Ω is intractable when Y has an exponential size.
Smoothing the dynamic program. As a tractable alternative, we propose an algorithmic smoothing. Namely, we replace max by max Ω locally within the DP recursion. Omitting the dependence on Ω, this defines a smoothed recursion over the new sequence (v i (θ))
N i=1 : v 1 (θ) 0 ∀i ∈ [2, . . . , N ] : v i (θ) max Ω j∈Pi θ i,j + v j (θ).(6)
The new algorithm outputs DP Ω (θ) v N (θ), the smoothed highest score. Smoothing the max operator locally brings the same benefit as before -DP Ω (θ) is smooth and ∇DP Ω (θ) is differentiable almost everywhere. However, computing DP Ω (θ) is now always tractable, since it simply requires to evaluate
(v i (θ)) N i=1
in topological order, as in the original recursion (3). Although LP Ω (θ) and DP Ω (θ) are generally different (in fact, LP Ω (θ) ≥ DP Ω (θ) for all θ ∈ Θ), we now show that DP Ω (θ) is a sensible approximation of LP(θ) in several respects.
Proposition 2. Properties of DP Ω 1. DP Ω (θ) is convex 2. LP(θ) − DP Ω (θ) is bounded above and below:
(N − 1)L Ω,N ≤ LP(θ) − DP Ω (θ) ≤ (N − 1)U Ω,N ,
where L Ω,N and U Ω,N are defined in Lemma 1.
3.
When Ω is separable, DP Ω (θ) = LP Ω (θ) if and only if Ω = −γH, where γ ≥ 0.
Proofs are given in §A.3. The first claim can be surprising due to the recursive definition of DP Ω (θ). The second claim implies that DP γΩ (θ) converges to LP(θ) when the regularization vanishes: DP γΩ (θ) → γ→0 LP(θ); LP γΩ (θ) also satisfies this property. The "if" direction of the third claim follows by showing that max −γH satisfies associativity. This recovers known results in the framework of message passing algorithms for probabilistic graphical models (e.g., Wainwright & Jordan, 2008, Section 4.1.3), with a more algebraic point of view. The key role that the distributive and associative properties play into breaking down large problems into smaller ones has long been noted (Verdu & Poor, 1987;Aji & McEliece, 2000). However, the "and only if" part of the claim is new to our knowledge. Its proof shows that max −γH is the only max Ω satisfying associativity, exhibiting a functional equation from information theory (Horibe, 1988). While this provides an argument in favor of entropic regularization, 2 2 regularization has different benefits in terms of sparsity of the solutions.
Relaxed argmax layers
It is easy to check that ∇LP Ω (θ) belongs to conv(Y) and can be interpreted as an expected path under some distribution induced by ∇max Ω , over all possible Y ∈ Y -see §A.4 for details. This makes ∇LP Ω (θ) interpretable as a continuous relaxation of the highest-scoring path Y (θ) defined in (4). However, like LP Ω (θ), computing ∇LP Ω (θ) is intractable in the general case. Fortunately, we now show that ∇DP Ω (θ) is always easily computable by backpropagation and enjoys similar properties.
Differentiable Dynamic Programming for Structured Prediction and Attention 6
Computing ∇DP Ω (θ). Computing ∇DP Ω (θ) can be broken down into two steps. First, we compute and record the local gradients alongside the recursive step (6):
∀ i ∈ [N ] : q i (θ) ∇max Ω (θ i + v(θ)) ∈ N , where v(θ) (v 1 (θ), . . . , v N (θ)). Since we assume that θ i,j = −∞ if (i, j) ∈ E, we have supp(q i (θ)) = P i .
This ensures that, similarly to v i (θ), q i (θ) exclusively depends on (v j (θ)) j∈Pi . Let C j be the children of node j ∈ [N ]. A straighforward application of backpropagation (cf. §A.5) yields a recursion run in reverse-topological order, starting from node j = N − 1 down to j = 1:
∀ i ∈ C j : e i,j ←ē i q i,j thenē j ← i∈Cj e i,j , whereē N ← 1 and e i,j ← 0 for (i, j) / ∈ E. The final output is ∇DP Ω (θ) = E.
Assuming max Ω can be computed in linear time, the total cost is O(|E|), the same as DP(θ). Pseudo-code is summarized in §A.5.
Associated path distribution. The backpropagation we derived has a probabilistic interpretation. Indeed, Q(θ) ∈ R N ×N can be interpreted as a transition matrix: it defines a random walk on the graph G, i.e., a finite Markov chain with states V and transition probabilities supported by E. The random walk starts from node N and, when at node i, hops to node j ∈ P i with probability q i,j . It always ends at node 1, which is absorbing. The walk follows the path Y ∈ Y with a probability p θ,Ω (Y ), which is simply the product of the q i,j of visited edges. Thus, Q(θ) defines a path distribution p θ,Ω . Our next proposition shows that ∇DP Ω (Y ) ∈ conv(Y) and is equal to the expected path E θ,Ω [Y ] under that distribution.
Proposition 3. ∇DP Ω (θ) as an expected path ∀ θ ∈ Θ : ∇DP Ω (θ) = E θ,Ω [Y ] = E ∈ conv(Y).
Proof is provided in §A.5. Moreover, ∇DP Ω (θ) is a principled relaxation of the highest-scoring path Y (θ), in the sense that it converges to a subgradient of LP(θ) as the regularization vanishes:
∀ θ ∈ Θ : ∇DP γΩ (θ) − −− → γ→0 Y (θ) ∈ ∂LP(θ).
When Ω = −γH, the distributions underpinning LP Ω (θ) and DP Ω (θ) coincide and reduce to the Gibbs distribution p θ,Ω (Y ) ∝ exp( θ, Y /γ). The value LP Ω (θ) = DP Ω (θ) is then equal to the log partition.
When Ω = γ · 2 , some transitions between nodes have zero probability and hence some paths have zero probability under the distribution p θ,Ω . Thus, ∇DP Ω (θ) is typically sparse -this will prove interesting to introspect the various models we consider (typically, the smaller γ, the sparser ∇DP Ω (θ)).
Multiplication with the Hessian ∇ 2 DP Ω (θ)Z
Using ∇DP Ω (θ) as a layer involves backpropagating through ∇DP Ω (θ). This requires computing the Jacobian ∇∇DP Ω (θ) or in other words the Hessian ∇ 2 DP Ω (θ), a linear map from R N ×N to R N ×N . Fortunately, a practical implementation of backpropagation only requires to apply that map to a provided matrix Z ∈ R N ×N , i. e., ∇ 2 DP Ω (θ)Z. We therefore focus on that term. Recall that the directional derivative of DP Ω at θ along Z can be computed by ∇DP Ω (θ), Z ∈ R. Our key technique, which is also at the heart of Pearlmutter's method (1994), is to observe that ∇ 2 DP Ω (θ)Z is the gradient of the directional derivative at θ along Z. Namely,
∇ 2 DP Ω (θ)Z = ∇ ∇DP Ω (θ), Z .
We therefore break down the computation of ∇ 2 DP Ω (θ)Z into two steps. First, we compute the directional derivative ∇DP Ω (θ), Z using the chain rule. It can be computed in one topologically-ordered pass over G. Similarly to the gradient computation, we record multiplications with the (generalized) local Hessian H i (θ) ∇ 2 max Ω (θ i + v(θ)) along the way. Second, we compute the gradient of the directional derivative using backpropagation. It yields a recursion for computing ∇ 2 DP Ω (θ)Z in reverse topological-order over G. The complete derivation and the pseudo-code are given in §A.7. The total computational cost is O(|E|), as for the gradient computation.
Performance. Using autodiff frameworks such as PyTorch (Paszke et al., 2017), it is possible to only implement DP Ω (θ) and rely on tape-based gradient computation to obtain ∇DP Ω (θ). Provided that we tape the backward pass as well, we can then backpropagate again through ∇DP Ω (θ) to obtain ∇ 2 DP Ω (θ)Z. In practice, however, implementing backpropagation without resorting to autodiff software is crucial, since the DAG structure can be directly harcoded in concrete cases -see §4. Moreover, our "reverse-over-forward" approach to compute the Hessian product (backpropagating over the directional derivative computation) yields a simpler computation graph than the "reverse-over-reverse" approach (backpropagation over taped backpropagation). In experiments, our approach is up to 50× faster than vanilla PyTorch on the Viterbi DAG. Note that our algorithms are readily vectorizable and can efficiently handle mini-batches with varying input lengths.
Summary. We have proposed DP Ω (θ), a smooth, convex and tractable relaxation to the value of LP(θ). We have also shown that ∇DP Ω (θ) belongs to conv(Y) and is therefore a sound relaxation to solutions of LP(θ). To conclude this section, we formally define our proposed two layers.
Definition 1. Differentiable dynamic programming layers
Value layer: DP Ω (θ) ∈ R Gradient layer: ∇DP Ω (θ) ∈ conv(Y)
Examples of computational graphs
We now illustrate two instantiations of our framework for specific computational graphs.
Sequence prediction
We demonstrate in this section how to instantiate DP Ω to the computational graph of the Viterbi algorithm (Viterbi, 1967;Rabiner, 1990), one of the most famous instances of DP algorithm. We call the resulting operator Vit Ω . We wish to tag a sequence X = (x 1 , . . . , x T ) of vectors in R D (e.g., word representations) with the most probable output sequence (e.g., entity tags) y = (y 1 , . . . , y T ) ∈ [S] T . This problem can be cast as finding the highest-scoring path on a treillis G. While y can always be represented as a sparse N × N binary matrix, it is convenient to represent it instead as a T × S × S binary tensor Y , such that y t,i,j = 1 if y transitions from node j to node i on time t, and 0 otherwise -we set y 0 = 1. The potentials can similarly be organized as a T × S × S real tensor, such that θ t,i,j = φ t (x t , i, j). Traditionally, the potential functions φ t were human-engineered (Sutton et al., 2012, §2.5). In recent works and in this paper, they are learned end-to-end (Lample et al., 2016).
Using the above binary tensor representation, the inner product Y , θ is equal to T t=1 φ t (x t , y t , y t−1 ), y's cumulated score. This is illustrated in Figure 2 on the task of part-of-speech tagging. The bold arrows indicate one possible output sequence y, i.e., one possible path in G.
When Ω = −H, we recover linear-chain conditional random fields (CRFs) (Lafferty et al., 2001) and the probability of y (Y in tensor representation) given X is
p θ,−H (y|X)∝ exp( Y , θ )= exp T t=1 φ t (x t , y t , y t−1 ) . From Prop. 3, the gradient ∇Vit −H (θ) = E ∈ R T ×S×S is such that e t,i,j = p θ,−H (y t = i, y t−1 = j|X). The marginal probability of state i at time t is simply p θ,−H (y t = i|X) = S j=1 e t,i,j .
Using a different Ω simply changes the distribution over state transitions. When Ω = · 2 , the marginal probabilities are typically sparse. Pseudo-code for Vit Ω (θ), as well as gradient and Hessian-product computations, is provided in §B.2. The case Ω = · 2 is new to our knowledge.
When Ω = −H, the marginal probabilities are traditionally computed using the forward-backward algorithm (Baum & Petrie, 1966). In contrast, we compute ∇Vit −H (θ) using backpropagation while efficiently Differentiable Dynamic Programming for Structured Prediction and Attention maintaining the marginalization. An advantage of our approach is that all operations are numerically stable.
The relation between forward-backward and backpropagation has been noted before (e.g., Eisner (2016)). However, the analysis is led using (+, ×) operations, instead of (max Ω , +) as we do. This Viterbi instantiation can immediately be generalized to graphical models with a tree structure, and to approximate inference in general graphical models, since unrolled loopy belief propagation (Pearl, 1988) yields a dynamic program.
Time-series alignment
We now demonstrate how to instantiate DP Ω to the computational graph of dynamic time warping (DTW) (Sakoe & Chiba, 1978), whose goal is to seek the minimal cost alignment between two time-series. We call the resulting operator DTW Ω . Formally, let N A and N B be the lengths of two time-series, A and B. Let a i and b j be the i th and j th observations of A and B, respectively. Since edge weights only depend on child nodes, it is convenient to rearrange Y and θ as N A × N B matrices. Namely, we represent an alignment Y as a N A × N B binary matrix, such that y i,j = 1 if a i is aligned with b j , and 0 otherwise. Likewise, we represent θ as a N A × N B matrix. A classical example is θ i,j = d(a i , b j ), for some differentiable discrepancy measure d. We write Y the set of all monotonic alignment matrices, such that the path that connects the upper-left (1, 1) matrix entry to the lower-right (N A , N B ) one uses only ↓, →, moves. The DAG associated with Y is illustrated in Figure 3 with N A = 4 and N B = 3 below.
start 1,1 Again, the bold arrows indicate one possible path Y ∈ Y from start to end in the DAG, and correspond to one possible alignment. Using this representation, the cost of an alignment (cumulated cost along the path) is conveniently computed by Y , θ . The value DTW Ω (θ) can be used to define a loss between alignments or between time-series. Following Proposition 3, ∇DTW Ω (θ) = E ∈ R N A ×N B can be understood as a soft Differentiable Dynamic Programming for Structured Prediction and Attention 9 alignment matrix. This matrix is sparse when Ω = · 2 , as illustrated in Figure 1 (right).
θ 1,1 2,1 1,2 3,1 1,3 2,2 θ 2,2 2,3 θ 2,3 3,2 1,4 2,4 3,3 θ 3,3 3,4 θ 3,4 end Y , θ = θ 1,1 + θ 2,2 + θ 2,3 + θ 3,3 + θ 3,4
Pseudo-code to compute DTW Ω (θ) as well as its gradient and its Hessian products are provided in §B.3.
When Ω = −H, DTW Ω (θ) is a conditional random field known as soft-DTW, and the probability p θ,Ω (Y |A, B) is a Gibbs distribution similar to §4.1 (Cuturi & Blondel, 2017). However, the case Ω = · 2 and the computation of ∇ 2 DTW Ω (θ)Z are new and allow new applications.
Differentiable structured prediction
We now apply the proposed layers, DP Ω (θ) and ∇DP Ω (θ), to structured prediction (Bakır et al., 2007), whose goal is to predict a structured output Y ∈ Y associated with a structured input X ∈ X . We define old and new structured losses, and demonstrate them on two structured prediction tasks: named entity recognition and time-series alignment.
Structured loss functions
Throughout this section, we assume that the potentials θ ∈ Θ have already been computed using a function from X to Θ and let C : Y × Y → R + be a cost function between the ground-truth output Y true and the predicted output Y .
Convex losses. Because C is typically non-convex, the cost-augmented structured hinge loss (Tsochantaridis et al., 2005) is often used instead for linear models
C (Y true ; θ) max Y ∈Y C(Y true , Y ) + Y , θ − Y true , θ .(7)
This is a convex upper-bound on C(Y true , Y (θ)), where Y (θ) is defined in (4). To make the costaugmented decoding tractable, it is usually assumed that C(Y true , Y ) is linear in Y , i. e., it can be written as C Ytrue , Y for some matrix C Ytrue . We can then rewrite (7) using our notation as
C (Y true ; θ) = LP(θ + C Ytrue ) − Y true , θ .
However, this loss function is non-differentiable. We therefore propose to relax LP by substituting it with DP Ω :
C,Ω (Y true ; θ) DP Ω (θ + C Ytrue ) − Y true , θ .
Losses in this class are convex, smooth, tractable for any Ω, and by Proposition 2 property 2 a sensible approximation of C . In addition, they only require to backpropagate through DP Ω (θ) at training time.
It is easy to check that we recover the structured hinge loss with C,0 (Tsochantaridis et al., 2005) and the CRF loss with 0,−H (Lafferty et al., 2001). The last one has been used on top of LSTMs in several recent works (Lample et al., 2016;Ma & Hovy, 2016). Minimizing 0,−H (θ) is equivalent to maximizing the likelihood p θ,−H (Y true ). However, minimizing 0, · 2 is not equivalent to maximizing p θ, · 2 (Y true ). In fact, the former is convex while the latter is not.
Non-convex losses. A direct approach that uses the output distribution p θ,Ω consists in minimizing the risk y∈Y p θ,−H (Y )C(Y true , Y ). As shown by Stoyanov & Eisner (2012), this can be achieved by backpropagating through the minimum risk decoder. However, the risk is usually non-differentiable, piecewise constant (Smith & Eisner, 2006) and several smoothing heuristics are necessary to make the method work (Stoyanov & Eisner, 2012).
Another principled approach is to consider a differentiable approximation ∆ : Y × conv(Y) → R + of the cost C. We can then relax C(Y true , Y (θ)) by ∆(Y true , ∇DP Ω (θ)). Unlike minimum risk training, this approach is differentiable everywhere when Ω = −H. Both approaches require to backpropagate through ∇DP Ω (θ), which is only roughly twice as costly as backpropagating through DP Ω (θ) using the approach outlined in §3.4. 
Named entity recognition
Let X = (x 1 , · · · , x T ) be an input sentence, where each word x t is represented by a vector in R D , computed using a neural recurrent architecture trained end-to-end. We wish to tag each word with named entities, i.e., identify blocks of words that correspond to names, locations, dates, etc. We use the specialized operator Vit Ω described in §4.1. In our experiments, we define the elements of the potential tensor θ(X) ∈ R T ×S×S when t > 1 by
θ(X) t,i,j w i x t + b i + t i,j and θ(X) 1,i,j w i x t + b i , where (w i , b i ) ∈ R D × R
is the linear classifier associated with tag i and T ∈ R S×S is a transition matrix. We learn W , b and T along with the network producing X, and compare two losses:
Surrogate convex loss: 0,Ω (Y true ; θ),
Relaxed loss: ∆(Y true , ∇DP Ω (θ)),
where ∆(Y true , Y ) is the squared 2 distance when Ω = · 2 2 and the Kullback-Leibler divergence when Ω = −H, applied row-wise to the marginalization of Y true and Y .  Table 1, along with (Lample et al., 2016) results with different pretrained embeddings. With proper parameter selections, all losses perform within 1% F 1 -score of each other, although entropy-regularized losses perform slightly better on 3 /4 languages. However, the 2 2 -regularized losses yield sparse predictions, whereas entropy regularization always yields dense probability vectors. Qualitatively, this allows to identify ambiguous predictions more easily -this is illustrated in §C.1 with additional figures.
Supervised audio-to-score transcription
We use our framework to perform supervised audio-to-score alignment on the Bach 10 dataset (Duan & Pardo, 2011). The dataset consists of 10 music pieces with audio tracks, MIDI transcriptions, and annotated alignments between them. We transform the audio tracks into a sequence of audio frames using a feature extractor (see §C.2) to obtain a sequence A ∈ R N A ×D , while the associated score sequence is represented by B ∈ R N B ×K (each row b j is a one-hot vector corresponding to one key b j ). Each pair (A, B) is associated to an alignment Y true ∈ R N A ×N B . As described in §4.2, we need to define a discrepancy matrix θ ∈ R N A ×N B between the elements of the two sequences. We set the cost between an audio frame and a key to be the log-likelihood of this key given a multinomial linear classifier. For all i ∈ [N A ], we define
l i − log(softmax(W a i + c)) ∈ R K ,(8)
and ∀ j ∈ [N B ], θ i,j l i,bj , where (W , c) ∈ R D×K × R K are learned classifier parameters. We predict a soft alignment by Y = ∇DTW −H (θ). Following (Garreau et al., 2014), we define the relaxed loss
∆(Y true , Y ) L(Y − Y true ) 2 F ,
where L a the lower triangular matrix filled with 1. When Y ∈ Y is a true alignement matrix, ∆(Y true , Y ) is the area between the path of Y true and Y , which corresponds to the mean absolute deviation in the audio literature. When Y ∈ conv(Y), it is a convex relaxation of the area. At test time, once θ is learned, we use the non-regularized DTW algorithm to output a hard alignment Y (θ) ∈ Y.
Results. We perform a leave-one-out cross-validation of our model performance, learning the multinomial classifier on 9 pieces and assessing the quality of the alignment on the remaining piece. We report the mean absolute deviation on both train and test sets. A solid baseline consists in learning the multinomial classifier (W , c) beforehand, i.e., without end-to-end training. We then use this model to compute θ as in (8) and obtain Y (θ). As shown in Table 2, our end-to-end technique outperforms this baseline by a large margin. We also demonstrate in §C.2 that the alignments obtained by end-to-end training are visibly closer to the ground truth. End-to-end training thus allows to fine-tune the distance matrix θ for the alignment task at hand.
Structured and sparse attention
We show in this section how to apply our framework to neural sequence-to-sequence models augmented with an attention mechanism (Bahdanau et al., 2015). An encoder first produces a list of vectors X = (x 1 , . . . , x T ) representing the input sequence. A decoder is then used to greedily produce the corresponding output sequence. To simplify the notation, we focus on one time step of the decoding procedure. Given the decoder's current hidden state z and X as inputs, the role of the attention mechanism is to produce a distribution w ∈ T over X, for the current time step. This distribution is then typically used to produce a context vector c X w, that is in turn invoved in the computation of the output sequence's next element.
Structured attention layers. Kim et al. (2017) proposed a segmentation attention layer, which is capable of taking into account the transitions between elements of X. They use a linear-chain CRF to model the probability p θ,−H (y|X) of a sequence y = (y 1 , . . . , y T ), where each y t is either 1 ("pay attention") or 0. They then propose to use normalized marginal probabilities as attention weights: w t ∝ p θ,−H (y t = 1|X). They show how to backpropagate gradients through the forward-backward algorithm, which they use to compute the marginal probabilities.
Generalizing structured attention. We now show how to generalize segmentation layers to any Ω and how to backpropagate through them efficiently. Using the notation from §4.1, any y can be represented as a tensor Y ∈ {0, 1} T ×2×2 and the potentials as a tensor θ ∈ R T ×2×2 . As in (Kim et al., 2017), we define θ t,1,j x t M z + t 1,j and θ t,0,j t 0,j ,
where xM z is a learned bilinear form and T ∈ R 2×2 is a learned transition matrix. Following §4.1, the gradient ∇Vit Ω (θ) is equal to the expected matrix E ∈ R T ×2×2 and the marginals are obtained by marginalizing that matrix. Hence, we can set w t ∝ p θ,Ω (y t = 1|X) = e t,1,0 + e t,1,1 . Backpropagating through ∇Vit Ω (θ) can be carried out using our approach outlined in §3.4. This approach is not only more general, but also simpler and more robust to underflow problems than backpropagating through the forward-backward algorithm as done in (Kim et al., 2017).
Experiments. We demonstrate structured attention layers with an LSTM encoder and decoder to perform French to English translation using data from a 1 million sentence subset of the WMT14 FR-EN challenge.
We illustrate an example of attenion map obtained with negentropy and 2 2 regularizations in Figure 4. Nonzero elements are underlined with borders: 2 2 -regularized attention maps are sparse and more interpretable -this provides a structured alternative to sparsemax attention (Martins & Astudillo, 2016). Results were all within 0.8 point of BLEU score on the newstest2014 dataset. For French to English, standard softmax attention obtained 27.96, while entropy and 2 2 regularized attention obtained 27.96 and 27.19introducing structure and sparsity therefore provides enhanced interpretability with comparable peformance. We provide model details, full results and further visualizations in §C.3.
Conclusion
We proposed a theoretical framework for turning a broad class of dynamic programs into convex, differentiable and tractable operators, using the novel point of view of smoothed max operators. Our work sheds a new light on how to transform dynamic programs that predict hard assignments (e.g., the maximum a-posteriori estimator in a probabilistic graphical model or an alignment matrix between two time-series) into continuous and probabilistic ones. We provided a new argument in favor of negentropy regularization by showing that it is the only one to preserve associativity of the smoothed max operator. We showed that different regularizations induce different distributions over outputs and that 2 2 regularization has other benefits, in terms of sparsity of the expected outputs. Generally speaking, performing inference in a graphical model and backpropagating through it reduces to computing the first and second-order derivatives of a relaxed maximum-likelihood estimation -leveraging this observation yields elegant and efficient algorithms that are readily usable in deep learning frameworks, with various promising applications.
Appendix A Proofs and detailed derivations
This section contains the proofs of the propositions and lemmas presented in the main text. It also contains derivations of gradient, directional derivative and Hessian-product computations.
A.1 Proof of Lemma 1 (properties of max Ω ) Property 1 (boundedness). Let q and q Ω be the solutions of max q∈ D q x and max q∈ D q x−Ω(q), respectively. Then, we have
max Ω (x) = q Ω , x − Ω(q Ω ) ≥ q , x − Ω(q ) = max(x) − Ω(q ) and max(x) − Ω(q Ω ) ≥ q Ω , x − Ω(q Ω ) = max Ω (x).
Combining the two and using L Ω,D ≤ Ω(q) ≤ U Ω,D ∀q ∈ D , we obtain
max(x) − U Ω,D ≤ max(x) − Ω(q ) ≤ max Ω (x) ≤ max(x) − Ω(q Ω ) ≤ max(x) − L Ω,D .
When Ω(q) = i q i log q i , we have the tight inequality − log D ≤ Ω(q) ≤ 0 ∀q ∈ D and hence
max(x) ≤ max Ω (x) ≤ max(x) + log D.
When Ω(q) = 1 2 q 2 , we have the tight inequality 1 2D ≤ Ω(q) ≤ 1 2 ∀q ∈ D and hence
max(x) − 1 2 ≤ max Ω (x) ≤ max(x) − 1 2D .
Note that the difference U Ω,D − L Ω,D is equal to log D when Ω is the negative entropy and to D−1 2D ≤ 1 2 when Ω is the squared 2 norm. Since log D > 1 2 for all integers D ≥ 2, we get a better approximation of the max operator using squared 2 norm than using negative entropy, whenever D ≥ 2.
Property 2 (distributivity of + over max Ω ). This follows immediately from
max Ω (x + c1) = max q∈ D q, x + c1 − Ω(q) = max q∈ D
q, x − Ω(q) + c = max Ω (x) + c.
Using our shorthand notation, this simply becomes max
Ω Y ∈Y (f (Y ) + c) = max Ω Y ∈Y f (Y ) + c.
Property 3 (commutativity). Assume Ω(P q) = Ω(q) for all permutation matrices P . Let P −1 be the inverse permutation matrix associated with P . Then we have
max Ω (P x) = max q∈ D q, P x − Ω(q) = max q∈ D P −1 q, x − Ω(q) = max q∈ D q, x − Ω(P q) = max q∈ D q, x − Ω(q).
Property 4 (non-decreasingness in each coordinate). If x ≤ y, then for all q ∈ D , x, q − Ω(q) ≤ y, q − Ω(q), as all q coordinates are non-negative. Thus max Ω (x) ≤ max Ω (y).
Property 5 (insensitivity to −∞). Since max Ω (x) = max q∈ D q, x − Ω(q), if x j = − ∞, then q j = ∇max Ω (x) j = 0 is the only feasible solution for the j th coordinate.
A.2 Proof of Proposition 1 (optimality of DP recursion)
Let v i (θ) be the highest-score path up to node i ∈ [N ]. Let Y i be the set of paths y = (y 1 , . . . , y L ) starting from node 1 and reaching node i, that is y 1 = 1 and y L = i. Note that L may depend on y but we do not make this dependency explicit. Because nodes are sorted in topological order, we can compute
v i (θ) by v i (θ) = max y∈Yi L t=2 θ yt,yt−1 = max y∈Yi L−1 t=2 θ yt,yt−1 + θ y L ,y L−1 = max y∈Yi L−1 t=2 θ yt,yt−1 + θ i,y L−1 .
Recall that P i is the set of parent nodes of node i. From the associativity of the max operator,
v i (θ) = max j∈Pi max y∈Yi y L−1 =j L−1 t=2 θ yt,yt−1 + θ i,y L−1 = max j∈Pi max y∈Yi y L−1 =j L−1 t=2 θ yt,yt−1 + θ i,j .
From the distributivity of + over max, we obtain
v i (θ) = max j∈Pi   max y∈Yi y L−1 =j L−1 t=2 θ yt,yt−1   + θ i,j = max j∈Pi v j (θ) + θ i,j ,
where we used the fact that the inner max operations are independent of y L = i. This concludes the proof of the optimality of (3).
A.3 Proof of Proposition 2 (properties of DP Ω (θ))
We prove in this section the three main claims of Proposition 2. For the first two claims, we rewrite (3) and (6) using the following notations:
v 0 i (θ) max(u 0 i (θ)) and v Ω i (θ) max(u Ω i (θ)), where u 0 i (θ) (θ i,1 + v 0 1 (θ), . . . , θ i,i−1 + v 0 i−1 (θ), −∞, −∞, . . . , −∞) ∈ R N and u Ω i (θ) (θ i,1 + v Ω 1 (θ), . . . , θ i,i−1 + v Ω i−1 (θ), −∞ i , −∞, . . . , −∞) ∈ R N .
These definitions are indeed valid as per Lemma 1, property 5.
Proof of DP Ω (θ) convexity. Since v Ω 1 (θ) = 0, it is trivially convex. Assume that v Ω 2 (θ), . . . , v Ω i−1 (θ) are convex. Then, v Ω i (θ) is the composition of max Ω and u Ω i , a convex function and a function which outputs a vector whose each coordinate is convex in θ. By induction, since max Ω is non-decreasing per coordinate (cf . Lemma 1 property 4), v Ω i (θ) is convex (e.g., Boyd & Vandenberghe, 2004, §3.2.4).
Therefore v Ω i (θ) is convex for all i ∈ [N ] and DP Ω (θ) = v Ω N (θ) is convex. Proof of DP Ω (θ) bound. We clearly have v Ω 1 (θ) ≥ v 0 1 (θ). Assume that v Ω j (θ) ≥ v 0 j (θ) − (j − 1)U Ω,N for all j ∈ {2, . . . , i − 1}. That is, u Ω i (θ) ≥ u 0 i (θ) − (i − 2)U Ω,N 1, where 1 ∈ R N is the unit vector. Then, by induction, we have max Ω (u Ω i (θ)) ≥ max Ω (u 0 i (θ)) − (i − 2)U Ω,N ≥ max(u 0 i (θ)) − (i − 1)U Ω,N ,
where we used Lemma 1, properties 1, 2 and 4. Therefore v
Ω i (θ) ≥ v 0 i (θ) − (i − 1)U Ω,N for all i ∈ [N ]
and hence, DP Ω (θ) ≥ LP(θ) − (N − 1)U Ω,N . Using a similar reasoning we obtain v 0 i (θ) − (i − 1)L Ω,N ≥ v Ω i (θ) and therefore LP(θ) − (N − 1)L Ω,N ≥ DP Ω (θ). To summarize, we obtain
LP(θ) − (N − 1)L Ω,N ≥ DP Ω (θ) ≥ LP(θ) − (N − 1)U Ω,N ,
which concludes the proof. Note that using property 1 of Lemma 1, this immediately implies a bound involving LP Ω (θ) instead of LP(θ).
Proof that Ω = −γH ⇒ DP Ω (θ) = LP Ω (θ).
We first show that max Ω is associative.
Lemma 2. Associativity of max Ω when Ω = −γH
We have max Ω (max Ω (x), c) = max Ω (x, c) ∀x ∈ R D , c ∈ R.
Proof. We simply use the closed form of max Ω when Ω = −γH (cf. §B.1):
max Ω (max Ω (x), c) = γ log(exp(max Ω (x)/γ) + exp(c/γ)) = γ log exp log D i=1 exp(x i /γ) + exp(c/γ) = γ log D i=1 exp(x i /γ) + exp(c/γ) = max Ω (x, c),
and the lemma follows.
Using our shorthand notation, Lemma 2 can be used to write max Ω (y1,...,yi,...,y L )
f (y) = max Ω v max Ω (y1,...,v,...,y L ) f (y).
This is precisely the associative property that we used in the proof of Proposition 1. The second property that we used, the distributivity of + over max, holds for any max Ω , as per Lemma 1 property 2. Thus, the same proof as Proposition 1 is also valid when we substitute max with max Ω , when Ω = −γH, which yields LP Ω (θ) = DP Ω (θ).
Proof that Ω = −γH ⇐ DP Ω (θ) = LP Ω (θ).
Mirroring the previous proof, we first characterize the regularizations Ω for which max Ω is associative.
Lemma 3. Let Ω : D → R be a regularization function, i. e., dom Ω = D . Assume that there exist ω convex lower-semi-continuous defined on [0, 1] such that Ω
(q) = d i=1 ω(q i ). If max Ω (max Ω (x), c) = max Ω (x, c) ∀x ∈ R D , c ∈ R, then Ω(q) = −γ d i=1 q i log(q i ) for some γ ≥ 0.
Proof. We start by writing the associativity property for three elements. For all x 1 , x 2 , x 3 ∈ R,
max Ω (x 1 , x 2 , x 3 ) = max Ω max Ω (x 1 , x 2 ), x 3 ) = max q+q3=1 q,q3≥0 q max q1+q2=1 qi≥0 q 1 x 1 +q 2 x 2 − ω(q 1 ) − ω(q 2 ) + q 3 x 3 − ω(q 3 ) − ω(q) = max q1+q2+q3=1 qi≥0 q 1 x 1 + q 2 x 2 + q 3 x 3 − Φ(q 1 , q 2 , q 3 ), where Φ(q 1 , q 2 , q 3 ) (q 1 + q 2 ) ω q 1 q 1 + q 2 + ω q 2 q 1 + q 2 + ω(q 1 + q 2 ) + ω(q 3 ).
We have performed a variable change q 1,2 = qq 1,2 at the second line, and noticed q = q 1 + q 2 . Therefore
max Ω (x 1 , x 2 , x 3 ) = Φ (x 1 , x 2 , x 3 ),
where Φ is the convex conjugate of Φ restricted to ]0, 1] 3 . By definition, we also have max Ω (x 1 , x 2 , x 3 ) = Ω (x 1 , x 2 , x 3 ), so that Ω = Φ on R 3 . As Ω is convex and lower semi-continous, we can apply Moreau-Yoshida theorem and obtain Ω = Ω = Φ ≤ Φ.
Suppose that there exists q = (q 1 , q 2 , q 3 ) ∈ 3 such that Φ(q 1 , q 2 , q 3 )  0, ∀ i ∈ [N ], ∀ j ∈ P i , P[w t = j|w t−1 = i] = q i,j .
We set y i,j 1{∃ t > 0 s.t. w t−1 = i, w t = j} where 1 is the characteristic function of an event, thereby defining a random variable Y ∈ Y, with distribution D. We leave implicit the dependency of P in θ and Ω. As the depth of w t (number of edges to connect to the root node) is stricly decreasing with t, (w t ) t reaches node 1 in finite time with probability one and is constant after this event. We introduce the random variables (ȳ j ) j , defined for all j ∈ [N ] as
y j 1{∃ t ≥ 0, w t = j} = i∈Cj y i,j if j = N , 0 otherwise.
By definition, using the fact that P[w t = j|w t−1 = i] is independent of t (Markov property), for all i ∈ C j and for all j ∈ [N − 1], we have
P[y i,j = 1] = E[y i,j ] = P[∃ t > 0, w t−1 = i]P[w t = j|w t−1 = i] = E[ȳ i ]q i,j .
Linearity of the expectation then provides
E[ȳ j ] = i∈Cj E[y i,j ],
with initialization E[ȳ N ] = 1. We recover the same two-step recursion as the one defining E andē, with the same initialization. Hence the probabilistic interpretation of the gradient, where the expectation is taken with respect to the distribution D of Y :
E = E θ,Ω [Y ] andē = E θ,Ω [ȳ].
Algorithm 1 Compute DP Ω (θ) and ∇DP Ω (θ)
Input: Edge weights θ ∈ R N ×N v 1 ← 0,ē N ← 1, Q, E ← 0 ∈ R N ×N for i ∈ [2, . . . , N ] do Topological order v i ← max Ω j∈Pi θ i,j + v j (q i,j ) j∈Pi ← ∇max Ω j∈Pi θ i,j + v j for j ∈ [N − 1, . . . , 1] do Reverse topological order ∀ i ∈ C j , e i,j ← e i,jēi ,ē j ← i∈Cj e i,j Return: DP Ω (θ) = v N , ∇DP Ω (θ) = E ∈ R N ×N Intermediate computation for Algorithm 2 e [ē] N i=1 ∈ R N , Q ∈ R N ×N
Algorithm 2 Compute ∇DP Ω (θ), Z and ∇ 2 DP Ω (θ)Z Input: Edge weights and perturbation θ, Z ∈ R N ×N Call Algorithm 1 with input θ to getē and Q
v 1 ← 0;ė N ← 0,Q,Ė ← 0 ∈ R N ×N for i ∈ [2, . . . , N ] do Topological ordeṙ v i ← j∈Pi q i,j (z i,j +v j ) (A1) (q i,j ) j∈Pi ← J Ω (q i,j ) j∈Pi (z i,j +v j ) j∈Pi (A2) for j ∈ [N − 1, . . . , 1] do Reverse topological order ∀ i ∈ C j ,ė i,j ←q i,jēi + q i,jėi (A3) e j ← i∈Cjė i,j Return: ∇DP Ω (θ), Z =v N ∇ 2 DP Ω (θ)Z =Ė ∈ R N ×N A.6 Computation of the directional derivative ∇DP Ω (θ), Z
The derivations of the following two sections allows to write Algorithm 2.
Letv i ∇v i (θ), Z , where v i (θ) is defined in (6). Since v i only directly depends on v j + θ i,j for j ∈ P i , a straighforward differentiation of ∇v i (θ), Z givesv i = j∈Pi ∂v i ∂v j (v j + z i,j ) .
Recall that ∂vi ∂vj = q i,j and has already been obtained when computing ∇DP Ω (θ). Hence equation (A1), reproduced here:
∀ i ∈ [2, . . . , N ] :v i = j∈Pi q i,j (v j + z i,j ).(13)
This recursion can be computed in topological order, starting fromv 1 = 0 to finish atv N = ∇DP Ω (θ), Z .
A.7 Computation of the Hessian-vector product ∇ 2 DP Ω (θ)Z For convenience, let us define ∇ 2 DP Ω (θ)Z Ė . For (i, j) / ∈ E, we evidently haveė i,j = 0. For (i, j) ∈ E, since θ i,j influences only v i andv i , we obtaiṅ
e i,j = ∂v N ∂θ i,j = ∂v N ∂v i ∂v i ∂θ i,j + ∂v N ∂v i ∂v i ∂θ i,j .
We will now show how to derive each of the right-hand side terms in turn. We already know that ∂vi ∂θi,j = q i,j . We also have ∂v N ∂vi = u i . Indeed, observe thatv j only directly influencesv i for i ∈ C i . Therefore, we have
∂v N ∂v j = i∈Cj ∂v N ∂v i q i,j ∀j ∈ [N − 1](14)
and ∂v N ∂v1 = 1. Comparing (12) and (14), we see that ( ∂v N ∂vi ) i follows the same recursion as ( ∂v N ∂vi ) i . Since ∂v N ∂vn = ∂v N ∂vn , both sequences are equal:
∂v N ∂v i = ∂v N ∂v i = e i .
Next, we derive ∂vi ∂θi,j . Since, for j ∈ P i ,v j + z i,j does not depend on θ i,j , differentiating (13) w.r.t. θ i,j , we obtain
∂v i ∂θ i,j = k∈Pi ∂q i,j ∂θ i,j (v k + z i,k ) = k∈Pi ∂ 2 v i ∂θ i,j ∂θ i,k (v k + z i,k ) q i,j .
This can be conveniently rewritten in a vectorial form aṡ
q i = ∇ 2 max Ω (θ i + v) (z i +v) = J Ω (q i ) (z i +v),
where we have definedv (v 1 , . . . ,v N ) and where we have used the function J Ω defined in §B.1, that conveniently computes the Hessian of max Ω from its gradient. The Hessian has this form for both negentropy and 2 2 regularizations. In a practical implementation, we only need to compute the coordinates (i, j) ofQ, for j ∈ P i . Namely, as specified in (A2),
(q i,j ) j∈Pi ← J Ω (q i,j ) j∈Pi (z i,j +v j ) j∈Pi .
Finally, we derive ∂v N ∂vi .
Since v j influences only v i andv i for i ∈ C j , the chain rule gives
∂v N ∂v i = j∈Ci ∂v N ∂v j ∂v j ∂v i + ∂v N ∂v j ∂v j ∂v i = j∈Cjė i,j ė i .
Combining the above, for any j ∈ [N − 1], we obtain the following two-step recursion (A3), reproduced here:
∀ i ∈ C j ,ė i,j =q i,j e i + q i,jėi andė j = i∈Cjė i,j .
Similarly to the computation of ∇DP Ω (θ), our algorithm computes this recursion in reverse topological order over the graph G, yielding ∇ 2 DP Ω (θ)Z =Ė.
B Examples of algorithm instantiations
We provide the explicit forms of max Ω and its derivative for the negentropy and 2 2 regularizations. Then, we provide details and pseudo-code for the two instances of differentiable dynamic programming presented in §4.
B.1 Examples of max Ω
Negative entropy. When Ω(q) = γ D i=1 q i log q i , where γ > 0 (smaller is less regularized), we obtain
max Ω (x) = γ log D i=1 exp(x i /γ) ∇max Ω (x) = exp(x/γ) D i=1 exp(x i /γ) ∇ 2 max Ω (x) = J Ω (∇max Ω (x)),
where J Ω (q) (Diag(q) − qq )/γ. Note that ∇max Ω (x) recovers the usual "softmax" with temperature γ = 1. For a proof of the expression of max Ω , see, e.g., (Boyd & Vandenberghe, 2004, Example 3.25).
Squared 2 norm. When Ω(x) = γ 2 x 2 2 with γ > 0, we obtain the following expressions
max Ω (x) = q , x − γ 2 q 2 2
∇max Ω (x) = argmin
q∈ D q − x/γ 2 2 = q ∇ 2 max Ω (x) = J Ω (∇max Ω (x)),
where J Ω (q) (Diag(s) − ss / s 1 )/γ and s ∈ {0, 1} D is a vector that indicates the support of q. Note that ∇max Ω (x) is precisely the Euclidean projection onto the simplex of x/γ and can be computed exactly Algorithm 5 Compute DTWΩ(θ) and ∇DTWΩ(θ)
Input: Distance matrix θ ∈ R N A ×N B Forward pass v 0,0 = 0; v i,0 = v 0,j = ∞, i ∈ [N A ], j ∈ [N B ] for i ∈ [1, . . . , N A ], j ∈ [1, . . . , N B ] do v i,j = d i,j + min Ω (v i,j−1 , v i−1,j−1 , v i−1,j ) q i,j = ∇min Ω (v i,j−1 , v i−1,j−1 , v i−1,j ) ∈ R 3 Backward pass q i,N B +1 = q N A +1,j = 0 3 , i ∈ [N A ], j ∈ [N B ] e i,N B +1 = e N A +1,j = 0, i ∈ [N A ], j ∈ [N B ] q N A +1,N B +1 = (0, 1, 0); e N A +1,N B +1 = 1 for j ∈ [N B , . . . , 1], i ∈ [N A , . . . , 1] do e i,j = q i,j+1,1 e i,j+1 + q i+1,j+1,2 e i+1,j+1 + q i+1,j,3 e i+1,j Return: DTW Ω (θ) = v N A ,N B ∇DTW Ω (θ) = (e) N A ,N B i,j=1
Intermediate computations for Algo. 6:
Q (q) N A +1,N B +1,3 i,j,k=1 ; E (e) N A +1,N B +1 i,j=1 Algorithm 6 Compute ∇DTWΩ(θ), Z , ∇ 2 DTWΩ(θ) Z Input: θ ∈ R N A ×N B , Z ∈ R N A ×N B Call Algo. 5 with input θ to retrieve Q and E Forward pasṡ v i,0 =v 0,j = 0, i ∈ [0, . . . , N A ], j ∈ [N B ] for i ∈ [1, . . . , N B ], j ∈ [1, . . . , N A ] dȯ v i,j = z i,j + q i,j,1vi,j−1 + q i,j,2vi−1,j−1 + q i,j,3vi−1,j q i,j = −J Ω (q i,j ) (v i,j−1 ,v i−1,j−1 ,v i−1,j ) ∈ R 3 Backward pasṡ q i,N B +1 =q N A +1,j = 0 3 , i ∈ [0, . . . , N A ], j ∈ [N B ] e i,N B +1 =ė N A +1,j = 0, i ∈ [0, . . . , N A ], j ∈ [N B ] for j ∈ [N B , . . . , 1], i ∈ [N A , . . . , 1] dȯ e i,j =q i,j+1
,1 e i,j+1 + q i,j+1,1ėi,j+1 + q i+1,j+1,2 e i+1,j+1 + q i+1,j+1,2ėi+1,j+1 + q i+1,j,3 e i+1,j + q i+1,j,3ėi+1,j
Return: ∇DTW Ω (θ), Z =v N A ,N B ∇ 2 DTW Ω (θ) Z = (ė) N A ,N B i,j=1
Representation. We follow the notation of §4, i.e. we represent Y and θ as N A × N B matrices. We represent Y as a binary matrix such that y i,j = 1 if a i is aligned with b j , and y i,j = 0 otherwise. Likewise, we represent θ as a real matrix such that θ i,j is a measure of "discrepancy" between a i and b j .
Algorithms. Following the DTW literature (Sakoe & Chiba, 1978), we seek an alignment with minimal cost. For that reason, we introduce the smoothed min operator, its gradient and its Hessian as follows
min Ω (x) −max Ω (−x) ∇min Ω (x) = ∇max Ω (−x) ∇ 2 min Ω (x) = −∇ 2 max Ω (−x) = −J Ω (∇max Ω (−x)) = −J Ω (∇min Ω (x)).
Applying (6) to the DTW DAG gives rise to a smoothed version of the algorithm. Let v i,j (θ) be the alignment cost up to cell (i, j). Then the smoothed DTW recursion is
v i,j (θ) = θ i,j + min Ω (v i,j−1 (θ), v i−1,j−1 (θ), v i−1,j (θ))
The value DTW Ω (θ) v N A ,N B (θ) can be computed in O(N A N B ) time. Applying the derivations of §3.3 and §3.4 to this specific DAG, we can compute ∇DTW Ω (θ), ∇DTW Ω (θ), Z and ∇ 2 DTW Ω (θ)Z with the same complexity. The procedures, with appropriate handling of the edge cases, are summarized in Algorithm 5 and 6, respectively.
Note that when Ω is the negative entropy, DTW Ω (θ) is known as soft-DTW (Cuturi & Blondel, 2017). While the DP computation of DTW Ω (θ) and of its gradient were already known, the generalization to any strongly convex Ω and the computation of ∇ 2 DTW Ω (θ)Z are new. From Proposition 2 property 1, DTW Ω (θ) is a concave function of the discrepancy matrix θ for any Ω. With respect to time-series, DTW Ω is neither convex nor concave.
C Experimental details and further results
We finally provide details on the architecture used in experiments, with additionnals figures. : Test predictions from the entropy and 2 2 regularized named entity recognition (NER) models. Red dots indicate ground truth. When using 2 2 regularization, model predictions are sparse (grey borders indicates non-zero cells). They are thus easier to introspect for ambiguities, as we can list a finite number of possible outputs.
C.1 Named entity recognition (section §5.2)
Our model extracts word embedding from a 300-dimensional lookup table concatenated with a 50-dimensional character embedding. This character embedding corresponds to the concatenation of the last hidden unit of a bi-directional character LSTM, as in Lample et al. (2016). Character embedding size is set to 50. A word LSTM then produces sentence-aware features for each word. This LSTM is bi-directional with 100dimensional hidden units per direction. The final features X used to build the potential tensor θ are thus 200-dimensional. Note that, in contrast with Lample et al. (2016):
• The look-up table is initialized with 300-dimensional embeddings from FastText (Joulin et al., 2016), trained on Wikipedia corpus. • We do not pad letters prior to feeding the character LSTM as it is not principled.
• We do not train the unknown word embedding as we found it had no effect.
We convert tags to the IOBES (Inside-Outside-Begin-End-Stop) scheme to build a richer Vit Ω model than if we used the simpler IOB (Inside-Outside-Begin) scheme, that has a lower number of tags. We performed a small grid-search to select the step-size and batch-size used for optimization: s ∈ {0.005, 0.01, 0.02}, b ∈ {8, 32, 128}. For each language and each loss, we select the highest-scoring model on the validation set, and report the test score.
The model is strongly subject to overfitting using the convex surrogate loss and the log likelihood. We have to use a small batch size (b = 8) and vanilla SGD with large step size (s = 0.01) to avoid this overfitting issue. For all losses, accelerated stochastic optimizers have all lower generalization performance than SGD, as also noticed in (Lample et al., 2016) when using the classical negative log-likelihood as a loss. Figure 7: Alignment maps between score onsets and audio frames on test data from the Bach10 dataset. Our end-to-end trained model qualitatively performs better than the baseline model.
Visualization. The models using 2 2 regularization perform nearly on par with the ones using negentropy, as demonstrated in Table 1. On the other hand, 2 2 regularization leads to tag probability vectors that are sparse and hence easier to parse. They allow to detect ambiguities more easily. We display a few tagged English sequences in Figure 6. The model using 2 2 regularization correctly identifies an ambiguous entity (Union Bank of Switzerland ) and can be used to propose two tag sequences: (B-ORG, I-ORG, I-ORG, E-ORG) or (B-ORG, E-ORG, O, S-LOC). Probabilities of every tag sequence can be computed using the matrix Q, as described in §3.3 -this remains tractable as long as the matrix Q is sparse enough, so that the number of non-zero probabilities sequence remains low. On the other hand, the model using negentropy regularization never assign a zero probability to any tag sequence -it is therefore not tractable to provide the user with a small set of interesting sequences.
C.2 Supervised audio-to-score transcription (section §5.3) Audio sequences, sampled at 22.05 kHz, are split into frames of 512 samples. We extract the following features from these sequences: energy, spectral centroid, spectral bandwidth, and the 5 first MFCC features. All features are centered around the median and normalized. The ∇DTW Ω layer is written in Cython 1 , and hence run on CPU. This technical choice was suggested by the fact that we have to write explicit loops to specify the topological and reverse topological pass over the DTW computation graph (see Algorithm 5). However, it is possible to use only contiguous vector operations and thus take advantage of GPU computations -this is left for future work. We use SciPy's 2 LBFGS-B solver to perform end-to-end training and multinomial regression. We use a 2 2 regularization on the weight W ,: we selected it using a grid search over {10 −5 , 10 −4 , . . . , 1} and selected 10 −3 .
Further vizualisation. In Figure 7, we display the alignment maps we obtained using our algorithm and using the baseline multinomial model followed by a hard-DTW alignment computation. These alignment maps correspond to the predicted onsets of keys. Our model (in orange) performs visibly better in predicting onsets. We use OpenNMT-py library 3 to fit our structured attention model. Model architecture and optimization details are as follow:
• We use a bidirectional LSTM encoder and decoder, with 500 units in each direction and a depth of 2 layers . • The decoder is fed with the input representation as in Luong et al. (2015).
• SGD training with s = 1 learning rate, decaying from epoch 8 to epoch 15 with rate 0.65, batch size of size 256. • Training sentence of lengths superior to 50 are ignored, and translated sentence are forced to a length inferior to 100. • The temperature parameter is set to γ = 2 for entropy, and γ = 10 for 2 2 . Performance is not affected much by this parameter, provided that it is not set too low in the 2 2 case -with a too small γ, Vit Ω reduces to unregularized MAP estimation and ∇Vit Ω has zero derivatives.
We use a 1-million sentence subject of WMT14 English-to-French corpus, available at http://nmt-benchmark.net/. We use Moses tokenizer and do not perform any post-processing, before computing BLEU score on detokenized sentences (multi bleu.perl script).
Implementation. We implemeted a batch version of the ∇Vit Ω layer on GPU, using the PyTorch tensor API. Model with negentropy-regularized attention mechanism runs 1/2 as fast as the softmax attention mechanism (approximately 7500 tokens/s vs 15000 tokens/s on a single Nvidia Titan X Pascal). With 2 2 regularization, it is only 1/3 as fast: approximately 5000 tokens/s. Although this remains reasonable, it could certainly be optimized by rewriting kernels using lower-level languages (e.g., using ATen API from PyTorch.) Further results. Table 3 provides BLEU scores for both translation directions on the 1 million sentence subset of WMT14 we used. We observe that the introduction of structure and sparsity does not hinder the general performance of the model. We provide several examples of attention maps in Figure 8, that illustrate the sparsity patterns 2 2 regularization uncovers. 
Figure 2 :
2Computational graph of the Viterbi algorithm.
Figure 3 :
3Computational graph of the DTW algorithm.
Experiments. We measure the performance of the different losses and regularizations on the four languages of the CoNLL 2003 (Tjong Kim Sang & De Meulder, 2003) dataset. Following Lample et al. (2016), who use the 0,−H loss, we use a character LSTM and pretrained embeddings computed using FastText (Joulin et al., 2016) on Wikipedia. Those are fed to a word bidirectional LSTM to obtain X. Architecture details are provided in §C.1. Results are reported in
Figure 4 :
4Attention maps obtained with structured attention. Although both regularizations led to the same translation (y-axis) in this example, attention is sparse and more interpretable with 2 2 .
Figure 5 :
5In general, v 6 (θ) = DP Ω (θ) = LP Ω (θ).
Figure 6
6Figure 6: Test predictions from the entropy and 2 2 regularized named entity recognition (NER) models. Red dots indicate ground truth. When using 2 2 regularization, model predictions are sparse (grey borders indicates non-zero cells). They are thus easier to introspect for ambiguities, as we can list a finite number of possible outputs.
Figure 8 :
8Attention on test samples from Newstest2014. Borders indicate non-zero cells. Translations (y-axis) are often qualitatively equivalent, while attentions maps are sparse in the 2 2 case.
Table 1 :
1F 1 score comparison on CoNLL03 NER datasets.Ω 
Loss 
English Spanish German Dutch 
Negentropy 
Surrogate 
90.80 
86.68 
77.35 87.56 
Relaxed 
90.47 
86.20 
77.56 
87.37 
2 
2 
Surrogate 
90.86 
85.51 
76.01 
86.58 
Relaxed 
89.49 
84.07 
76.91 
85.90 
(Lample et al., 2016) 
90.96 
85.75 
78.76 
81.74 
Table 2 :
2Mean absolute deviation of alignment using an end-to-end trained multinomial classifier and a pre-trained one.Linear model 
Train 
Test 
End-to-end trained 0.17 ± 0.01 1.07 ± 0.61 
Pretrained 
1.80 ± 0.14 
3.69 ± 2.85 
Random θ 
14.64 ± 2.63 14.64 ± 0.29 
Table 3 :
3Detokenized BLEU score on newstest2014 data using regularized and unregularized attention.Attention model 
WMT14 1M fr→en WMT14 en→fr 
Softmax 
27.96 
28.08 
Entropy regularization 
27.96 
27.98 
2 
2 reg. 
27.21 
27.28 
C.3 Structured and sparse attention (section  §6) 
other nodes. The cardinality of Y is the delannoy(N A −1, N B −1) number(Sulanke, 2003;Banderier & Schwer, 2005). That number grows exponentially with N A and N B .
http://cython.org/ 2 http://scipy.org/
http://opennmt.net/
AcknowledgementsMB thanks Vlad Niculae and Marco Cuturi for many fruitful discussions. AM thanks Julien Mairal, Inria Thoth and Inria Parietal for lending him the computational resources necessary to run the experiments. He Differentiable Dynamic Programming for Structured Prediction and Attention 13 thanks University Paris-Saclay and his Ph.D. supervisors Bertrand Thirion and Gaël Varoquaux for allowing him to do an internship at NTT, and Olivier Grisel for his insightful comments.Algorithm 3 Compute Vit Ω (θ) and ∇Vit Ω (θ)Backward pass u T +1 = (1, 0, . . . , 0) ∈ R S for t ∈ [T, . . . , 0], j ∈ [S] do e t,·,j = q t+1,·,j • u t+1 ; u t,j = e t,·,j , 1 SReturn:Vit Ω (θ) = v T +1,1 ∇Vit Ω (θ) = (e t−1,i,j ) T,S,S t=1,i,j=1Intermediary computations for Alg. 4:Algorithm 4 Compute ∇Vit Ω (θ), Z and ∇ 2 Vit Ω (θ)Zin worst-case O(D log D) time using the algorithm of(Michelot, 1986)or in expected O(D) time using the randomized pivot algorithm of(Duchi et al., 2008). It can be efficiently performed on Nvidia GPUs since recently. An important benefit of the squared 2 norm, compared to the negative entropy, is that ∇max Ω (x) tends to be sparse. This is useful, among other things, to define sparse attention mechanisms(Martins & Astudillo, 2016;Niculae & Blondel, 2017).B.2 Sequence prediction with the smoothed Viterbi algorithmComputational graph. As illustrated in §4, the DAG contains a start node, S nodes for each time step and end node. Therefore |V| = N = T S +2. Only nodes from consecutive time steps are connected to each other. Taking into account the start and end nodes, the total number of edges is therefore |E| = (T − 1)S 2 + 2S.Representation. We follow the notation of §4, i.e. we represent Y and θ as T × S × S tensors (we can safely ignore the edges connected to the end node since their value is 0). We represent Y as a binary tensor such that y t,i,j = 1 if Y is in states i and j in time steps t and t − 1, and y t,i,j = 0 otherwise. Likewise, we represent the potentials θ as a real tensor such that θ t,i,j contains the potential of transitioning from state j to state i on time t.Algorithms. Applying recursion (6) to this specific DAG, we obtain a smoothed version of the Viterbi algorithm. Let v t,i be the score of being in state i up to time t. We can rewrite the smoothed Bellman recursion asThe value Vit Ω (θ) max Ω (v T (θ)) can be computed in topological order, starting from v 0 (θ). The total computational cost is O(T S 2 ). Using the computations of §3.3 and §3.4 to this specific DAG, we can compute ∇Vit Ω (θ), ∇Vit Ω (θ), Z and ∇ 2 Vit Ω (θ)Z with the same complexity. The procedures are summarized in Algorithm 3 and Algorithm 4, respectively. From Proposition 2 property 1, Vit Ω (θ) is a convex function for any Ω.B.3 Monotonic aligment prediction with the smoothed DTWComputational graph. As illustrated in §4, the DAG contains a start node and N A N B nodes. Therefore, |V| = N = N A N B + 1. Due to the monotonic constraint, each node may only be connected with at most
The generalized distributive law. Srinivas M Aji, Robert J Mceliece, IEEE Transactions on Information Theory. 462Aji, Srinivas M and McEliece, Robert J. The generalized distributive law. IEEE Transactions on Information Theory, 46(2):325-343, 2000.
OptNet: Differentiable optimization as a layer in neural networks. Brandon Amos, J Kolter, Zico, Proc. of ICML. of ICMLAmos, Brandon and Kolter, J. Zico. OptNet: Differentiable optimization as a layer in neural networks. In Proc. of ICML, pp. 136-145, 2017.
Neural Machine Translation by Jointly Learning to Align and Translate. Dzmitry Bahdanau, Kyunghyun Cho, Yoshua Bengio, Proc. of ICLR. of ICLRBahdanau, Dzmitry, Cho, Kyunghyun, and Bengio, Yoshua. Neural Machine Translation by Jointly Learning to Align and Translate. In Proc. of ICLR, 2015.
Predicting Structured Data. Gökhan Bakır, Hofmann, Thomas, Schölkopf, Bernhard, Alexander J Smola, Ben Taskar, S V Vishwanathan, The MIT PressBakır, Gökhan, Hofmann, Thomas, Schölkopf, Bernhard, Smola, Alexander J, Taskar, Ben, and Vish- wanathan, S. V. N. Predicting Structured Data. The MIT Press, 2007.
Why Delannoy numbers?. Cyril Banderier, Sylviane Schwer, Journal of Statistical Planning and Inference. 1351Banderier, Cyril and Schwer, Sylviane. Why Delannoy numbers? Journal of Statistical Planning and Inference, 135(1):40-54, 2005.
Statistical inference for probabilistic functions of finite state markov chains. The Annals of Mathematical Statistics. Leonard E Baum, Ted Petrie, 37Baum, Leonard E. and Petrie, Ted. Statistical inference for probabilistic functions of finite state markov chains. The Annals of Mathematical Statistics, 37(6):1554-1563, 1966.
Smoothing and First Order Methods: A Unified Framework. Amir Beck, Marc Teboulle, SIAM Journal on Optimization. 222Beck, Amir and Teboulle, Marc. Smoothing and First Order Methods: A Unified Framework. SIAM Journal on Optimization, 22(2):557-580, 2012.
On the theory of dynamic programming. Richard Bellman, Proc. of the National Academy of Sciences. of the National Academy of Sciences38Bellman, Richard. On the theory of dynamic programming. Proc. of the National Academy of Sciences, 38 (8):716-719, 1952.
Control of uncertain systems with a set-membership description of the uncertainty. Dimitri P Bertsekas, Massachusetts Institute of TechnologyPhD thesisBertsekas, Dimitri P. Control of uncertain systems with a set-membership description of the uncertainty. PhD thesis, Massachusetts Institute of Technology, 1971.
Convex optimization. Stephen Boyd, Lieven Vandenberghe, Cambridge university pressBoyd, Stephen and Vandenberghe, Lieven. Convex optimization. Cambridge university press, 2004.
Soft-DTW: a Differentiable Loss Function for Time-Series. Marco Cuturi, Mathieu Blondel, Proc. of ICML. of ICMLCuturi, Marco and Blondel, Mathieu. Soft-DTW: a Differentiable Loss Function for Time-Series. In Proc. of ICML, pp. 894-903, 2017.
The theory of max-min, with applications. John M Danskin, SIAM Journal on Applied Mathematics. 144Danskin, John M. The theory of max-min, with applications. SIAM Journal on Applied Mathematics, 14 (4):641-664, 1966.
Differentiable learning of submodular functions. Josip Djolonga, Andreas Krause, Proc. of NIPS. of NIPSDjolonga, Josip and Krause, Andreas. Differentiable learning of submodular functions. In Proc. of NIPS, pp. 1014-1024, 2017.
. Zhiyao Duan, Bryan Pardo, Bach 10 datasetDuan, Zhiyao and Pardo, Bryan. Bach 10 dataset. http://music.cs.northwestern.edu/data/Bach10. html.
Soundprism: An online system for score-informed source separation of music audio. Zhiyao Duan, Bryan Pardo, IEEE Journal of Selected Topics in Signal Processing. 56Duan, Zhiyao and Pardo, Bryan. Soundprism: An online system for score-informed source separation of music audio. IEEE Journal of Selected Topics in Signal Processing, 5(6):1205-1215, 2011.
Efficient projections onto the 1 -ball for learning in high dimensions. John Duchi, Shalev-Shwartz, Shai, Yoram Singer, Tushar Chandra, Proc. of ICML. of ICMLDuchi, John, Shalev-Shwartz, Shai, Singer, Yoram, and Chandra, Tushar. Efficient projections onto the 1 -ball for learning in high dimensions. In Proc. of ICML, pp. 272-279, 2008.
Inside-outside and forward-backward algorithms are just backprop (tutorial paper). Jason Eisner, Proc. of the Workshop on Structured Prediction for NLP. of the Workshop on Structured Prediction for NLPEisner, Jason. Inside-outside and forward-backward algorithms are just backprop (tutorial paper). In Proc. of the Workshop on Structured Prediction for NLP, pp. 1-17, 2016.
Metric learning for temporal sequence alignment. Damien Garreau, Lajugie, Rémi, Arlot, Sylvain, Francis Bach, Proc. of NIPS. of NIPSGarreau, Damien, Lajugie, Rémi, Arlot, Sylvain, and Bach, Francis. Metric learning for temporal sequence alignment. In Proc. of NIPS, pp. 1817-1825, 2014.
Entropy functions and functional equations. Eszter Gselmann, Mathematical Communications. 16Gselmann, Eszter. Entropy functions and functional equations. Mathematical Communications, (16):347- 357, 2011.
Entropy of terminal distributions and the Fibonnacci trees. Yasuichi Horibe, The Fibonacci Quarterly. 26Horibe, Yasuichi. Entropy of terminal distributions and the Fibonnacci trees. The Fibonacci Quarterly, (26):135-140, 1988.
Differentiable Dynamic Programming for Structured Prediction and Attention 14. Differentiable Dynamic Programming for Structured Prediction and Attention 14
Armand Joulin, Grave, Edouard, Bojanowski, Piotr, Douze, Matthijs, Hérve Jégou, Tomas Mikolov, Fasttext, arXiv:1612.03651zip: Compressing text classification models. arXiv preprintJoulin, Armand, Grave, Edouard, Bojanowski, Piotr, Douze, Matthijs, Jégou, Hérve, and Mikolov, Tomas. Fasttext.zip: Compressing text classification models. arXiv preprint arXiv:1612.03651, 2016.
Structured Attention Networks. Yoon Kim, Denton, Carl, Luong Hoang, Alexander M Rush, Proc. of ICLR. of ICLRKim, Yoon, Denton, Carl, Hoang, Luong, and Rush, Alexander M. Structured Attention Networks. In Proc. of ICLR, 2017.
Conditional random fields: Probabilistic models for segmenting and labeling sequence data. John Lafferty, Andrew Mccallum, Pereira, C N Fernando, Proc. of ICML. of ICMLLafferty, John, McCallum, Andrew, and Pereira, Fernando CN. Conditional random fields: Probabilistic models for segmenting and labeling sequence data. In Proc. of ICML, pp. 282-289, 2001.
Neural architectures for named entity recognition. Guillaume Lample, Ballesteros, Miguel, Subramanian, Sandeep, Kazuya Kawakami, Chris Dyer, Proc. of NAACL. of NAACLLample, Guillaume, Ballesteros, Miguel, Subramanian, Sandeep, Kawakami, Kazuya, and Dyer, Chris. Neural architectures for named entity recognition. In Proc. of NAACL, pp. 260-270, 2016.
The representation of the cumulative rounding error of an algorithm as a Taylor expansion of the local rounding errors. Seppo Linnainmaa, Univ. HelsinkiPhD thesisLinnainmaa, Seppo. The representation of the cumulative rounding error of an algorithm as a Taylor expansion of the local rounding errors. PhD thesis, Univ. Helsinki, 1970.
Effective Approaches to Attention-based Neural Machine Translation. Thang Luong, Hieu Pham, Manning, D Christopher, Proc. of EMNLP. of EMNLPLuong, Thang, Pham, Hieu, and Manning, Christopher D. Effective Approaches to Attention-based Neural Machine Translation. In Proc. of EMNLP, pp. 1412-1421, 2015.
End-to-end sequence labeling via bi-directional lstm-cnns-crf. Xuezhe Ma, Eduard Hovy, Proc. of ACL. of ACLMa, Xuezhe and Hovy, Eduard. End-to-end sequence labeling via bi-directional lstm-cnns-crf. In Proc. of ACL, pp. 1064-1074, 2016.
From softmax to sparsemax: A sparse model of attention and multi-label classification. André F T Martins, Ramón Astudillo, Fernandez, Proc. of ICML. of ICMLMartins, André F.T. and Astudillo, Ramón Fernandez. From softmax to sparsemax: A sparse model of attention and multi-label classification. In Proc. of ICML, pp. 1614-1623, 2016.
A finite algorithm for finding the projection of a point onto the canonical simplex of R n. Christian Michelot, Journal of Optimization Theory and Applications. 501Michelot, Christian. A finite algorithm for finding the projection of a point onto the canonical simplex of R n . Journal of Optimization Theory and Applications, 50(1):195-200, 1986.
. Jean-Jacques Moreau, 93Proximit et dualit dans un espace hilbertien. Bullet de la Société Mathémathique de FranceMoreau, Jean-Jacques. Proximit et dualit dans un espace hilbertien. Bullet de la Société Mathémathique de France, 93(2):273-299, 1965.
Smooth minimization of non-smooth functions. Yurii Nesterov, Mathematical Programming. 1031Nesterov, Yurii. Smooth minimization of non-smooth functions. Mathematical Programming, 103(1):127- 152, 2005.
A Regularized Framework for Sparse and Structured Neural Attention. Vlad Niculae, Mathieu Blondel, Proc. of NIPS. of NIPSNiculae, Vlad and Blondel, Mathieu. A Regularized Framework for Sparse and Structured Neural Attention. In Proc. of NIPS, pp. 3340-3350, 2017.
Pytorch: Tensors and dynamic neural networks in Python with strong GPU acceleration. Adam Paszke, Gross, Sam, Soumith Chintala, Gregory Chanan, Paszke, Adam, Gross, Sam, Chintala, Soumith, and Chanan, Gregory. Pytorch: Tensors and dynamic neural networks in Python with strong GPU acceleration, 2017.
Probabilistic Reasoning in Intelligent Systems: Networks of Plausible Inference. Judea Pearl, ElsevierPearl, Judea. Probabilistic Reasoning in Intelligent Systems: Networks of Plausible Inference. Elsevier, 1988.
Fast exact multiplication by the Hessian. Pearlmutter, A Barak, Neural computation. 61Pearlmutter, Barak A. Fast exact multiplication by the Hessian. Neural computation, 6(1):147-160, 1994.
A tutorial on hidden markov models and selected applications in speech recognition. Lawrence R Rabiner, Readings in Speech Recognition. Rabiner, Lawrence R. A tutorial on hidden markov models and selected applications in speech recognition. In Readings in Speech Recognition, pp. 267-296. 1990.
Dynamic programming algorithm optimization for spoken word recognition. Hiroaki Sakoe, Seibi Chiba, IEEE Transactions on Acoustics, Speech, and Signal Processing. 26Sakoe, Hiroaki and Chiba, Seibi. Dynamic programming algorithm optimization for spoken word recognition. IEEE Transactions on Acoustics, Speech, and Signal Processing, 26:43-49, 1978.
Minimum risk annealing for training log-linear models. David A Smith, Jason Eisner, Proc. of COLING/ACL. of COLING/ACLSmith, David A. and Eisner, Jason. Minimum risk annealing for training log-linear models. In Proc. of COLING/ACL, pp. 787-794, 2006.
Minimum-risk training of approximate crf-based nlp systems. Veselin Stoyanov, Jason Eisner, Proc. of NAACL. of NAACLStoyanov, Veselin and Eisner, Jason. Minimum-risk training of approximate crf-based nlp systems. In Proc. of NAACL, pp. 120-130, 2012.
Objects counted by the central Delannoy numbers. Robert A Sulanke, Journal of Integer Sequences. 613Sulanke, Robert A. Objects counted by the central Delannoy numbers. Journal of Integer Sequences, 6(1): 3, 2003.
An introduction to conditional random fields. Foundations and Trends in Machine Learning. Charles Sutton, Andrew Mccallum, 4Sutton, Charles, McCallum, Andrew, et al. An introduction to conditional random fields. Foundations and Trends in Machine Learning, 4(4):267-373, 2012.
Introduction to the CoNLL-2003 shared task: Languageindependent named entity recognition. Tjong Kim Sang, Erik F , De Meulder, Fien , Proc. of CoNLL. of CoNLLTjong Kim Sang, Erik F. and De Meulder, Fien. Introduction to the CoNLL-2003 shared task: Language- independent named entity recognition. In Proc. of CoNLL, pp. 142-147, 2003.
Differentiable Dynamic Programming for Structured Prediction and Attention 15. Differentiable Dynamic Programming for Structured Prediction and Attention 15
Large margin methods for structured and interdependent output variables. Tsochantaridis, Ioannis, Joachims, Thorsten, Thomas Hofmann, Yasemin Altun, Journal of Machine Learning Research. 6Tsochantaridis, Ioannis, Joachims, Thorsten, Hofmann, Thomas, and Altun, Yasemin. Large margin meth- ods for structured and interdependent output variables. Journal of Machine Learning Research, 6:1453- 1484, 2005.
Abstract dynamic programming models under commutativity conditions. Sergio Verdu, Poor, Vincent, SIAM Journal on Control and Optimization. 254Verdu, Sergio and Poor, H Vincent. Abstract dynamic programming models under commutativity conditions. SIAM Journal on Control and Optimization, 25(4):990-1006, 1987.
Error bounds for convolutional codes and an asymptotically optimum decoding algorithm. Andrew Viterbi, IEEE Transactions on Information Theory. 132Viterbi, Andrew. Error bounds for convolutional codes and an asymptotically optimum decoding algorithm. IEEE Transactions on Information Theory, 13(2):260-269, 1967.
Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning. Martin J Wainwright, Michael I Jordan, 1Wainwright, Martin J. and Jordan, Michael I. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning, 1(12):1-305, 2008.
Differentiable Dynamic Programming for Structured Prediction and Attention
https://arxiv.org/abs/1802.03676v1
https://arxiv.org/abs/2112.14570
Lyapunov Exponents for Diversity in Differentiable Games
2021-12-24
Lorraine; Jonathan; Vicol; Paul; Parker-Holder; Jack; Kachman; Tal; Metz; Luke; Foerster; Jakob
<div><div>
 
 
 
 
 
 
 <p><a href="/pdf/2112.14570">Download PDF</a></p><blockquote>
 Abstract: Ridge Rider (RR) is an algorithm for finding diverse solutions to
optimization problems by following eigenvectors of the Hessian ("ridges"). RR
is designed for conservative gradient systems (i.e., settings involving a
single loss function), where it branches at saddles - easy-to-find bifurcation
points. We generalize this idea to non-conservative, multi-agent gradient
systems by proposing a method - denoted Generalized Ridge Rider (GRR) - for
finding arbitrary bifurcation points. We give theoretical motivation for our
method by leveraging machinery from the field of dynamical systems. We
construct novel toy problems where we can visualize new phenomena while giving
insight into high-dimensional problems of interest. Finally, we empirically
evaluate our method by finding diverse solutions in the iterated prisoners'
dilemma and relevant machine learning problems including generative adversarial
networks.
 </blockquote>
 
 
 </div><div>
 <h2>Submission history</h2><p> From: Jonathan Lorraine [<a href="/show-email/76b75863/2112.14570">view email</a>]
 <br /><strong>[v1]</strong>
Fri, 24 Dec 2021 22:48:14 UTC (12,450 KB)<br /></p></div>||||I|||| Skip to main content
 We gratefully acknowledge support from
 the Simons Foundation and member institutions.
 > cs > arXiv:2112.14570

 Help | Advanced Search

 All fields Title Author Abstract Comments Journal reference ACM classification MSC classification Report number arXiv identifier DOI ORCID arXiv author ID Help pages Full text
 Search
 GO

 quick links

 * Login
 * Help Pages
 * About

 Computer Science > Computer Science and Game Theory

 arXiv:2112.14570 (cs)
 [Submitted on 24 Dec 2021]

 Title: Lyapunov Exponents for Diversity in Differentiable Games

 Authors: Jonathan Lorraine, Paul Vicol, Jack Parker-Holder, Tal Kachman, Luke Metz, Jakob Foerster
 Download PDF
 Abstract: Ridge Rider (RR) is an algorithm for finding diverse solutions to optimization problems by following eigenvectors of the Hessian ("ridges"). RR is designed for conservative gradient systems (i.e., settings involving a single loss function), where it branches at saddles - easy-to-find bifurcation points. We generalize this idea to non-conservative, multi-agent gradient systems by proposing a method - denoted Generalized Ridge Rider (GRR) - for finding arbitrary bifurcation points. We give theoretical motivation for our method by leveraging machinery from the field of dynamical systems. We construct novel toy problems where we can visualize new phenomena while giving insight into high-dimensional problems of interest. Finally, we empirically evaluate our method by finding diverse solutions in the iterated prisoners' dilemma and relevant machine learning problems including generative adversarial networks.
 Comments: AAMAS2022, 24 pages  
 Subjects: Computer Science and Game Theory (cs.GT) ; Machine Learning (cs.LG); Multiagent Systems (cs.MA)
 Cite as: arXiv:2112.14570 [cs.GT]  
 (or arXiv:2112.14570v1 [cs.GT] for this version) 
 https://doi.org/10.48550/arXiv.2112.14570 
 Focus to learn more  
 arXiv-issued DOI via DataCite 
 

 Submission history

 From: Jonathan Lorraine [view email]
 [v1] Fri, 24 Dec 2021 22:48:14 UTC (12,450 KB)
 Full-text links:

 Download:

 * PDF
 * Other formats
 (license)
 Current browse context:
 cs.GT
 < prev | next >
 new | recent | 2112
 Change to browse by:
 cs
 cs.LG
 cs.MA

 References & Citations

 * NASA ADS
 * Google Scholar
 * Semantic Scholar

 DBLP - CS Bibliography

 listing | bibtex
 Jonathan Lorraine
 Paul Vicol
 Jack Parker-Holder
 Tal Kachman
 Luke Metz
 …
 a export bibtex citation Loading...

 Bibtex formatted citation

 ×
 loading...
 Data provided by:

 Bookmark

 Bibliographic Tools

 Bibliographic and Citation Tools

 Bibliographic Explorer Toggle
 Bibliographic Explorer (What is the Explorer?)
 Litmaps Toggle
 Litmaps (What is Litmaps?)
 scite.ai Toggle
 scite Smart Citations (What are Smart Citations?)
 Code, Data, Media

 Code, Data and Media Associated with this Article

 Links to Code Toggle
 Papers with Code (What is Papers with Code?)
 ScienceCast Toggle
 ScienceCast (What is ScienceCast?)
 Demos

 Demos

 Replicate Toggle
 Replicate (What is Replicate?)
 Spaces Toggle
 Hugging Face Spaces (What is Spaces?)
 Related Papers

 Recommenders and Search Tools

 Connected Papers Toggle
 Connected Papers (What is Connected Papers?)
 Core recommender toggle
 CORE Recommender (What is CORE?)
 About arXivLabs

 arXivLabs: experimental projects with community collaborators

 arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

 Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

 Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.

 Which authors of this paper are endorsers? | Disable MathJax (What is MathJax?)
 * About
 * Help
 * Click here to contact arXiv Contact
 * Click here to subscribe Subscribe
 * Copyright
 * Privacy Policy
 * Web Accessibility Assistance

 * arXiv Operational Status
 Get status notifications via email or slack

Lyapunov Exponents for Diversity in Differentiable Games
May 9-13, 2022
Jonathan Lorraine 
Vector Institute
Google Research, Brain Team
Exponents for Diversity in Differentiable Games
University of Toronto
Vector Institute
University of Toronto
University of Oxford
Radboud University
University of Oxford
Paul Vicol 
Vector Institute
Google Research, Brain Team
Exponents for Diversity in Differentiable Games
University of Toronto
Vector Institute
University of Toronto
University of Oxford
Radboud University
University of Oxford
Jack Parker-Holder 
Vector Institute
Google Research, Brain Team
Exponents for Diversity in Differentiable Games
University of Toronto
Vector Institute
University of Toronto
University of Oxford
Radboud University
University of Oxford
Tal Kachman 
Vector Institute
Google Research, Brain Team
Exponents for Diversity in Differentiable Games
University of Toronto
Vector Institute
University of Toronto
University of Oxford
Radboud University
University of Oxford
Luke Metz 
Vector Institute
Google Research, Brain Team
Exponents for Diversity in Differentiable Games
University of Toronto
Vector Institute
University of Toronto
University of Oxford
Radboud University
University of Oxford
Jakob Foerster 
Vector Institute
Google Research, Brain Team
Exponents for Diversity in Differentiable Games
University of Toronto
Vector Institute
University of Toronto
University of Oxford
Radboud University
University of Oxford
Jonathan Lorraine 
Vector Institute
Google Research, Brain Team
Exponents for Diversity in Differentiable Games
University of Toronto
Vector Institute
University of Toronto
University of Oxford
Radboud University
University of Oxford
Paul Vicol 
Vector Institute
Google Research, Brain Team
Exponents for Diversity in Differentiable Games
University of Toronto
Vector Institute
University of Toronto
University of Oxford
Radboud University
University of Oxford
Jack Parker-Holder 
Vector Institute
Google Research, Brain Team
Exponents for Diversity in Differentiable Games
University of Toronto
Vector Institute
University of Toronto
University of Oxford
Radboud University
University of Oxford
Tal Kachman 
Vector Institute
Google Research, Brain Team
Exponents for Diversity in Differentiable Games
University of Toronto
Vector Institute
University of Toronto
University of Oxford
Radboud University
University of Oxford
Luke Metz 
Vector Institute
Google Research, Brain Team
Exponents for Diversity in Differentiable Games
University of Toronto
Vector Institute
University of Toronto
University of Oxford
Radboud University
University of Oxford
Jakob 2022 Foerster 
Vector Institute
Google Research, Brain Team
Exponents for Diversity in Differentiable Games
University of Toronto
Vector Institute
University of Toronto
University of Oxford
Radboud University
University of Oxford
Lyapunov 
Vector Institute
Google Research, Brain Team
Exponents for Diversity in Differentiable Games
University of Toronto
Vector Institute
University of Toronto
University of Oxford
Radboud University
University of Oxford
Lyapunov Exponents for Diversity in Differentiable Games
Proc. of the 21st International Conference on Autonomous Agents and Multiagent Systems (AAMAS 2022), Online
. of the 21st International Conference on Autonomous Agents and Multiagent Systems (AAMAS 2022), OnlineMay 9-13, 2022ACM Reference Format:
Ridge Rider (RR) is an algorithm for finding diverse solutions to optimization problems by following eigenvectors of the Hessian ("ridges"). RR is designed for conservative gradient systems (i.e., settings involving a single loss function), where it branches at saddles -easy-to-find bifurcation points. We generalize this idea to nonconservative, multi-agent gradient systems by proposing a method -denoted Generalized Ridge Rider (GRR) -for finding arbitrary bifurcation points. We give theoretical motivation for our method by leveraging machinery from the field of dynamical systems. We construct novel toy problems where we can visualize new phenomena while giving insight into high-dimensional problems of interest. Finally, we empirically evaluate our method by finding diverse solutions in the iterated prisoners' dilemma and relevant machine learning problems including generative adversarial networks.
INTRODUCTION
In machine learning it is often useful to select particular solutions with desirable properties that an arbitrary (global or local) minimum might not have. For example, finding solutions in image classification using shapes which generalize more effectively than textures. Important instances of this in single-objective minimization are seeking solutions that generalize to unseen data in supervised learning [1,2], in policy optimization [3], and generative models [4]. Many real-world systems are not so simple and instead involve multiple agents each of which uses a different subset of parameters to minimize their own objective. Some examples are generative adversarial networks (GANs) [5,6], actor-critic models [6], curriculum learning [7][8][9][10], hyperparameter optimization [11][12][13][14][15][16][17], adversarial examples [18,19], learning models [20][21][22], domain adversarial adaptation [23], neural architecture search [24][25][26][27][28], multi-agent settings [29] and meta-learning [30][31][32]. In these settings, the aim is to find one equilibrium (of potentially many equilibria) where agents exhibit some desired behavior.
For example, in the iterated prisoners' dilemma (Sec. 5), solutions favoring reciprocity over unconditional defection result in higher returns for all agents. In GANs, solutions often generate a subset of the modes from the target distribution [33], and in Hanabi, some solutions coordinate far better with humans [34]. Existing methods often find solutions in small subspaces -even after many random restarts, as shown in Table 1. By finding a diverse set of equilibria in these games, we may be able to (a) find solutions with a better joint outcome, (b) develop stronger generative models in adversarial learning, or (c) find solutions that coordinate better with humans.
Recently, Ridge Rider (RR) [35] proposed a general method for finding diverse solutions in single-objective optimization. RR is a branching tree search, which starts at a stationary point and then follows different eigenvectors of the Hessian ("ridges") with negative eigenvalues, moving downhill from saddle point to saddle point. In settings where multiple agents each have their own objective (i.e., games), the relevant generalization of the Hessian -the game Hessian [36] in Eq. 6 -is not symmetric. Thus, in general, the game Hessian has complex eigenvalues (EVals) with associated complex eigenvectors (EVecs), making RR not directly applicable.
In this paper, we generalize RR to multi-agent settings by leveraging machinery from dynamical systems. We connect RR with methods for finding bifurcation points, i.e. points where small changes in the initial parameters lead to very different learning dynamics and optimization outcomes. We propose novel metrics, inspired by Lyapunov exponents [37] that measure how quickly learning trajectories separate. These metrics generalize the branching criterion from RR, allowing us to locate a broad class of potential bifurcations, and enabling us to find more diverse behavior.
Starting points with rapid trajectory separation can be found via gradient-based optimization by differentiating through the exponent calculation, which is implemented efficiently for large models with Jacobian-vector products.
Our contributions include:
• Connections between finding diverse solutions and Lyapunov exponents, allowing us to leverage the broad body of work in dynamical systems. • Proposing a method, Generalized Ridge-Rider (GRR), that scales to high-dimensional differentiable games. • Presenting novel diagnostic problem settings based on high dimensional games like the iterated prisoners dilemma (IPD) and GANs to study a range of different bifurcation types. • Compared to existing methods, GRR finds diverse solutions in the IPD, spanning cooperation, defection and reciprocity. • Compared to RR, our GRR method also provides a more efficient estimator of the negative EVecs. • Lastly, we present larger-scale experiments on GANs -a model class of high interest to the ML community.
BACKGROUND
App. Table 4 summarizes our notation. Consider the single-objective optimization problem: * := arg min L ( )
We denote the gradient of the loss at parameters by := ( ) := ∇ L ( )| . We can locally minimize the loss L using gradient descent with step size :
+1 = −(2)
Due to the potential non-convexity of the L, multiple stationary points can exist, and gradient descent will only find a particular solution based on the initialization 0 . arXiv:2112.14570v1 [cs.GT] 24 Dec 2021
Ridge Rider
Ridge Rider (RR) [35] finds diverse solutions in single-objective minimization problems. The method first finds a saddle point, either analytically, e.g. in tabular reinforcement learning, or by minimizing the gradient norm. Then, RR branches the optimization procedure following different directions (or "ridges") given by the EVecs of the Hessian H = ∇ = ∇ (∇ L). Full computation of the eigendecomposition of H , i.e. its EVecs and EVals, is often prohibitively expensive; however, we can efficiently access a subset of the eigenspaces via Hessian-vector products H = ∇ ((∇ L) ) [38][39][40][41].
Optimization in Games
Instead of simply optimizing a single loss, optimization in games involves multiple agents, each with a loss function that can depend on other agents. For simplicity, we look at 2-player games with players (denoted by and ) who want to minimize their loss -L ( , ) or L ( , ) -with their parameters -or . * ∈ arg min L ( , * ), * ∈ arg min L ( * , )
If L and L are differentiable in and , we say the game is differentiable. One of the simplest optimization methods is to find local solutions by simply following the players' gradients, but this is unstable when the game Hessian has complex EVals [42]. Here, := ( , ) is an estimator for ∇ L | , , and the simultaneous gradient update is:
+1 = − , +1 = − (SimSGD)
We simplify notation with the concatenation of all players' parameters (or joint-parameters) := [ , ] ∈ R and the joint-gradient vector fieldˆ: R → R , denoted at the ℎ iteration by:
:=ˆ( ) := [ ( ), ( )] = [ , ](4)
We write the next iterate in (SimSGD) with fixed-point operator :
+1 = ( ) = −ˆ(5)
The Jacobian of the fixed point operator -denoted -is useful for analysis, including bounding convergence rates near fixed points [43] and finding points where local changes to parameters may cause convergence to qualitatively different solutions [44]. The fixed point operator's Jacobian crucially depends on the Jacobian of the joint-gradientˆ, which is called the game Hessian [45] because it generalizes the Hessian:
H := ∇ˆ= ∇ 2 L ∇ ∇ L ∇ ∇ L ⊤ ∇ 2 L(6):= ∇ ( ) = −Ĥ(7)
In Fig. 2 we show a game with a solution that we can only converge to by using an appropriate optimizer. Thus, we need to incorporate information about the optimizer when generalizing RR, which we do by working with the (largest) EVals/EVecs of instead of (most negative) EVals/EVecs ofĤ . We can understand the difference between optimization with a single and multiple objectives as follows: In single-objective optimization following the gradient forms trajectories from a conservative vector field becauseĤ = H is the Hessian of the loss which is symmetric and has real EVals. However, in games with multiple objectives,Ĥ can be non-symmetric and have complex EVals, resulting in a non-conservative vector field from optimization, opening the door to many new phenomena.
METHODS FOR GENERALIZING RR
Here, we present two key contributions for generalizing RR to games. We first connect diversity in optimization to the general concept of bifurcations, places where a small change to the parameters causes a large change to the optimization trajectories. Second, we introduce Lyapunov exponents [37] and easy-to-optimize variants as a tool for finding these aforementioned bifurcations.
Connecting Diversity and Bifurcations
In dynamical systems, bifurcations are regions of the parameter space where small perturbations result in very different optimization trajectories, and in general a dynamical system can contain a variety of different types of bifurcations. In contrast, in conservative gradient vector fields, saddle points are the only relevant class of bifurcations. As a consequence, their EVecs play a key role in the shape of the phase portraits, which are geometric representations of the underlying dynamics. In particular, the negative EVecs are orthogonal to separatrices [46], boundaries between regions in our system with different dynamical behavior, thus providing directions to move in for finding different solutions. This perspective provides a novel view on RR. RR branches at saddle points, the only relevant class of bifurcation points in single loss optimization.
However, in the dynamical systems literature, many bifurcation types have been studied [37]. This inspires generalizing RR to nonconservative gradient fields (e.g. multi-agent settings) where a broad variety of bifurcations occur. See Fig. 2 for a Hopf bifurcation [44] or Fig. 7 for various others.
Lyapunov Exponents for Bifurcations
Using tools from dynamical systems research we look at how to find general bifurcation points. Our objectives are inspired by the Lyapunov exponent, which measures asymptotic separation rates of optimization trajectories for small perturbations. We propose a similar quantity, but for finite length trajectories. Given a -step trajectory generated by our fixed point operator -i.e., optimizerwith initialization 0 and Jacobian at iteration of , we measure the separation rate for an initial, normalized displacement with:
( 0 , ) = 1 ∑︁ =0 ( 0 , ),(8)
where
( 0 , ) := log( ⊤ ( ( 0 )) ⊤ ( 0 ) )(9)
We call the ℎ Lyapunov term. When = 0, theˆare called the local Lyapunov exponents, while as → ∞ these are called the (global) Lyapunov exponents [47]. For a variable , we denote this as the -step or truncated Lyapunov exponent. Fig. 1 visualizes the exponents' calculation providing additional intuition. For notational simplicity, we suppress the dependency of on 0 going forward.
Within a basin of attraction to a given fixed point the global Lyapunov exponent is constant [37]. Intuitively, this is because an arbitrarily high number of Lyapunov terms near the fixed point dominate the average defining the exponent in Eq. 8. This property prevents us from optimizing the global exponent using gradients, making it a poor objective for bifurcations. As such, our interest in the truncated exponent is motivated from multiple directions: Figure 1: Visualization of the components to a Lyapunov ex-ponentˆ( 0 , ) described in Eqs. 8, 9, which measures how quickly trajectories separate starting at a point 0 in direction . Here, the optimization trajectory iterates accumulate at a fixed point * . We show two displacements -1 and 2 -resulting in separate "perturbed" trajectories shown in and . We measure the separation rate between the true and perturbed trajectories at the ℎ optimizer iteration with the Lyapunov term ( 0 , ). The exponentˆ( 0 , ) is the average of the first terms. See Fig. 6 for actual trajectories on a toy problem used in an exponent calculation.
(1) Non-zero gradient signals for finding bifurcations (2) Computationally tractability (3) A better separation rate description for the finite trajectories used in practice However, unlike the global exponent, the truncated version lacks theoretical results.
In more than one dimension, the -step Lyapunov exponent is a function of the specific direction of the perturbation [46]. We look at using the direction for maximal separation -i.e., the max -step Lyapunov exponent:
( 0 ) = max , ∥ ∥=1ˆ( 0 , )(10)
For dynamical systems with basins of attraction, common in optimization, the max exponent is largest on the boundary between basins, which motivates maximizing the max exponent to find bifurcations. The max exponent can be evaluated by finding the largest EVal of an average of the Jacobians over the optimization steps [37]:
† := 1 ∑︁ =0 ( ) ⊤ ( ),(11)( 0 ) = max ∈Sp( † ) | |(12)
Importantly, in higher dimensions one eigenvalue dominates the spectrum of after a large number of steps [48,49] and is thus a point of maximal exploration in our solution space. We note some practical points for computing these exponents: when = 0 the max exponent is the max eigenvalue of 0 . As → ∞ and our fixed point operator converges to a fixed point * , the max exponent is the max EVal of at * . Calculatingˆis easiest when = 0 or → ∞, e.g., by power iteration on the relevant . For intermediary , directly using leading EVals of † involves reevaluating the entire optimization trajectory many times. Instead, it is often easier to work with bounds. A simple lower bound is formed by using the leading EVec at any single step, or an upper bound by using the leading EVec at each step, which are tight as → ∞ [48]. We investigate these strategies in App. Fig. 9. Commonly, our goal is to obtain many qualitatively different solutions from a single starting point, which motivates simultaneously optimizing the exponents corresponding to multiple different directions. Relatedly, the sum of positive global Lyapunov exponents gives an estimate of the Kolmogorov-Sinai or metric entropy by Pesin's theorem [50]. We use this to motivate different performance metrics in Section 4.2.
PROPOSED ALGORITHMS
Having given an overview of the key mathematical concepts, we now present our overall algorithm. First, we introduce a general branching-tree search framework for finding diverse solutions in differentiable games. Next, we present our method -Generalized Ridge Rider (GRR) -which implements this framework using truncated Lyapunov exponents (Eq. 10) as the branching criterion. Lastly, we highlight the differences between GRR and RR.
Branching Optimization Tree Searches
Our framework is a generalized version of RR and contains the following components:
(1) A method for finding a suitable starting point for our branching process -see Fig. 6. (2) A process for selecting branching directions (or perturbations) from a given branching point -see Fig. 4. (3) A prescription for how to continue the optimization process along a given branch after the initial perturbation. (4) A re-branching decision rule, i.e., when to go back to step (2). This was important in RR because optimizers in highdimensional non-convex ML problems often finish at saddle points [51]. (5) Lastly, a metric to rank the different solutions. We visualize this process in Fig. 3. RR is an instance of this general process, where each component is suitable for single-objective optimization. In the next section, we present another instance of this method, designed for optimization in games. We include a more detailed description of branching tree searches in App. Alg. 1, highlighting the important changes compared to RR.
Generalized Ridge Rider (GRR)
Starting point: Motivated by Section 3.2, we look at optimizing the maximal -step Lyapunov exponent from Eq. 10 to obtain our starting point:
L ( 0 ) = −ˆ( 0 ) = − max , ∥ ∥=1ˆ( 0 , )(13)
However, using a single exponent only guarantees trajectory separation in a single direction. If we want to branch across multiple bifurcations in different directions, we need an objective using exponents in multiple, different directions. We look at the simple objective choice summing over exponents:
L sum ( 0 ) = − max 1 ,..., ∑︁ =1ˆ( 0 , ),(14)
such that ∥ ∥ = 1, ⊤ = 0 for all , ∈ 1, . . . , , ≠
Intuitively, the constraint guarantees we have different directions to separate in, by making them orthogonal. It is straightforward to evaluate this objective, by evaluating the top-EVals of the matrix from Eq. 11. More generally, convex functions of the -step exponents in different directions form reasonable objectives that are more amenable to optimization. Specifically, we also look at:
L min ( 0 ) = − max 1 ,..., min =1...ˆ( 0 , )(16)such that ∥ ∥ = 1, ⊤ = 0 for all , ∈ 1, . . . , , ≠(17)
Branching the parameter optimization:
We must choose what direction to branch in; our procedure for evaluating Lyapunov exponent objectives creates natural candidates. Specifically, evaluating the max exponent involves finding the direction maximizing trajectory separation, which we re-use for branching. Notably, this is the most negative EVal of the Hessian if we start at a saddle point and use SGD when calculating the trajectories, generalizing the choice from RR. For each direction, we can move in both a positive and negative direction, giving two branches.  Fig. 6. In each setup, we have two EVecs and branch in opposite directions, giving four paths, displayed in different colors. Steps with the eigenvector have magenta circles marking boundaries. Top: In the small IPD, finding, then branching at a saddle -where the joint-player grad. lognorm log(∥ˆ∥) is 0 -allows us to find defect-defect and titfor-tat solutions. Bottom: In the Mixed Problem of Small IPD and Matching Pennies, branching at the Hopf bifurcation allows us to find both solutions. Here, there are no saddle points near the bifurcation, so RR's starting point does not allow branching to find both solutions.
Also, we must choose how far to move in the directions. If we move too far, we may leap into entirely different parts of the parameter space -e.g., missing interesting regions and recovering similar solution modes. If we are not exactly at a bifurcation -only near it -then we may need to move some minimum distance to cross the separatrix and find a new solution. We look at two simple strategies to move sufficiently far. First, we try taking a single step with the normalized exponent direction scaled according to the exponent. Second, we look at taking small steps in the exponent direction until the alignment with the joint-gradient flips, which generalizes RR's "riding a ridge" (following an EVec of the Hessian) while it is a descent direction.
Optimizing each branch: For optimization in games, the stability properties of solutions can crucially depend on optimizer choice [52]. One should choose an optimizer suited to the problem. In our experiments, we use Learning with Opponent Learning Awareness (LOLA) [29] which can converge to periodic solutions and is attracted to high-welfare solutions in the IPD. In Section 6.1.1 we contrast finding diverse solutions using LOLA with simultaneous SGD (simSGD) -a method that works well for single-objective optimization, but cannot find periodic solutions. App. A summarizes other optimizer choices.
Re-branching: In single-objective optimization in ML, our optimizer often finishes at a high-dimensional saddle [51], which makes re-branching important. Specifically, we can re-branch at the saddle in negative EVec directions to try to find critical points with less negative EVecs. In our setup, we are interested in re-branching if our optimizer finishes at a point where EVals of the Jacobian of the fixed point operator are greater than 1. These are directions where our optimizer will continue moving the parameters. [53] observed EVals of larger than 1 at the end of GAN training, indicating that we may want to re-branch for games in machine learning.
Comparing GRR and RR
RR is a branching optimization search specifically for single-objective optimization -which is less general than optimization in games -so it can outperform GRR. E.g., non-conservative systems have more bifurcation types than conservative ones. If we are only concerned with saddle bifurcations, we can just find a saddle stationary point by minimizing the gradient norm. We know that this (relatively) easy-to-find stationary point lies on the separatrix. However, Hopf bifurcations are not necessarily near stationary points. Thus optimizing gradient norms does not work in general, while optimizing a Lyapunov exponent does (Fig. 2).
While it might be overkill to find a separatrix with Lyapunov exponents in a single-objective setting, we take some lessons from GRR back to RR. It is useful to view RR as a method for finding bifurcations and branching across them. This motivates ways to sort between different stationary points to start at -an open problem from RR. For example, using the point with the largest Kolmogorov-Sinai entropy [50]. At stationary points, this is simply the (negative) sum of negative EVals. Another limitation of RR is effectively estimating the most negative EVals of the Hessian. It is often simplerin computation and implementation -to estimate the leading EVals of the Jacobian of the fixed point operator instead of the most negative eigenvalues of the HessianĤ . In Section 6.2.3 we show that our method saves Hessian-vector product evaluations when estimating EVecs in setups from RR.
EXPERIMENTAL SETTING
We experimentally investigate GRR on a variety of problems summarized in this section and described in detail in App. C.1. We chose these as they cover different types of dynamics and contain different kinds of bifurcations. Some are standard benchmarks, while others -i.e., Random Subspaces -are novel to this work. We also summarize our gradient computation for these problems.
Matching Pennies is a simplified 2-parameter version of rockpaper-scissors. This problem's game Hessian has purely imaginary EVals unlike the small IPD, but only a single solution. Thus, by itself, is a poor fit for evaluating methods for a diversity of solutions, but nevertheless a useful test when probing GRR behavior.
The Iterated Prisoners' Dilemma (IPD) is the discounted, infinitely iterated Prisoner's Dilemma [54]. Each agent's policy conditions on the actions in the prior time step, so there are 5 parameters for each agent -the probability of cooperating initially and those given both agents' preceding actions. There are several different relevant equilibria in the IPD, including unconditional defection (DD), leading to the worst-case joint outcome, and tit-for-tat (TT), where agents initially cooperate, then copy the opponents' action (giving a higher reward). We turn the IPD into a differentiable game by calculating the analytical expected return as a function of the joint policy of the two agents.
The Small IPD is a 2-parameter simplification of IPD, which has both DD and TT Nash equilibria, allowing us to visualize optimization difficulties from the full-scale IPD. However, the game Hessian has strictly real EVals, unlike the full-scale IPD.
Mixing Small IPD and Matching Pennies interpolates between the Small IPD and Matching pennies games with an interpolation factor ∈ [0, 1]. This problem has two solutions -one where both players cooperate and one where both players select actions uniformly, with a Hopf bifurcation separating these.
Generative Adversarial Networks (GANs): We use a setup from [55,36,56], where the task is to learn a Gaussian mixture distribution using GANs. The data is sampled from a multimodal distribution to investigate the tendency to collapse on a subset of modes during training -see App. Fig. 16 for the ground truth.
Random Subspace IPD/GAN: To see how robustly we can find bifurcations with the exponents, we construct more complicated toy problems by taking higher-dimensional problems and optimizing in a random subspace. For each player, we select a random direction to optimize in, by sampling a vector v with entries from [0, 1] and normalizing it. Additionally, we select a random offset b from whatever an appropriate initialization is for the higher-dimensional problem. So, the first player controls -coordinate and optimizes the loss L (v + b , v + b ), while the second player controls the -coordinate and optimizes
L (v + b , v + b ).
Single-objective problems: We apply our method to find bifurcations on single-objective optimization problems. There are various relevant problems in machine learning, but we focus on comparisons with RRs EVec estimation in MNIST classification.
Optimizing the starting point objective: To optimize these objectives, we use automatic differentiation libraries (like Jax [41] or PyTorch [40]) to compute gradients through methods that calculate our Lyapunov exponent-inspired objectives. The scalability of this approach depends on the implementation of our exponent calculation, which can depend on estimating the top eigenvalues of the positive semi-definite (PSD) symmetric matrix in Eq. 11. In simple settings we can differentiate through the full spectrum calculation via jax.linalg.numpy.eigh; we investigate this on the toy experiment in Sec. 6.1.2 and the IPD in Sec. 6.2.1. However, in ML, the matrix from Eq. 11 is typically too large for the full spectrum. Directly estimating the top EVals with an iterative method allows us to (automatically) differentiate through them. We differentiate through jax.numpy.linalg.eigh in Fig.6 and investigate using power iteration with Hessian-vector products in App. Fig. 9. 
EXPERIMENTAL RESULTS
First, in Sec. 6 Here, we give a system with complex EVecs showing (a) the importance of selecting a convergent optimizer in GRR, and (b) an example task where RR cannot be applied. Fig. 2 shows the phase portrait for baseline methods on our Mixed Problem. LOLA (and other game optimizers) can find both solutions, while naïvely following the gradient always finds a single solution. 6.1.2 Starting Point Selection. Fig. 6 shows the effect of optimizing the starting point for the max 10-step Lyapunov exponent on the Mixed Problem. We find gradient-based optimization can find bifurcations. Next, Fig. 9 contrasts different direction choices for the exponent calculation. We find that re-estimating the top EVecs at each iteration performs best, though the simple methods also work.
App. Fig. 10 shows the max -step exponent for multiple numbers of steps , showing that a moderate number of steps -e.g., 10 -allows us to find bifurcations. App. Fig. 15 shows different Lyapunov exponent objectives, trying to guarantee trajectory divergence in multiple directions. We can find bifurcations while guaranteeing trajectory separation in every direction. Impact of inner optimizer choices on bifurcation structure: App. Fig. 11 contrasts the exponents for LOLA and simSGD, showing that we find optimizer-dependent bifurcations. App. Fig. 12 investigates the impact of optimization-algorithm parameter choices on bifurcation structure. This shows that if the step size is too large, the optimizer does not converge, resulting in bifurcations between complicated limit cycle trajectories [57], and making GRR difficult to apply.
Starting points on single-objective problems: App. Fig. 13 investigates our algorithm in single-objective problems, showing that our method finds bifurcations in the same setup as RR. App. Fig. 14 shows our method on the logistic map, giving intuition for our method on a canonical example for bifurcations. Fig. 4 we demonstrate branching at bifurcations to find multiple solutions to toy problems. This shows the branching process, and an explicit example where RR's starting point does not work, but GRR's does. The small IPD has a saddle bifurcation, while the Mixed Problem has a Hopf bifurcation. Fig. 7 we look at calculating Lyapunov exponents on toy problems with more complicated bifurcation structures. We create a variety of more complex toy problems by taking the high-dimensional IPD and GAN problems and selecting a random subspace to optimize in. We can effectively highlight bifurcations in this setup.
Branching at Bifurcations. In
A Range of Complicated Toy Problems. In
Scaling the Results
Optimizing Lyapunov Exponents on IPD.
Here, we investigate our ability to use gradient-based optimizers on Lyapunov exponents in the IPD. Fig. 5 shows the feasibility of using gradient descent to tune the 1-step max Lyapunov exponent. App. Fig. 18 optimizes an objective using multiple exponents, showing that we effectively optimize multiple exponents, which gives trajectory separation in multiple directions. App. Fig. 19 compares objectives using multiple exponents, showing that using the minimum of the top exponents gives trajectory separation in all directions, unlike the naïve choice of optimizing their sum. The sum of exponents finds solutions separating extremely fast in the top directions, while (slowly) converging in the bottom directions. In contrast, the min of the exponents does not allow convergence in the bottom directions. App. Fig. 20 compares optimizing the -step max Lyapunov exponent for variable , showing that we effectively minimize multi-step exponents in higher-dimensional problems if required.
Player 1 Strategy
Player 2 Strategy
Optimizing the exponent Figure 6: Calculation and optimization of a max 10-step Lyapunov exponent from Eq. 10 on the mixed small IPD and Matching Pennies problem. Gradient-based optimization on this objective effectively finds the bifurcation. Top: We show a heatmap of the exponent, and visualize the calculation of each exponent in two regions. This involves simulating 10step trajectories shown in red starting at the black points, then finding a direction that maximizes trajectory separation. We use this exponent to find bifurcations -in this case between the solution in the top right and the center. Bottom: We show optimization trajectories for gradient ascent on the exponent for a grid of initializations. For a variety of starting points, the optimization procedure finds large exponent locations (final iterate shown with red circles).
GRR Applied to the IPD.
Here, we use our method on the IPD, where existing methods have difficulty finding diverse solutions. There are two solution modes: ones where both agents end up defecting and cooperating respectively. Table 1 compares our method to baselines of following gradients and LOLA, each run with random initializations. Our method finds both solutions modes, unlike existing approaches. We found that it was sufficient to use the max Lyapunov exponent as our objective, which only guarantees separation in 1 direction. Similarly, we found that it was sufficient to use a 1-step or local Lyapunov exponent objective, though we may require more steps to find bifurcations in other problems.
Player 1 Strategy
Max 10-step Lyapunov Exponent
Player 2 Strategy
Random subspace test problems Discriminator Generator Figure 7: We show the Lyapunov exponent heatmap (as in Fig. 6) on more complicated toy problems to see how robustly we can find different bifurcations. The exponent peaks near where trajectories (shown in red) separate, showing that we find various bifurcations. See App. Fig. 17 for other sampled subspaces. Sec. 5 describes how we construct these examples by taking higher-dimensional problems and optimizing them in a random subspace. Top: An IPD subspace with multiple Hopf bifurcations. Bottom: A GAN subspace with various bifurcations.
Improving RR's EVec Estimation.
We investigate efficiently finding the most negative EVecs in RR by estimating the largest EVecs of the Jacobian of our fixed point operator. We measure efficiency by comparing the number of Hessian-vector product (HVP) evaluations because HVP evaluations dominate the cost of EVec estimation here. Table 2 shows how many HVP evaluations we require to reach different MNIST classifier accuracies by following EVecs. Our method can more efficiently use HVP evaluations than the RR method because we do not need to repeatedly re-estimate the most negative EVal. We stress that this problem is not designed to train a single, strong classifier; it is easy to simply train our network by following the gradient to 100% train accuracy. This problem was selected from RR's experiments because it requires us to accurately and efficiently estimate negative EVecs many times. A downstream use of this is training an ensemble of classifiers for generalization. . Takeaway: Our method finds solutions at both loss modes, while existing approaches of using random initializations, then following the gradient or using LOLA do not find diverse solutions. The IPD has two solution modes -i.e., solutions where both agents end up defecting with a loss of 2 and where both agents end up cooperating with a loss of 1 (like tit-for-tat). We assess which modes were found by showing (P)layer 1's strategy, which is the chance of (C)ooperating given both players' last action -ex., ( | ) is the chance if previously P1 defected and P2 cooperated. We compare GRR flavors with just following gradients via SimSGD and LOLA [29] from random (init)ializations. We compare with 20 random inits because GRR follows at most 20 branches, and because we have 10 EVecs in either direction (+/-). GRR only branches in directions where EVals of the Jacobian of the fixed point operator are greater than 1 (i.e., trajectories locally diverge) as visualized in Fig. 5 (middle). We look at the impact of starting at an approximate bifurcation in GRR, by branching on the EVecs at a random init. If the max Lyapunov exponent is not tuned, then each branch finds the same solution.
MNIST Accuracy # HVP Evaluations Our method Method from RR 10 000 19%(+8%) 11% 100 000 89%(+6%) 83% 1 000 000 93%(+2%) 91% 
Calculating Lyapunov Exponents for GANs.
Here, we investigate scaling our exponent calculations to machine learning models where the (game) Hessian is so large we cannot materialize it and can only use Hessian-vector products. Specifically, we use the GAN described in Section 5. We look at calculating our exponent for various hyperparameters and random re-starts. We evaluate the quality of using our exponent to find diverse solutions, by calculating the log-probability of samples from an ensemble of GANs from the top 5 optimization branches. Table 3 shows the mean and standard deviation (over 10 random restarts) of the max 10-step Lyapunov exponent and the resulting ensemble's log-probability. Each GAN was trained for 10 000 updates, so evaluating each ensemble cost approximately 50 000 evaluations of both players' gradients. In contrast, each exponent cost less than 1000 evaluations of both gradients to compute. This shows that we effectively scale our exponent calculation to larger models of interest from machine learning, and find that a large (mean) exponent aligns with regions where we can branch to train the strongest ensemble of GANs.  Table 3: We display the mean and standard deviation (over 10 random restarts) of the max 10-step Lyapunov exponent and the log-probability of an ensemble of 5 GANs obtained by branching in the top 5 directions at the initialization. We show that the better performing ensembles also have higher Lyapunov coefficients as well as demonstrating that our exponent calculation is scalable to larger problems. The best GANs log-prob. from the best ensemble was −12 861 ± 356, which is worse than the ensemble's performance of −10 846± 256. This indicates that each GAN may be learning a different part of the data distribution (samples in App. Fig. 16).
CONCLUSION
In this paper we introduced Generalized Ridge Rider, an extension of the Ridge Rider algorithm to settings with multiple losses. We showed that, in these settings, a broader class of bifurcation points needs to be considered, and that GRR indeed discovers them in a variety of problems. Experimentally, we isolate each component of GRR demonstrating their effectiveness, and show that -in contrast to baseline methods -GRR obtains a diversity of qualitatively different solutions in multi-agent settings such as the iterated prisoner's dilemma. We also provide empirical justification for our method by using tools from the dynamical systems literature, allowing us to find arbitrary bifurcations. This hints at a multitude of approaches and tools from dynamical systems, that can be used for understanding game dynamics and learning diversity.
Diversity in machine learning. Finding diverse solutions is often desirable in machine learning, for example improving performance for model ensembles [65], with canonical approaches directly optimizing for negative correlation amongst model predictions [66]. In recent times these ideas have begun to re-emerge, improving ensemble performance [67][68][69], robustness [70,3] and boosting exploration in reinforcement learning [71][72][73][74].
Many of these approaches seek to find diverse solutions by following gradients of an altered, typically multi-objective loss function. By contrast, the recent Ridge Rider (RR, [35]) algorithm searches for diverse solutions by following EVecs of the Hessian with respect to the original loss function, producing orthogonal (loss reducing) search directions.
Finding solutions in games. There are first-order methods for finding solutions in games including extragradient [75,76], optimistic gradient [77,78], negative momentum [52], complex momentum [53], and iterate averaging [79]. There are also higherorder methods like consensus optimization [80], symplectic gradient adjustment (SGA) [45], local symplectic surgery (LSS) [81], competitive gradient descent (CGD) [82]. Recently, [83] looked at the effects of overparameterization while doing online learning in games.
Diversity in multi-agent RL. In recent times a series of works have explore the benefit of diversity in both competitive [9,84,85] and cooperative [86,87] multi-agent RL. However, once again these approaches all consider augmented loss functions. Instead, we take inspiration from RR and extend it to the multi-agent setting.
Tree searches in MDP. Tree searches are a classic technique for planning and control in discrete Markov Decision Processes (MDP). In these settings, the locations to branch from, and the potential branch choices are already given, and all that remains is choosing which to explore. Much interest has been developed in performing these searches ranging from depth/breadth first search, A* [88], to more sophisticated methods such as Monte Carlo Tree Search [89] potentially with learned value functions [90]. Searching in continuous action spaces has been explored in [91][92][93]. All of these methods search in the action space of the agents. In contrast, we seek to employ these search methods to find the parameters of the agents.
Rapidly-exploring random tree is an algorithm that builds a space-filling tree to cover a given continuous optimization space [94,95]. Unlike RR-based methods, the branching points are not determined by the underlying properties of the loss landscape (e.g. only at saddles, or bifurcation). This technique is typically used in robotic motion planning [96].
Branching in evolutionary optimization. Designing optimization algorithms in a multi-modal loss landscape has been the focus of the evolutionary optimization community [97,98]. These methods implicitly build a tree over candidate solutions. Evolutionary algorithms specifically to encourage diversity have been explored in the context of multi player games [9,85,99] by encouraging populations of agents to learn different strategies which perform well against other agents resulting in a growing collection of strategies.
Bifurcations. This work builds on the prior workshop submission of [100]. The works of [101][102][103][104][105][106] leverage the bifurcations for learning in games. [102] introduces two mechanisms -hysteresis and optimal control mechanisms -to control equilibrium selection, with the goal of improving social welfare. We do not focus on maximizing social welfare and instead focus just on finding multiple solutions. [103] studies the properties of bifurcations in routing games, showing various analyses for the social cost. [104] combines catastrophe theory with mechanism design to destabilize inefficient solutions. [105] studies how changing an exploration parameter in multi-agent learning can affect equilibrium selection. Our work focuses on connecting bifurcations with finding diverse solutions in optimization with Ridge Rider, and finding those bifurcationspotentially with gradient-based methods.
Lyapunov exponents. The works of [107][108][109] use Lyapunov exponents when learning in games. [107] shows that some learning algorithms are Lyapunov chaotic in payoff space. Our work focuses on using Lyapunov exponents variants to identify bifurcations (which are used to find diverse solutions in optimization), and potentially using optimization methods on the exponent.
Separatrices. The works of [110][111][112] look at computing separatrices in games which are leveraged for various purposes. [112] finds the shape of separatrices using invariant functions with online learning dynamics.  Simultaneous SGD := Defined to be equal to , , , · · · ∈ C Scalars , , , · · · ∈ C Vectors , , , · · · ∈ C × Matrices ⊤ The transpose of matrix The identity matrix ℜ( ), ℑ( )
The real or imaginary component of ∈ C The imaginary unit. ∈ C =⇒ = ℜ( ) + ℑ( )
The complex conjugate of ∈ C | | := √¯T he magnitude or modulus of ∈ C arg( )
The argument or phase of ∈ C =⇒ = | | exp( arg( )) ,
A symbol for the outer/inner players , ∈ N The number of weights for the outer/inner players A symbol for the parameters or weights of a player ∈ R , ∈ R The outer/inner parameters or weights L : R → R A symbol for a loss L ( , ), L ( , )
The outer/inner losses -R + ↦ → R The joint-gradient vector field at weights ∇ˆ:= ∇ˆ| ∈ R × The Jacobian of the joint-gradientˆat weightŝ H
The game Hessian ∈ C
The step size or learning rate ∈ C, Notation for an arbitrary Eval Sp( ) ∈ C
The spectrum -or set of eigenvalues -of ∈ R × ( ) := max ∈Sp( ) | |
The spectral radius in R + of ∈ R × ( ) Fixed point operator for our optimization The Jacobian of the fixed point operator A displacement for a Lyapunov exponent ( 0 , ) := log( ⊤ ( ( 0 )) ⊤ ( 0 ) ) A Lyapunov term for a Lyapunov exponent
( 0 , ) = 1 =0 ( 0 , ) A -step Lyapunov exponent ( 0 ) = max , ∥ ∥=1ˆ( 0 , )
The max -step Lyapunov exponent L ( 0 )
A starting point loss using Lyapunov exponents L sum ( 0 )
A loss using the sum of top exponents L min ( 0 )
A loss using the min of top exponents rewards for each player summed via ∞
=0
. Each agent is conditioned on the actions in the prior state ( ). Thus, there are 5 parameters for each agent : ( | ) is the probability of cooperating at start state 0 = ∅ or state = ( 1 −1 , 2 −1 ) for > 0. There are two Nash equilibria which we are interested in: Defect-Defect (DD) where agents are selfish (resulting in poor reward), and tit-for-tat (TT) where agents initially cooperate, then copy the opponents' action (resulting in higher reward).
Small IPD: This is a 2-parameter simplification of IPD, which allows DD and TT Nash equilibria. We fix the strategy if our opponent defects, to defect with high probability. We also constrain the probability of cooperating to only depend on if the opponent cooperates, and in the initial state we assume our opponent cooperated. This game allows us to visualize some of the optimization difficulties for the full-scale IPD, however, the game Hessian has strictly real EVals unlike the full-scale IPD. See Fig 2 top for a visualization of the strategy space.
Matching Pennies: This is a simplified 2-parameter version of rock-paper-scissors, where each player selects Cooperate or Defect. This game has a Nash equilibrium where each player selects its action with uniform probability. Notably, this problem's game Hessian has purely imaginary EVals, so following the gradient does not converge to solutions and we need a method for learning in games like LOLA. Also, this game only has a single solution; thus it is a poor fit for evaluating RR, which finds a diversity of solutions. See Fig 2, bottom for a visualization of the strategy space.
Mixing Small IPD and Matching Pennies: This game interpolates between the Small IPD and Matching Pennies games, with the loss for player :
L , , = L , + (1 − )L ℎ ,(18)
.  Takeaway: Re-estimating the top EVecs at each iteration performs best, but is most expensive. Left: We sample a random normalized direction uniformly for our displacement at each exponent calculation, which does not clearly show the bifurcation. Middle: We perform 10 steps of power iteration to tune . First, we tune the displacement at only the first iteration, which shows the bifurcation. Next, we tune the displacement at every iteration in the exponent calculation, which shows the bifurcation more clearly. Right: We use jax.numpy.linalg.eigh to tune , which also clearly shows the bifurcation. Figure 10: We show the calculation for the -step Lyapunov exponent for various on the small IPD. The -step optimization trajectories used for the exponent calculation are shown in red, allowing us to see the horizon that the associated exponent measures separation over. Takeaway: We can effectively find bifurcations with a -step exponent for various . Using only 1step does not show the bifurcation, while using more does. As gets larger, the scale of the gradients changes, causing difficult optimization in the limit -note the changing color bar scale. Step size = 10 Player 1 Strategy
Player 2 Strategy
Step size = 20    Figure 7 with more random subspaces. The optimization trajectories used for the exponent calculation are shown in red, allowing us to see the trajectories the associated exponent measures separation over. Takeaway: The exponent is peaked near where trajectories separate for each subspace, showing that these strategies can find various bifurcations. We display the Lyapunov exponent calculation -as in Fig. 6 -on more complicated toy problems, to see how robustly we can find different bifurcations. Section 5 describes how we construct these examples by taking higher-dimensional problems and optimizing in a random subspace. Top: Different IPD subspace: we effectively highlight bifurcations -i.e., regions where trajectory behavior qualitatively changes. Bottom: Different GAN subspaces: we are able to find bifurcations, but the highlighted structure is less crisp in this more complex toy example. Our objective is the minimum of the top local Lyapunov exponents, with a different 's optimization trajectory shown in each row. To review this visualization also in Figure 5: The spectrum is shown with a scatter-plot in blue, with a progressively larger alpha at each iteration. The final spectrum is shown in red. For the Jacobian of the fixed-point operator , a vertical red line is shown where the EVal norm equals 1, signifying the cutoff between (locally) convergent and divergent eigenspaces.
Top: Optimizing the max local Lyapunov exponent from Figure 5, where the largest EVal is effectively maximized. Middle: Optimizing the top 5 local exponents, where we find 5 moderately divergent directions instead of 1 extremely divergent direction as in the top. Bottom: Optimizing the top 10 local exponents, which results in trajectory separation for every direction. We study different objectives of the top exponents in Appendix Figure 19, and study non-local -step exponents in Appendix Figure 20.  Figure 19: We compare the optimization for different losses using the top 10 local Lyapunov exponents. Takeaway: Using the minimum of the top exponents is a strong method to guarantee separation in many directions. To review this visualization also in Figure 5: The spectrum is shown with a scatter-plot in blue, with a progressively larger alpha at each iteration. The final spectrum is shown in red. For the Jacobian of the fixed point operator , a vertical red line is shown where the EVal norm equals 1, signifying the cutoff between (locally) convergent and divergent eigenspaces. Top: We optimize the sum of the top 10 exponents. This is a simple choice, but it does not guarantee trajectory separation in every direction. Our optimizer is happy with solutions that are diverging rapidly in the top directions, while converging (slowly) in the bottom directions. Bottom: To guarantee trajectory separation is every direction, we look at optimizing the minimum of the top exponents. Here, at the end of training, all EVals of are greater than 1, signifying local trajectory separation in every direction.
loss L (
0 ) = −ˆ1 0 ( 0 )
optimization iteration Figure 20: We investigate optimizing a loss of the max -step Lyapunov exponent for > 1. Takeaway: We are able to effectively minimize multi-step exponents in higher-dimensional problems if required. To review this visualization also in Figure 5: The spectrum is shown with a scatter-plot in blue, with a progressively larger alpha at each iteration. The final spectrum is shown in red. For the Jacobian of the fixed point operator , a vertical red line is shown where the EVal norm equals 1, signifying the cutoff between (locally) convergent and divergent eigenspaces.
Correspondence to: Jonathan Lorraine lorraine@cs.toronto.edu Proc. of the 21st International Conference on Autonomous Agents and Multiagent Systems (AAMAS 2022), P. Faliszewski, V. Mascardi, C. Pelachaud, M.E. Taylor (eds.), May 9-13, 2022, Online. © 2022 International Foundation for Autonomous Agents and Multiagent Systems (www.ifaamas.org). All rights reserved.
Figure 2 :
2The phase portrait for two standard optimization algorithms on the mixed small IPD and Matching pennies problem. We show trajectories following the gradient with simSGD in red and LOLA[29] -a method for learning in games -in blue. All initializations of SimSGD only find the solution in the top right, because the center solution has imaginary EVals, while LOLA finds all solutions. For comparisons over more test problems see App.Fig. 8.
Figure 3 :
3Visualization of branching optimization tree search. The key components are: (1) selecting the starting point, (2) creating different branches, (3) optimizing each branch, and (4) choosing when to re-branch.
Figure 4 :
4We show branching at different types of bifurcations, obtained by optimizing a Lyapunov exponent as shown in
Figure 5 :
5We display gradient descent optimization on the 1-step max Lyapunov exponent objective (Eq. 13) on the IPD. Takeaway: We effectively reduce our loss and correspondingly raise the max EVal of . Left: We display our loss -i.e., the negative Lyapunov exponent objective-as optimization progresses. Middle: We visualize the spectrum of the Jacobian of our fixed point operator in log-polar coordinates as optimization progresses. The spectrum is shown with a scatter-plot in blue, with a progressively larger alpha at each iteration. The final spectrum is shown in red. A vertical red line is shown where the EVal norm equals 1, signifying the cutoff between (locally) convergent and divergent eigenspaces. We effectively maximize the norm of the largest EVal. Right: We display the spectrum of the game Hessian. A horizontal red line is shown where the real part of the EVal transitions from negative to positive, signifying the cutoff between (locally) convergent and divergent eigenspaces under gradient flow. Log-polar coordinates are required to see structure in the spectrum.
the set of branches B = SplitBranch( ) 8: Initialize the set of solutions S = ∅ 9: while Branches B non-empty do 10: , B = ChooseBranch(B) 11: * = Optimize( . , . ) Prisoners' Dilemma (IPD): This game is an infinite sequence of the Prisoner's Dilemma, where the future payoff is discounted by a factor ∈ [0, 1). In other words, the sequences of
( , ), ( , ) Gradient of outer/inner losses w.r.t. their weights in R / * ( ) := arg minL ( , ) The best-response of the inner player to the outer player L * ( ) := L ( , * ( )) The outer loss with a best-responding inner player * := arg min L * ( ) Outer optimal weights with a best-responding inner player := + The combined number of weights for both players := [ , ] ∈ R A concatenation of the outer/inner weightŝ ( ) := [ ( ), ( )] ∈ R A concatenation of the outer/inner gradients 0 = [ 0 , 0 ] ∈ R The initial parameter values An iteration number :=ˆ( ) ∈ R
Figure 9 :
9We compare different methods for choosing a direction in the max 10-step Lyapunov exponent calculation on the Mixed Objective.
Figure 11 :
11We contrast the max 20-step Lyapunov exponent calculations between LOLA and simSGD. The 20-step optimization trajectories used for the exponent calculation are shown in red, allowing us to see the trajectories the associated exponent measures separation over. Takeaway: We can find optimizer-dependant bifurcations with our method. Top: 10 steps from optimization with LOLA at each point where we calculate an exponent for the heatmap. Bottom: The same visualization as top, except with an optimizer of simSGD.
Figure 12 :Figure 13 : 0 Figure 14 :
1213014We investigate the impact of optimization algorithm parameters on the 10-step max Lyapunov exponent calculation with LOLA. The 10-step optimization trajectories used for the exponent calculation are shown in red, allowing us to see the trajectories the associated exponent measures separation over. Takeaway: If the step size is too large, the optimizer does not converge to solutions, resulting in complicated limit cycle bifurcations. Top: 10 steps from the optimization with a step size of = 10 at each point where we calculate an exponent for the heatmap. Bottom: The same visualization as top, except with step size of = 10. The optimization trajectories only converge at the solution in the top right. In the center, the trajectories accumulate at a limit cycle, rotating around the solution.Max 10-step Lyapunov Exponent We investigate finding bifurcations with a 10-step max Lyapunov exponent on a single-objective optimization problem from RR. The optimization trajectories (used for the exponent calculation) are shown in red, allowing us to verify where bifurcations are. Takeaway: Our method effectively highlights bifurcations in RRs example. We display the Lyapunov exponent on the logistic map: ( + 1) = ( ) + + ( ) 2 . Takeaway: Intuition for Lyapunov exponents on a canonical 1-dimensional example for bifurcations.
Figure 15 :
15We compare calculating different 10-step Lyapunov exponent objectives trying to guarantee divergence in multiple directions on the small IPD. The optimization trajectories used for the exponent calculation are shown in red, allowing us to see the trajectories for the associated objective. Takeaway: We can find bifurcations, while guaranteeing trajectory separation in every direction. Interestingly, local maxima -not saddles -allow trajectory separation in all directions here. Top: The max 10-step Lyapunov exponent. Bottom: The sum of the top 2 10-step Lyapunov exponents.
Figure 16 :Figure 17 :
1617Ground truth samples for our Mixture of Gaussian experiment. This reproduces
Figure 18 :
18We compare the optimization for different losses using the local Lyapunov exponent with different numbers of exponents. Takeaway: We can effectively optimize multiple exponents, which give trajectory separation in multiple directions.
EVal arg( ) log-norm of EVal log(| |)
.1 we use the diagnostic problems to demonstrate and ablate the key parts of our algorithms -i.e.optimizer choice 
(Sec. 6.1.1), starting point selection (Sec. 6.1.2), and branching (Sec. 
6.1.3). Next, in Sec. 6.2 we scale GRR to large-scale problem settings 
by (a) demonstrating that we improve RR's EVec estimation for 
neural network classifiers in Sec. 6.2.3, and (b) calculating Lyapunov 
exponents for GANs in Sec. 6.2.4. 
6.1 Diagnostic Experiments 
6.1.1 Optimizer Choice. 
Player 1 Loss
1Player 1 Strategy Distribution, [min, max]Search Strategy 
L [min, max] 
( 0 ) 
( | ) 
( | ) 
( | ) 
( | 
) 
GRR: tune max Lyap + top EVec branch + simSGD [1.000, 2.000] [.003, .999] [.032, .999] [.004, .884] [.001, .912] [.000, .013] 
GRR: tune max Lyap + top EVec branch + LOLA 
[1.000, 2.000] [.002, .999] [.063, .993] [.001, .910] [.000, .922] [.005, .103] 
20 Random init + SimSGD 
[1.997, 1.998] [.043, .194] [.142, .480] [.041, .143] [.055, .134] [.001, .001] 
20 Random init + LOLA 
[1.000, 1.396] [.000, 1.00] [.093, 1.00] [.000, .966] [.057, 1.00] [.000, .947] 
1 Random init + top EVec branch 
[2.000, 2.000] [.001, .003] [.027, .030] [.003, .007] [.008, .009] [.000, .000] 
Table 1: We show strategies for finding diverse solutions in the iterated prisoner's dilemma (IPD)
Table 2 :
2We show how many HVP evaluations we require 
to reach different MNIST classifier accuracies by following 
EVecs, repeating the exp. in RR's Fig. 4. This experiment is 
not designed to train a single strong classifier, but to test our 
ability to efficiently follow negative EVecs -see Sec. 6.2.3. 
Table 4 :
4NotationRR 
Ridge Rider [35] 
IPD 
Iterated Prisoners' Dilemma 
GAN 
Generative Adversarial Network [5] 
LOLA 
Learning with opponent learning awareness [29] 
EVec, EVal 
Shorthand for Eigenvector or Eigenvalue 
SGD 
Stochastic Gradient Descent 
SimSGD 
Figure 8: This shows the phase portrait for two standard optimization algorithms on a range of problems. Following the gradient is shown in red, while LOLA -a method for learning in games -is shown in blue. Left: The small IPD, which has solutions in the top right and bottom left. Middle: Matching Pennies, which has a single solution in the middle. Following the gradient does not find this solution because it has imaginary EVals, so we must use a method like LOLA. Right: A mixture of small IPD and Matching Pennies. Following the gradient only finds the solution in the top right, because the center solution has imaginary EVals; LOLA can find either solution. Takeaway: The mixture game has a range of phenomena, including an imaginary EVal solution, a real EVal solution, and a Hopf bifurcation. We may want to use a method for learning in games, so we can robustly converge to different solutions.Different direction estimation strategies in 10-step Lyapunov exponent calculationRandom DirectionPower Iteration, first step Power Iteration, every step jax.numpy.linalg.eigh every step Max 10-step Lyapunov ExponentThis problem has two solutions -one where both players coop-
erate, and one where both players select actions uniformly. The 
uniform action solution has imaginary EVals, so it is only stable 
under a method for learning in games, while the both cooperate 
solution has real EVals. There is a Hopf bifurcation separating these 
solutions. See Fig 2 for standard methods on this problem and Ap-
pendix Fig. 8 to contrast this problem with Small IPD or Matching 
Pennies. 
Player 1 Strategy 
Player 2 Strategy 
Small IPD 
Player 2 Strategy 
Mixed Problem 
Joint-player grad. log-norm log(∥ˆ∥) 
Player 2 Strategy 
Matching Pennies 
Player 1 Strategy 
Player 2 Strategy 
ACKNOWLEDGMENTSResources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute. Work, in part, was done while Jonathan Lorraine and Jack Parker-Holder were on internship at FAIR. We would also like to thank C. Daniel Freeman, Hérve Jégou, Noam Brown, and David Acuna for feedback on this work and acknowledge the Python community[58,59]for developing the tools that enabled this work, including numpy[60][61][62], Matplotlib[63]and SciPy[64].
Imagenet-trained cnns are biased towards texture; increasing shape bias improves accuracy and robustness. Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Wieland Felix A Wichmann, Brendel, ICLR. Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A Wichmann, and Wieland Brendel. Imagenet-trained cnns are biased towards texture; increasing shape bias improves accuracy and robustness. In ICLR, 2018.
Shortcut learning in deep neural networks. Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, Felix A Wichmann, Nature Machine Intelligence. 211Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelli- gence, 2(11):665-673, 2020.
Robots that can adapt like animals. Antoine Cully, Jeff Clune, Danesh Tarapore, Jean-Baptiste Mouret, Nature. 5217553Antoine Cully, Jeff Clune, Danesh Tarapore, and Jean-Baptiste Mouret. Robots that can adapt like animals. Nature, 521(7553):503-507, 2015.
Score-based generative modeling through stochastic differential equations. Yang Song, Jascha Sohl-Dickstein, P Diederik, Abhishek Kingma, Stefano Kumar, Ben Ermon, Poole, ICLR. Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Ku- mar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In ICLR, 2020.
Generative adversarial nets. Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, Yoshua Bengio, Advances in Neural Information Processing Systems (NeurIPS). Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Gen- erative adversarial nets. In Advances in Neural Information Processing Systems (NeurIPS), pages 2672-2680, 2014.
Connecting generative adversarial networks and actor-critic methods. David Pfau, Oriol Vinyals, arXiv:1610.01945David Pfau and Oriol Vinyals. Connecting generative adversarial networks and actor-critic methods. arXiv:1610.01945, 2016.
Ronan Collobert, and Jason Weston. Curriculum learning. Yoshua Bengio, Jérôme Louradour, International Conference on Machine Learning (ICML). Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason We- ston. Curriculum learning. In International Conference on Machine Learning (ICML), pages 41-48, 2009.
Emergent tool use from multi-agent autocurricula. Bowen Baker, Ingmar Kanitscheider, Todor Markov, Yi Wu, Glenn Powell, Bob Mcgrew, Igor Mordatch, ICLR. Bowen Baker, Ingmar Kanitscheider, Todor Markov, Yi Wu, Glenn Powell, Bob McGrew, and Igor Mordatch. Emergent tool use from multi-agent autocurricula. In ICLR, 2019.
Open-ended learning in symmetric zero-sum games. David Balduzzi, Marta Garnelo, Yoram Bachrach, Wojciech Czarnecki, Julien Perolat, Max Jaderberg, Thore Graepel, International Conference on Machine Learning (ICML). David Balduzzi, Marta Garnelo, Yoram Bachrach, Wojciech Czarnecki, Julien Perolat, Max Jaderberg, and Thore Graepel. Open-ended learn- ing in symmetric zero-sum games. In International Conference on Machine Learning (ICML), pages 434-443, 2019.
Intrinsic motivation and automatic curricula via asymmetric self-play. Sainbayar Sukhbaatar, Zeming Lin, Ilya Kostrikov, Gabriel Synnaeve, Arthur Szlam, Rob Fergus, International Conference on Learning Representations (ICLR). Sainbayar Sukhbaatar, Zeming Lin, Ilya Kostrikov, Gabriel Synnaeve, Arthur Szlam, and Rob Fergus. Intrinsic motivation and automatic curricula via asymmetric self-play. In International Conference on Learning Representations (ICLR), 2018.
Generic methods for optimization-based modeling. Justin Domke, Artificial Intelligence and Statistics. Justin Domke. Generic methods for optimization-based modeling. In Artificial Intelligence and Statistics, pages 318-326, 2012.
Gradientbased hyperparameter optimization through reversible learning. Dougal Maclaurin, David Duvenaud, Ryan Adams, International Conference on Machine Learning (ICML). Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient- based hyperparameter optimization through reversible learning. In International Conference on Machine Learning (ICML), pages 2113- 2122, 2015.
Jonathan Lorraine, David Duvenaud, arXiv:1802.09419Stochastic hyperparameter optimization through hypernetworks. arXiv preprintJonathan Lorraine and David Duvenaud. Stochastic hyperparameter optimization through hypernetworks. arXiv preprint arXiv:1802.09419, 2018.
Self-tuning networks: Bilevel optimization of hyperparameters using structured best-response functions. Matthew Mackay, Paul Vicol, Jon Lorraine, David Duvenaud, Roger Grosse, International Conference on Learning Representations (ICLR). Matthew MacKay, Paul Vicol, Jon Lorraine, David Duvenaud, and Roger Grosse. Self-tuning networks: Bilevel optimization of hyperpa- rameters using structured best-response functions. In International Conference on Learning Representations (ICLR), 2019.
Teaching with commentaries. Aniruddh Raghu, Maithra Raghu, Simon Kornblith, David Duvenaud, Geoffrey Hinton, International Conference on Learning Representations. Aniruddh Raghu, Maithra Raghu, Simon Kornblith, David Duvenaud, and Geoffrey Hinton. Teaching with commentaries. In International Conference on Learning Representations, 2020.
Optimizing millions of hyperparameters by implicit differentiation. Jonathan Lorraine, Paul Vicol, David Duvenaud, International Conference on Artificial Intelligence and Statistics (AISTATS). Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing mil- lions of hyperparameters by implicit differentiation. In International Conference on Artificial Intelligence and Statistics (AISTATS), pages 1540-1552, 2020.
Meta-learning to improve pretraining. Aniruddh Raghu, Jonathan Lorraine, Simon Kornblith, Matthew Mc-Dermott, David K Duvenaud, Advances in Neural Information Processing Systems. 342021Aniruddh Raghu, Jonathan Lorraine, Simon Kornblith, Matthew Mc- Dermott, and David K Duvenaud. Meta-learning to improve pre- training. Advances in Neural Information Processing Systems, 34, 2021.
Avishek Joey Bose, Gauthier Gidel, Hugo Berrard, Andre Cianflone, Pascal Vincent, Simon Lacoste-Julien, William L Hamilton, arXiv:2007.00720Adversarial example games. arXiv preprintAvishek Joey Bose, Gauthier Gidel, Hugo Berrard, Andre Cianflone, Pascal Vincent, Simon Lacoste-Julien, and William L Hamilton. Ad- versarial example games. arXiv preprint arXiv:2007.00720, 2020.
Adversarial examples: Attacks and defenses for deep learning. Xiaoyong Yuan, Pan He, Qile Zhu, Xiaolin Li, IEEE Transactions on Neural Networks and Learning Systems. 309Xiaoyong Yuan, Pan He, Qile Zhu, and Xiaolin Li. Adversarial exam- ples: Attacks and defenses for deep learning. IEEE Transactions on Neural Networks and Learning Systems, 30(9):2805-2824, 2019.
A game theoretic framework for model based reinforcement learning. Aravind Rajeswaran, Igor Mordatch, Vikash Kumar, International Conference on Machine Learning. Aravind Rajeswaran, Igor Mordatch, and Vikash Kumar. A game theoretic framework for model based reinforcement learning. In International Conference on Machine Learning, 2020.
A Lagrangian method for inverse problems in reinforcement learning. Pierre-Luc Bacon, Florian Schäfer, Clement Gehring, Animashree Anandkumar, Emma Brunskill, lis.csail.mit.edu/pubsPierre-Luc Bacon, Florian Schäfer, Clement Gehring, Animashree Anandkumar, and Emma Brunskill. A Lagrangian method for inverse problems in reinforcement learning. lis.csail.mit.edu/pubs, 2019.
Evgenii Nikishin, Romina Abachi, Rishabh Agarwal, Pierre-Luc Bacon, arXiv:2106.03273Control-oriented model-based reinforcement learning with implicit differentiation. arXiv preprintEvgenii Nikishin, Romina Abachi, Rishabh Agarwal, and Pierre-Luc Bacon. Control-oriented model-based reinforcement learning with implicit differentiation. arXiv preprint arXiv:2106.03273, 2021.
f-Domainadversarial learning: Theory and algorithms for unsupervised domain adaptation with neural networks. David Acuna, Guojun Zhang, T Marc, Sanja Law, Fidler, David Acuna, Guojun Zhang, Marc T Law, and Sanja Fidler. f-Domain- adversarial learning: Theory and algorithms for unsupervised domain adaptation with neural networks, 2021. URL https://openreview.net/ forum?id=WqXAKcwfZtI.
Barret Zoph, V Quoc, Le, arXiv:1611.01578Neural architecture search with reinforcement learning. arXiv preprintBarret Zoph and Quoc V Le. Neural architecture search with rein- forcement learning. arXiv preprint arXiv:1611.01578, 2016.
Regularized evolution for image classifier architecture search. Esteban Real, Alok Aggarwal, Yanping Huang, Quoc V Le, AAAI Conference on Artificial Intelligence. 33Esteban Real, Alok Aggarwal, Yanping Huang, and Quoc V Le. Regu- larized evolution for image classifier architecture search. In AAAI Conference on Artificial Intelligence, volume 33, pages 4780-4789, 2019.
DARTS: Differentiable architecture search. Hanxiao Liu, Karen Simonyan, Yiming Yang, ICLR. Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: Differen- tiable architecture search. In ICLR, 2019.
Seyed Kamyar Seyed Ghasemipour, and Richard Zemel. Gradient-based optimization of neural network architecture. Will Grathwohl, Elliot Creager, Workshop ICLR. Will Grathwohl, Elliot Creager, Seyed Kamyar Seyed Ghasemipour, and Richard Zemel. Gradient-based optimization of neural network architecture. In Workshop ICLR, 2018.
Understanding neural architecture search techniques. George Adam, Jonathan Lorraine, arXiv:1904.00438arXiv preprintGeorge Adam and Jonathan Lorraine. Understanding neural archi- tecture search techniques. arXiv preprint arXiv:1904.00438, 2019.
Learning with opponentlearning awareness. Jakob Foerster, Y Richard, Maruan Chen, Shimon Al-Shedivat, Pieter Whiteson, Igor Abbeel, Mordatch, International Conference on Autonomous Agents and MultiAgent Systems (AAMA). Jakob Foerster, Richard Y Chen, Maruan Al-Shedivat, Shimon White- son, Pieter Abbeel, and Igor Mordatch. Learning with opponent- learning awareness. In International Conference on Autonomous Agents and MultiAgent Systems (AAMA), pages 122-130, 2018.
Meta-learning for semi-supervised few-shot classification. Mengye Ren, Eleni Triantafillou, Sachin Ravi, Jake Snell, Kevin Swersky, Joshua B Tenenbaum, Hugo Larochelle, Richard S Zemel, International Conference on Learning Representations. Mengye Ren, Eleni Triantafillou, Sachin Ravi, Jake Snell, Kevin Swer- sky, Joshua B Tenenbaum, Hugo Larochelle, and Richard S Zemel. Meta-learning for semi-supervised few-shot classification. In Inter- national Conference on Learning Representations, 2018.
Meta-learning with implicit gradients. Aravind Rajeswaran, Chelsea Finn, M Sham, Sergey Kakade, Levine, Proceedings of the 33rd International Conference on Neural Information Processing Systems. the 33rd International Conference on Neural Information Processing SystemsAravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine. Meta-learning with implicit gradients. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, pages 113-124, 2019.
Flexible few-shot learning with contextual similarity. Mengye Ren, Eleni Triantafillou, Kuan-Chieh Wang, James Lucas, Jake Snell, Xaq Pitkow, Andreas S Tolias, Richard Zemel, arXiv:2012.05895arXiv preprintMengye Ren, Eleni Triantafillou, Kuan-Chieh Wang, James Lucas, Jake Snell, Xaq Pitkow, Andreas S Tolias, and Richard Zemel. Flex- ible few-shot learning with contextual similarity. arXiv preprint arXiv:2012.05895, 2020.
Wasserstein generative adversarial networks. Martin Arjovsky, Soumith Chintala, Léon Bottou, International Conference on Machine Learning (ICML). Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Ma- chine Learning (ICML), pages 214-223, 2017.
Other-Play" for zero-shot coordination. Hengyuan Hu, Adam Lerer, Alex Peysakhovich, Jakob Foerster, International Conference on Machine Learning (ICML). Hengyuan Hu, Adam Lerer, Alex Peysakhovich, and Jakob Foerster. "Other-Play" for zero-shot coordination. In International Conference on Machine Learning (ICML), pages 4399-4410, 2020.
. Jack Parker-Holder, Luke Metz, Cinjon Resnick, Hengyuan Hu, Adam Lerer, Alistair Letcher, Alexander Peysakhovich, Aldo Pacchiano, andJack Parker-Holder, Luke Metz, Cinjon Resnick, Hengyuan Hu, Adam Lerer, Alistair Letcher, Alexander Peysakhovich, Aldo Pacchiano, and
Ridge Rider: Finding diverse solutions by following eigenvectors of the Hessian. Jakob Foerster, Advances in Neural Information Processing Systems (NeurIPS). Jakob Foerster. Ridge Rider: Finding diverse solutions by follow- ing eigenvectors of the Hessian. In Advances in Neural Information Processing Systems (NeurIPS), pages 753-765, 2020.
The mechanics of n-player differentiable games. David Balduzzi, Sebastien Racaniere, James Martens, Jakob Foerster, Karl Tuyls, Thore Graepel, International Conference on Machine Learning (ICML). David Balduzzi, Sebastien Racaniere, James Martens, Jakob Foerster, Karl Tuyls, and Thore Graepel. The mechanics of n-player differen- tiable games. In International Conference on Machine Learning (ICML), pages 354-363, 2018.
Introduction to the Modern Theory of Dynamical Systems. Anatole Katok, Boris Hasselblatt, Cambridge University PressAnatole Katok and Boris Hasselblatt. Introduction to the Modern Theory of Dynamical Systems. Cambridge University Press, 1997.
Fast exact multiplication by the Hessian. A Barak, Pearlmutter, Neural Computation. 61Barak A Pearlmutter. Fast exact multiplication by the Hessian. Neural Computation, 6(1):147-160, 1994.
TensorFlow: Large-scale machine learning on heterogeneous systems. Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané. Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang ZhengMartín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015. URL https://www.tensorflow.org/.
Automatic differentiation in PyTorch. Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Devito, Zeming Lin, Alban Desmaison, Luca Antiga, Adam Lerer, Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Ed- ward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in PyTorch. 2017.
JAX: Composable transformations of Python+NumPy programs. James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake Vanderplas, Skye Wanderman-Milne, Qiao Zhang, James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: Com- posable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax.
Finite regret and cycles with fixed step-size via alternating gradient descent-ascent. P James, Gauthier Bailey, Georgios Gidel, Piliouras, Conference on Learning Theory. James P Bailey, Gauthier Gidel, and Georgios Piliouras. Finite regret and cycles with fixed step-size via alternating gradient descent-ascent. In Conference on Learning Theory, pages 391-407, 2020.
. Dimitri Bertsekas. Nonlinear Programming. Athena Scientific. Dimitri Bertsekas. Nonlinear Programming. Athena Scientific, 2008.
. K Jack, Hüseyin Hale, Koçak, Dynamics and Bifurcations. 3Springer Science & Business MediaJack K Hale and Hüseyin Koçak. Dynamics and Bifurcations, volume 3. Springer Science & Business Media, 2012.
Differentiable game mechanics. Alistair Letcher, David Balduzzi, Sébastien Racaniere, James Martens, Jakob Foerster, Karl Tuyls, Thore Graepel, The Journal of Machine Learning Research. 20Alistair Letcher, David Balduzzi, Sébastien Racaniere, James Martens, Jakob Foerster, Karl Tuyls, and Thore Graepel. Differentiable game mechanics. The Journal of Machine Learning Research, 20, 2019.
Chaos and Integrability in Nonlinear Dynamics: An Introduction. Michael Tabor, Wiley-InterscienceMichael Tabor. Chaos and Integrability in Nonlinear Dynamics: An Introduction. Wiley-Interscience, 1989.
Local Lyapunov exponents: Looking closely at chaos. C L Rodney, Wolff, Journal of the Royal Statistical Society. Rodney CL Wolff. Local Lyapunov exponents: Looking closely at chaos. Journal of the Royal Statistical Society, 1992.
Characterization of chaos in random maps. Vittorio Loreto, Giovanni Paladin, Michele Pasquini, Angelo Vulpiani, Physica A: Statistical Mechanics and its Applications. 2321-2Vittorio Loreto, Giovanni Paladin, Michele Pasquini, and Angelo Vulpiani. Characterization of chaos in random maps. Physica A: Statistical Mechanics and its Applications, 232(1-2):189-200, 1996.
Numerical implementation of the multiscale and averaging methods for quasi periodic systems. Tal Kachman, Shmuel Fishman, Avy Soffer, Computer Physics Communications. 221Tal Kachman, Shmuel Fishman, and Avy Soffer. Numerical implemen- tation of the multiscale and averaging methods for quasi periodic systems. Computer Physics Communications, 221:235-245, 2017.
Characteristic Lyapunov exponents and smooth ergodic theory. Uspekhi Matematicheskikh Nauk. Yakov Borisovich Pesin, Yakov Borisovich Pesin. Characteristic Lyapunov exponents and smooth ergodic theory. Uspekhi Matematicheskikh Nauk, 1977.
PyHessian: Neural networks through the lens of the Hessian. Zhewei Yao, Amir Gholami, Kurt Keutzer, Michael W Mahoney, IEEE International Conference on Big Data (Big Data). Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. PyHessian: Neural networks through the lens of the Hessian. In IEEE International Conference on Big Data (Big Data), pages 581-590, 2020.
Negative momentum for improved game dynamics. Gauthier Gidel, Askari Reyhane, Mohammad Hemmat, Rémi Le Pezeshki, Gabriel Priol, Simon Huang, Ioannis Lacoste-Julien, Mitliagkas, The 22nd International Conference on Artificial Intelligence and Statistics (AISTATS). Gauthier Gidel, Reyhane Askari Hemmat, Mohammad Pezeshki, Rémi Le Priol, Gabriel Huang, Simon Lacoste-Julien, and Ioannis Mitliagkas. Negative momentum for improved game dynamics. In The 22nd Inter- national Conference on Artificial Intelligence and Statistics (AISTATS), pages 1802-1811, 2019.
Jonathan Lorraine, David Acuna, Paul Vicol, David Duvenaud, arXiv:2102.08431Complex momentum for learning in games. arXiv preprintJonathan Lorraine, David Acuna, Paul Vicol, and David Duve- naud. Complex momentum for learning in games. arXiv preprint arXiv:2102.08431, 2021.
Prisoner's Dilemma/John Von Neumann, Game Theory and the Puzzle of the Bomb. William Poundstone, Anchor PressWilliam Poundstone. Prisoner's Dilemma/John Von Neumann, Game Theory and the Puzzle of the Bomb. Anchor Press, 1993.
Luke Metz, Ben Poole, David Pfau, Jascha Sohl-Dickstein, arXiv:1611.02163Unrolled generative adversarial networks. arXiv preprintLuke Metz, Ben Poole, David Pfau, and Jascha Sohl-Dickstein. Unrolled generative adversarial networks. arXiv preprint arXiv:1611.02163, 2016.
Stable opponent shaping in differentiable games. Alistair Letcher, Jakob Foerster, David Balduzzi, Tim Rocktäschel, Shimon Whiteson, International Conference on Learning Representations. Alistair Letcher, Jakob Foerster, David Balduzzi, Tim Rocktäschel, and Shimon Whiteson. Stable opponent shaping in differentiable games. In International Conference on Learning Representations, 2018.
H Steven, Strogatz, Chemistry, and Engineering. CRC PressBiologySteven H Strogatz. Nonlinear Dynamics and Chaos with Student Solutions Manual: With Applications to Physics, Biology, Chemistry, and Engineering. CRC Press, 2018.
Python reference manual. Centrum voor Wiskunde en Informatica Amsterdam. Guido Van Rossum, Fred L Drake Jr, Guido Van Rossum and Fred L Drake Jr. Python reference manual. Centrum voor Wiskunde en Informatica Amsterdam, 1995.
Python for scientific computing. E Travis, Oliphant, Computing in Science & Engineering. 93Travis E Oliphant. Python for scientific computing. Computing in Science & Engineering, 9(3):10-20, 2007.
A guide to NumPy. E Travis, Oliphant, Trelgol Publishing USA1Travis E Oliphant. A guide to NumPy, volume 1. Trelgol Publishing USA, 2006.
The NumPy array: A structure for efficient numerical computation. Stefan Van Der Walt, Chris Colbert, Gael Varoquaux, Computing in Science & Engineering. 132Stefan Van Der Walt, S Chris Colbert, and Gael Varoquaux. The NumPy array: A structure for efficient numerical computation. Com- puting in Science & Engineering, 13(2):22-30, 2011.
Array programming with NumPy. Jarrod Charles R Harris, Millman, J Stéfan, Ralf Van Der Walt, Pauli Gommers, David Virtanen, Eric Cournapeau, Julian Wieser, Sebastian Taylor, Nathaniel J Berg, Smith, Nature. 5857825Charles R Harris, K Jarrod Millman, Stéfan J van der Walt, Ralf Gom- mers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J Smith, et al. Array programming with NumPy. Nature, 585(7825):357-362, 2020.
Matplotlib: A 2d graphics environment. D John, Hunter, Computing in Science & Engineering. 93John D Hunter. Matplotlib: A 2d graphics environment. Computing in Science & Engineering, 9(3):90-95, 2007.
SciPy: Open source scientific tools for Python. Eric Jones, Travis Oliphant, Pearu Peterson, Eric Jones, Travis Oliphant, Pearu Peterson, et al. SciPy: Open source scientific tools for Python. 2001.
Neural network ensembles. Lars Kai Hansen, Peter Salamon, IEEE Transactions on Pattern Analysis and Machine Intelligence. Lars Kai Hansen and Peter Salamon. Neural network ensembles. IEEE Transactions on Pattern Analysis and Machine Intelligence, 1990.
Ensemble learning via negative correlation. Yong Liu, Xin Yao, Neural Networks. 1210Yong Liu and Xin Yao. Ensemble learning via negative correlation. Neural Networks, 12(10):1399 -1404, 1999.
Diversity inducing information bottleneck in model ensembles. Samarth Sinha, Homanga Bharadhwaj, Anirudh Goyal, Hugo Larochelle, Animesh Garg, Florian Shkurti, AAAISamarth Sinha, Homanga Bharadhwaj, Anirudh Goyal, Hugo Larochelle, Animesh Garg, and Florian Shkurti. Diversity inducing information bottleneck in model ensembles. AAAI, 2021.
Ensembles of locally independent prediction models. Andrew Slavin Ross, Weiwei Pan, Leo A Celi, Finale Doshi-Velez, The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI. Andrew Slavin Ross, Weiwei Pan, Leo A. Celi, and Finale Doshi- Velez. Ensembles of locally independent prediction models. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI, 2020.
Diversity networks: Neural network compression using determinantal point processes. Zelda Mariet, Suvrit Sra, International Conference on Learning Representations (ICLR). Zelda Mariet and Suvrit Sra. Diversity networks: Neural network compression using determinantal point processes. In International Conference on Learning Representations (ICLR), May 2016.
Improving adversarial robustness via promoting ensemble diversity. Tianyu Pang, Kun Xu, Chao Du, Ning Chen, Jun Zhu, International Conference on Machine Learning (ICML). Tianyu Pang, Kun Xu, Chao Du, Ning Chen, and Jun Zhu. Improving adversarial robustness via promoting ensemble diversity. In Interna- tional Conference on Machine Learning (ICML), 2019.
Exploiting open-endedness to solve problems through the search for novelty. Joel Lehman, Kenneth O Stanley, Proceedings of the Eleventh International Conference on Artificial Life. the Eleventh International Conference on Artificial LifeAlife XIMIT PressJoel Lehman and Kenneth O. Stanley. Exploiting open-endedness to solve problems through the search for novelty. In Proceedings of the Eleventh International Conference on Artificial Life (Alife XI). MIT Press, 2008.
Diversity is all you need: Learning skills without a reward function. Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, Sergey Levine, International Conference on Learning Representations (ICLR). Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. In International Conference on Learning Representations (ICLR), 2019.
Effective diversity in population-based reinforcement learning. Jack Parker-Holder, Aldo Pacchiano, Krzysztof Choromanski, Stephen Roberts, Advances in Neural Information Processing Systems (NeurIPS). Jack Parker-Holder, Aldo Pacchiano, Krzysztof Choromanski, and Stephen Roberts. Effective diversity in population-based reinforce- ment learning. In Advances in Neural Information Processing Systems (NeurIPS). 2020.
Quality diversity: A new frontier for evolutionary computation. Justin K Pugh, Lisa B Soros, Kenneth O Stanley, Frontiers in Robotics and AI. 340Justin K. Pugh, Lisa B. Soros, and Kenneth O. Stanley. Quality di- versity: A new frontier for evolutionary computation. Frontiers in Robotics and AI, 3:40, 2016.
The extragradient method for finding saddle points and other problems. M Galina, Korpelevich, Matecon. 12Galina M Korpelevich. The extragradient method for finding saddle points and other problems. Matecon, 12:747-756, 1976.
A tight and unified analysis of gradient-based methods for a whole spectrum of differentiable games. Waïss Azizian, Ioannis Mitliagkas, Simon Lacoste-Julien, Gauthier Gidel, International Conference on Artificial Intelligence and Statistics (AISTATS). 2020Waïss Azizian, Ioannis Mitliagkas, Simon Lacoste-Julien, and Gauthier Gidel. A tight and unified analysis of gradient-based methods for a whole spectrum of differentiable games. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2020.
Optimization, learning, and games with predictable sequences. Sasha Rakhlin, Karthik Sridharan, Advances in Neural Information Processing Systems. 26Sasha Rakhlin and Karthik Sridharan. Optimization, learning, and games with predictable sequences. Advances in Neural Information Processing Systems, 26:3066-3074, 2013.
Training gans with optimism. Constantinos Daskalakis, Andrew Ilyas, Vasilis Syrgkanis, Haoyang Zeng, International Conference on Learning Representations (ICLR 2018). Constantinos Daskalakis, Andrew Ilyas, Vasilis Syrgkanis, and Haoyang Zeng. Training gans with optimism. In International Con- ference on Learning Representations (ICLR 2018), 2018.
A variational inequality perspective on generative adversarial networks. Gauthier Gidel, Hugo Berard, Gaëtan Vignoud, Pascal Vincent, Simon Lacoste-Julien, International Conference on Learning Representations (ICLR). Gauthier Gidel, Hugo Berard, Gaëtan Vignoud, Pascal Vincent, and Simon Lacoste-Julien. A variational inequality perspective on gener- ative adversarial networks. In International Conference on Learning Representations (ICLR), 2018.
The numerics of gans. Lars Mescheder, Sebastian Nowozin, Andreas Geiger, Proceedings of the 31st International Conference on Neural Information Processing Systems. the 31st International Conference on Neural Information Processing SystemsLars Mescheder, Sebastian Nowozin, and Andreas Geiger. The nu- merics of gans. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pages 1823-1833, 2017.
On finding local Nash equilibria (and only local Nash equilibria) in zero-sum games. V Eric, Michael I Mazumdar, S Shankar Jordan, Sastry, arXiv:1901.00838arXiv preprintEric V Mazumdar, Michael I Jordan, and S Shankar Sastry. On finding local Nash equilibria (and only local Nash equilibria) in zero-sum games. arXiv preprint arXiv:1901.00838, 2019.
Competitive mirror descent. Florian Schäfer, Anima Anandkumar, Houman Owhadi, arXiv:2006.10179arXiv preprintFlorian Schäfer, Anima Anandkumar, and Houman Owhadi. Com- petitive mirror descent. arXiv preprint arXiv:2006.10179, 2020.
Implicit regularization in overparameterized bilevel optimization. Paul Vicol, Jonathan Lorraine, David Duvenaud, Roger Grosse, ICML 2021 Beyond First Order Methods Workshop. Paul Vicol, Jonathan Lorraine, David Duvenaud, and Roger Grosse. Implicit regularization in overparameterized bilevel optimization. In ICML 2021 Beyond First Order Methods Workshop, 2021.
Pick your battles: Interaction graphs as population-level objectives for strategic diversity. Marta Garnelo, Wojciech Marian Czarnecki, Siqi Liu, Dhruva Tirumala, Junhyuk Oh, Gauthier Gidel, David Hado Van Hasselt, Balduzzi, Proceedings of the 20th International Conference on Autonomous Agents and MultiAgent Systems (AAMA). the 20th International Conference on Autonomous Agents and MultiAgent Systems (AAMA)Marta Garnelo, Wojciech Marian Czarnecki, Siqi Liu, Dhruva Tiru- mala, Junhyuk Oh, Gauthier Gidel, Hado van Hasselt, and David Balduzzi. Pick your battles: Interaction graphs as population-level ob- jectives for strategic diversity. In Proceedings of the 20th International Conference on Autonomous Agents and MultiAgent Systems (AAMA), pages 1501-1503, 2021.
Grandmaster level in StarCraft II using multi-agent reinforcement learning. Oriol Vinyals, Igor Babuschkin, Wojciech M Czarnecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, H David, Richard Choi, Timo Powell, Petko Ewalds, Georgiev, Nature. 575Oriol Vinyals, Igor Babuschkin, Wojciech M Czarnecki, Michaël Math- ieu, Andrew Dudzik, Junyoung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in StarCraft II using multi-agent reinforcement learning. Nature, 575, 2019.
Multi-agent determinantal q-learning. Yaodong Yang, Ying Wen, Jun Wang, Liheng Chen, Kun Shao, David Mguni, Weinan Zhang, International Conference on Machine Learning (ICML). 2020Yaodong Yang, Ying Wen, Jun Wang, Liheng Chen, Kun Shao, David Mguni, and Weinan Zhang. Multi-agent determinantal q-learning. In International Conference on Machine Learning (ICML), 2020.
Trajectory diversity for zero-shot coordination. Andrei Lupu, Hengyuan Hu, Jakob Foerster, Proceedings of the 20th International Conference on Autonomous Agents and MultiAgent Systems. the 20th International Conference on Autonomous Agents and MultiAgent SystemsAndrei Lupu, Hengyuan Hu, and Jakob Foerster. Trajectory diversity for zero-shot coordination. In Proceedings of the 20th International Conference on Autonomous Agents and MultiAgent Systems, 2021.
A formal basis for the heuristic determination of minimum cost paths. Peter Hart, Nils Nilsson, Bertram Raphael, IEEE Transactions on Systems Science and Cybernetics. 42Peter Hart, Nils Nilsson, and Bertram Raphael. A formal basis for the heuristic determination of minimum cost paths. IEEE Transactions on Systems Science and Cybernetics, 4(2):100-107, 1968.
The Expected-Outcome Model of Two-Player Games. Bruce Abramson, Morgan KaufmannBruce Abramson. The Expected-Outcome Model of Two-Player Games. Morgan Kaufmann, 2014.
Mastering the game of Go with deep neural networks and tree search. David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den, Julian Driessche, Ioannis Schrittwieser, Veda Antonoglou, Marc Panneershelvam, Lanctot, Nature. David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of Go with deep neural networks and tree search. Nature, 2016.
A0c: Alpha zero in continuous action space. Joost Thomas M Moerland, Aske Broekens, Catholijn M Jonker Plaat, arXiv:1805.09613arXiv preprintThomas M Moerland, Joost Broekens, Aske Plaat, and Catholijn M Jonker. A0c: Alpha zero in continuous action space. arXiv preprint arXiv:1805.09613, 2018.
Monte Carlo tree search in continuous spaces using Voronoi optimistic optimization with regret bounds. Beomjoon Kim, Kyungjae Lee, Sungbin Lim, Leslie Kaelbling, Tomás Lozano-Pérez, Proceedings of the AAAI Conference on Artificial Intelligence. the AAAI Conference on Artificial Intelligence34Beomjoon Kim, Kyungjae Lee, Sungbin Lim, Leslie Kaelbling, and Tomás Lozano-Pérez. Monte Carlo tree search in continuous spaces using Voronoi optimistic optimization with regret bounds. In Pro- ceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 9916-9924, 2020.
Polyhoot: Monte-carlo planning in continuous space mdps with nonasymptotic analysis. Weichao Mao, Kaiqing Zhang, Qiaomin Xie, Tamer Basar, Advances in Neural Information Processing Systems. 33Weichao Mao, Kaiqing Zhang, Qiaomin Xie, and Tamer Basar. Poly- hoot: Monte-carlo planning in continuous space mdps with non- asymptotic analysis. Advances in Neural Information Processing Sys- tems, 33, 2020.
Rapidly-exploring random trees: A new tool for path planning. Steven M Lavalle, USAIowa State UniversityTechnical ReportSteven M LaValle et al. Rapidly-exploring random trees: A new tool for path planning. Technical Report, Iowa State University, USA, 1998.
Randomized kinodynamic planning. M Steven, James J Kuffner LavalleJr, The International Journal of Robotics Research. Steven M LaValle and James J Kuffner Jr. Randomized kinodynamic planning. The International Journal of Robotics Research, 2001.
An obstacle-based rapidly-exploring random tree. Samuel Rodriguez, Xinyu Tang, Jyh-Ming Lien, Nancy M Amato, IEEE International Conference on Robotics and Automation (ICRA). Samuel Rodriguez, Xinyu Tang, Jyh-Ming Lien, and Nancy M Amato. An obstacle-based rapidly-exploring random tree. In IEEE Interna- tional Conference on Robotics and Automation (ICRA), 2006.
Comparison of multi-modal optimization algorithms based on evolutionary algorithms. Gulshan Singh, Kalyanmoy Deb, Proceedings of the 8th Annual Conference on Genetic and Evolutionary Computation. the 8th Annual Conference on Genetic and Evolutionary ComputationGulshan Singh and Kalyanmoy Deb. Comparison of multi-modal optimization algorithms based on evolutionary algorithms. In Pro- ceedings of the 8th Annual Conference on Genetic and Evolutionary Computation, pages 1305-1312, 2006.
Evolutionary multimodal optimization: A short survey. Ka-Chun Wong, Advances in Evolutionary Algorithms Research. Ka-Chun Wong. Evolutionary multimodal optimization: A short survey. In Advances in Evolutionary Algorithms Research. 2015.
Alphastar: An evolutionary computation perspective. Kai Arulkumaran, Antoine Cully, Julian Togelius, Proceedings of the Genetic and Evolutionary Computation Conference Companion. the Genetic and Evolutionary Computation Conference CompanionKai Arulkumaran, Antoine Cully, and Julian Togelius. Alphastar: An evolutionary computation perspective. In Proceedings of the Genetic and Evolutionary Computation Conference Companion, 2019.
Using bifurcations for diversity in differentiable games. Jonathan Lorraine, Jack Parker-Holder, Paul Vicol, Aldo Pacchiano, Luke Metz, Tal Kachman, Jakob Foerster, ICML 2021 Beyond First Order Methods Workshop. Jonathan Lorraine, Jack Parker-Holder, Paul Vicol, Aldo Pacchiano, Luke Metz, Tal Kachman, and Jakob Foerster. Using bifurcations for diversity in differentiable games. In ICML 2021 Beyond First Order Methods Workshop, 2021.
Population dynamics from game theory. Christopher Zeeman, Global Theory of Dynamical Systems. SpringerE Christopher Zeeman. Population dynamics from game theory. In Global Theory of Dynamical Systems, pages 471-497. Springer, 1980.
Bifurcation mechanism design-from optimal flat taxes to better cancer treatments. Ger Yang, David Basanta, Georgios Piliouras, Games. 9221Ger Yang, David Basanta, and Georgios Piliouras. Bifurcation mech- anism design-from optimal flat taxes to better cancer treatments. Games, 9(2):21, 2018.
The route to chaos in routing games: When is price of anarchy too optimistic?. Thiparat Chotibut, Fryderyk Falniowski, Michał Misiurewicz, Georgios Piliouras, Advances in Neural Information Processing Systems (NeurIPS). 2020Thiparat Chotibut, Fryderyk Falniowski, Michał Misiurewicz, and Georgios Piliouras. The route to chaos in routing games: When is price of anarchy too optimistic? Advances in Neural Information Processing Systems (NeurIPS), 2020.
Catastrophe by design in population games: Destabilizing wasteful locked-in technologies. Georgios Piliouras, Web and Internet Economics: 16th International Conference (WINE). 473Georgios Piliouras. Catastrophe by design in population games: Destabilizing wasteful locked-in technologies. In Web and Internet Economics: 16th International Conference (WINE), page 473, 2020.
Exploration-exploitation in multi-agent learning: Catastrophe theory meets game theory. Stefanos Leonardos, Georgios Piliouras, arXiv:2012.03083arXiv preprintStefanos Leonardos and Georgios Piliouras. Exploration-exploitation in multi-agent learning: Catastrophe theory meets game theory. arXiv preprint arXiv:2012.03083, 2020.
Jakub Bielawski, Thiparat Chotibut, Fryderyk Falniowski, Grzegorz Kosiorowski, arXiv:2102.07974Michał Misiurewicz, and Georgios Piliouras. Follow-theregularized-leader routes to chaos in routing games. arXiv preprintJakub Bielawski, Thiparat Chotibut, Fryderyk Falniowski, Grzegorz Kosiorowski, Michał Misiurewicz, and Georgios Piliouras. Follow-the- regularized-leader routes to chaos in routing games. arXiv preprint arXiv:2102.07974, 2021.
Vortices instead of equilibria in minmax optimization: Chaos and butterfly effects of online learning in zero-sum games. Georgios Yun Kuen Cheung, Piliouras, Conference on Learning Theory. Yun Kuen Cheung and Georgios Piliouras. Vortices instead of equi- libria in minmax optimization: Chaos and butterfly effects of online learning in zero-sum games. In Conference on Learning Theory, pages 807-834, 2019.
Chaos of learning beyond zerosum and coordination via game decompositions. Yixin Yun Kuen Cheung, Tao, International Conference on Learning Representations (ICLR. 2020Yun Kuen Cheung and Yixin Tao. Chaos of learning beyond zero- sum and coordination via game decompositions. In International Conference on Learning Representations (ICLR), 2020.
Chaos in learning a simple two-person game. Yuzuru Sato, Eizo Akiyama, J Doyne Farmer, Proceedings of the National Academy of Sciences. 997Yuzuru Sato, Eizo Akiyama, and J Doyne Farmer. Chaos in learning a simple two-person game. Proceedings of the National Academy of Sciences, 99(7):4748-4751, 2002.
Average case performance of replicator dynamics in potential games via computing regions of attraction. Ioannis Panageas, Georgios Piliouras, Proceedings of the 2016 ACM Conference on Economics and Computation. the 2016 ACM Conference on Economics and ComputationIoannis Panageas and Georgios Piliouras. Average case performance of replicator dynamics in potential games via computing regions of attraction. In Proceedings of the 2016 ACM Conference on Economics and Computation, pages 703-720, 2016.
Equilibrium selection via replicator dynamics in 2x2 coordination games. Boyu Zhang, Josef Hofbauer, International Journal of Game Theory. 442Boyu Zhang and Josef Hofbauer. Equilibrium selection via replicator dynamics in 2x2 coordination games. International Journal of Game Theory, 44(2):433-448, 2015.
From chaos to order: Symmetry and conservation laws in game dynamics. David Sai Ganesh Nagarajan, Georgios Balduzzi, Piliouras, International Conference on Machine Learning (ICML). 2020Sai Ganesh Nagarajan, David Balduzzi, and Georgios Piliouras. From chaos to order: Symmetry and conservation laws in game dynamics. In International Conference on Machine Learning (ICML), 2020.
Lyapunov Exponents for Diversity in Differentiable Games
https://arxiv.org/abs/2112.14570
https://blog.kevmo314.com/finite-state-machines.html
On the Unreasonable Effectiveness of Finite State Machines
2020-09-23
None
<div>
 <a href="/">Back to the blog</a>
 
 <em>September 23rd, 2020</em>
 <p>
 Finite state machines are a powerful, but underutilized model for
 computation. They are the most flexible and "perfect" abstraction of how
 practical code behaves that I've run across, even moreso than my love of
 everything-is-a-function. There are more flexible models, such as Turing
 machines, but finite state machines strike a balance between being overly
 abstract and being grounded in application.
 </p>
 <p>
 A finite state machine is a <strong>machine</strong> that transitions
 between <strong>states</strong>, of which there are
 <strong>finitely many</strong>.
 </p>
 <h2>A simple example</h2>
 <p>
 We'll take this example straight from Wikipedia. A door has two states:
 <strong>open</strong> and <strong>closed</strong>. Transitioning between
 them is simply <strong>opening</strong> a closed door or
 <strong>closing</strong> an open door.
 </p>
 
 <p>
 With these states and transitions, the door can be modeled as a
 <strong>finite state machine</strong>. In this example, we don't dictate
 what opens or closes the door, but we can also specify that too. For
 example, in an automatically closing door the
 <strong>closing</strong> action is simply labeled as an automatic
 transition.
 </p>
 
 <h2>Implementing a finite state machine</h2>
 <p>
 Now that we have a door, let's implement it. Some tutorials choose to use
 a library like
 <a href="https://github.com/davidkpiano/xstate">XState</a>, but I think
 finite state machines can be implemented more elegantly without any
 libraries. Moreover, bare-bones implementations get to the real power of
 such models.
 </p>
 <p>
 First, declare our states. Each state is represented by a
 <code>type</code>.
 </p>
 <pre>type OpenState = "open";
type ClosedState = "closed";
type DoorState = OpenState | ClosedState;</pre>
 <p>
 Second, declare our transitions. Each transition is represented by a
 <code>function</code>.
 </p>
 <pre>function close(state: OpenState): ClosedState {
 return "closed";
}
function open(state: ClosedState): OpenState {
 return "open";
}</pre>
 <p>
 And that's it! To use this implementation, we can maintain our door as a
 <code>variable</code>.
 </p>
 <pre>let door: DoorState = "open";
door = close(door); // Close the door.
door = open(door); // Open it again.
door = open(door); // This won't compile, can't open an open door.</pre>
 <p>
 Try this on the
 <a href="https://www.typescriptlang.org/play?#code/C4TwDgpgBA8pB2BlYBDY0C8UBEB7B2A3AFCiRQDCANrgM4QAmyamOAxjfQ0caeNABFcuAE7N0ULHAhJUEgD6VOjcRBLEAZgFd4bYAEtc8KBzoQAFLTkQAXLASqAlHepmm1qAG9iUKCIjAWiLG2KZcPAC+vNq6BkZQ+DKW1i7K7izO9jKqXj5+AUEhifCRvFQBUAzCInZCojlYeATqVaKSJsrmrSKOhFAA9P1KZlDAABbQ3QB0xN3txV3Vvb6DWcb6wFAoAOYo+vAzc1gL3csDQwAqY-q0UADuRgDkm2y4ALZg+uUANCYo8M8EggtsZipVqlMgA">TypeScript playground</a>.
 </p>
 <h2>Adding complexity</h2>
 <p>
 The door example is simple, however. Practical finite state machines have
 many more states, transitions, and behaviors. But these can still be
 modeled with the above pattern. Let's consider an HTML button, which can
 be modeled with the following diagram.
 </p>
 
 <p>
 Play around with the behavior
 <a href="https://web.stanford.edu/class/archive/cs/cs103/cs103.1142/button-fsm/">here</a>, taking note of how the transitions work.
 </p>
 <p>Using the above framework, first declare our states.</p>
 <pre>type IdleUpState = "idle, up";
type IdleDownState = "idle, down";
type HoverState = "hover";
type HoverPressedState = "hover, pressed";
type HeldOutsideState = "held outside";
type PressedState = "pressed";
type FireState = "fire!";
type ButtonState = 
 | IdleUpState
 | IdleDownState
 | HoverState
 | HoverPressedState
 | HeldOutsideState
 | PressedState
 | FireState;</pre>
 <p>
 Then declare our transitions. I've omitted <code>leave</code>,
 <code>press</code>, and <code>release</code> to keep this readable.
 </p>
 <pre>function enter(state: IdleUpState | IdleDownState | HeldOutsideState): ButtonState {
 switch (state) {
 case "idle, up":
 return "hover";
 case "idle, down":
 return "hover, pressed";
 case "held, outside":
 return "pressed";
 }
}
...</pre>
 <p>
 Unfortunately, when we try using this implementation, something goes
 wrong!
 </p>
 <pre>let button: ButtonState = "idle, up";
button = enter(button); // This is ok.
button = leave(button); // This won't compile, but should.</pre>
 <p>
 What happened here? Sadly, the compiler doesn't know that
 <code>button</code> is a valid state when it gets passed to
 <code>leave</code>. There are three ways to solve this problem:
 </p>
 <ol>
 <li>
 Separate each transition into a different function for every input
 state.
 </li>
 <li>
 Loosen the input state requirements and don't transition if an invalid
 state is passed.
 </li>
 <li>Improve the compiler.</li>
 </ol>
 <p>
 #1 results in a lot of boilerplate but will work and #3 is outside the
 scope of this post, so I'll focus on #2. In order to not transition if the
 state is invalid, we adjust the transition so nothing happens if an
 incorrect state is passed in.
 </p>
 <pre>function enter(state: ButtonState): ButtonState {
 switch (state) {
 case "idle, up":
 return "hover";
 case "idle, down":
 return "hover, pressed";
 case "held, outside":
 return "pressed";
 default:
 return state;
 }
}</pre>
 <p>And with that simple change, our code now works.</p>
 <pre>let button = "idle, up";
button = enter(button); // This is ok.
button = leave(button); // Also now works.</pre>
 <p>Lastly, we want to bind our transitions to events.</p>
 <pre>let state = "idle, up";
button.on("mouseenter", () =&gt; state = enter(state));
button.on("mouseleave", () =&gt; state = enter(state));
document.on("mousedown", () =&gt; state = press(state));
document.on("mouseup", () =&gt; state = release(state));</pre>
 <p>
 And we're done! <code>state</code> now tracks the state of the button
 following the behavior of the state machine. Try it below, view the source
 with your browser.
 </p>
 
 <div>
 <p></p>
 </div>
 
 <h2>Why is this significant?</h2>
 <p>
 Another approach to solving #2 would be to use a higher order function
 called a <code>reducer</code>. This would remove the need for default
 statements, and might look something like this.
 </p>
 <pre>type EnterAction = "enter";
type LeaveAction = "leave";
type PressAction = "press";
type ReleaseAction = "release";
type Action =
 | EnterAction
 | LeaveAction
 | PressAction
 | ReleaseAction;
function reducer(action: Action, state: ButtonState): ButtonState {
 switch(action) {
 case "enter": {
 if (state === "idle, up" || state === "idle, down" || state === "held, outside") {
 return enter(state);
 }
 return state;
 }
 ...
 }
}</pre>
 <p>
 The other cases have been omitted for brevity, but if you're familiar with
 the direction that frontend development has been going for the last few
 years, this might seem
 <a href="https://redux.js.org/">incredibly familiar</a>. This is no
 surprise: finite state machines are an incredibly good model for
 applications. UI development in general has captured a lot of this
 intuition by decoupling state representation from rendering.
 </p>
 <p>
 If we take a step back and look again at the button finite state machine
 as an example, suppose we only look at the subset of states in the red
 box.
 </p>
 
 <p>
 This is still a finite state machine! In fact, finite state machines are
 composed <em>of finite state machines</em>. And when we look at the
 implementation of <code>reducer</code>, it looks suspiciously similar to
 the code for a transition, because it is, in fact, a transition.
 </p>
 <p>
 So finite state machines are just composed of finite state machines. Take
 this to an extreme and look at a computer. A
 <a href="https://en.wikipedia.org/wiki/Von_Neumann_architecture">von Neumann</a>
 architected computer is composed of memory, which stores state, and a
 processor, which performs transitions. Therefore,
 <strong>a physical computer is nothing but a finite state machine.</strong>
 Functions are nothing but transitions between memory states. And I think
 this realization is why writing a finite state machine from scratch is so
 insightful, because viewing code as nothing but state transitions offers a
 way to dramatically reduce complexity and make code easier to understand.
 </p>
 <p>
 Finite state machines aren't just good models for computing, they are an
 <em>exceptionally good</em> model for practical software engineering. The
 next time you run into a complex problem and it's hard to rack your head
 around the complexity, try decomposing it into simple states and
 transitions between them. As we've seen, finite state machines are
 composed of finite state machines, so not only can you start by modeling a
 small subset of states and build up from there, but also every problem is
 guaranteed to be modelable.
 </p>
 <p>
 After all, if it runs on a computer, it runs on a finite state machine, so
 it has to be, fundamentally, a transition between states. That is, in my
 opinion, unreasonably effective.
 </p>
 <hr />
 <p>
 Want more finite state machines? Read on with
 <a href="/time-invariant-finite-state-machines.html">Time-Invariant Finite State Machines!</a>
 </p>
 <p>
 Comment on this post on
 <a href="https://twitter.com/muxfdz/status/1308777691655663616">Twitter</a>
 or <a href="https://news.ycombinator.com/item?id=24568078">Hacker News</a>
 </p>
 
||||I||||Back to the blog

On the Unreasonable Effectiveness of Finite State Machines

September 23rd, 2020

Finite state machines are a powerful, but underutilized model for computation. They are the most flexible and "perfect" abstraction of how practical code behaves that I've run across, even moreso than my love of everything-is-a-function. There are more flexible models, such as Turing machines, but finite state machines strike a balance between being overly abstract and being grounded in application.

A finite state machine is a machine that transitions between states, of which there are finitely many.

A simple example

We'll take this example straight from Wikipedia. A door has two states: open and closed. Transitioning between them is simply opening a closed door or closing an open door.

With these states and transitions, the door can be modeled as a finite state machine. In this example, we don't dictate what opens or closes the door, but we can also specify that too. For example, in an automatically closing door the closing action is simply labeled as an automatic transition.

Implementing a finite state machine

Now that we have a door, let's implement it. Some tutorials choose to use a library like XState, but I think finite state machines can be implemented more elegantly without any libraries. Moreover, bare-bones implementations get to the real power of such models.

First, declare our states. Each state is represented by a type.


type OpenState = "open";
type ClosedState = "closed";

type DoorState = OpenState | ClosedState;

Second, declare our transitions. Each transition is represented by a function.


function close(state: OpenState): ClosedState {
 return "closed";
}

function open(state: ClosedState): OpenState {
 return "open";
}

And that's it! To use this implementation, we can maintain our door as a variable.


let door: DoorState = "open";

door = close(door); // Close the door.
door = open(door); // Open it again.
door = open(door); // This won't compile, can't open an open door.

Try this on the TypeScript playground.

Adding complexity

The door example is simple, however. Practical finite state machines have many more states, transitions, and behaviors. But these can still be modeled with the above pattern. Let's consider an HTML button, which can be modeled with the following diagram.

Play around with the behavior here, taking note of how the transitions work.

Using the above framework, first declare our states.


type IdleUpState = "idle, up";
type IdleDownState = "idle, down";
type HoverState = "hover";
type HoverPressedState = "hover, pressed";
type HeldOutsideState = "held outside";
type PressedState = "pressed";
type FireState = "fire!";

type ButtonState = 
 | IdleUpState
 | IdleDownState
 | HoverState
 | HoverPressedState
 | HeldOutsideState
 | PressedState
 | FireState;

Then declare our transitions. I've omitted leave, press, and release to keep this readable.


function enter(state: IdleUpState | IdleDownState | HeldOutsideState): ButtonState {
 switch (state) {
 case "idle, up":
 return "hover";
 case "idle, down":
 return "hover, pressed";
 case "held, outside":
 return "pressed";
 }
}

...

Unfortunately, when we try using this implementation, something goes wrong!


let button: ButtonState = "idle, up";

button = enter(button); // This is ok.
button = leave(button); // This won't compile, but should.

What happened here? Sadly, the compiler doesn't know that button is a valid state when it gets passed to leave. There are three ways to solve this problem:

 1. Separate each transition into a different function for every input state.
 2. Loosen the input state requirements and don't transition if an invalid state is passed.
 3. Improve the compiler.

#1 results in a lot of boilerplate but will work and #3 is outside the scope of this post, so I'll focus on #2. In order to not transition if the state is invalid, we adjust the transition so nothing happens if an incorrect state is passed in.


function enter(state: ButtonState): ButtonState {
 switch (state) {
 case "idle, up":
 return "hover";
 case "idle, down":
 return "hover, pressed";
 case "held, outside":
 return "pressed";
 default:
 return state;
 }
}

And with that simple change, our code now works.


let button = "idle, up";

button = enter(button); // This is ok.
button = leave(button); // Also now works.

Lastly, we want to bind our transitions to events.


let state = "idle, up";

button.on("mouseenter", () => state = enter(state));
button.on("mouseleave", () => state = enter(state));
document.on("mousedown", () => state = press(state));
document.on("mouseup", () => state = release(state));

And we're done! state now tracks the state of the button following the behavior of the state machine. Try it below, view the source with your browser.

Or not, needs javascript :(
 Click me!

Why is this significant?

Another approach to solving #2 would be to use a higher order function called a reducer. This would remove the need for default statements, and might look something like this.


type EnterAction = "enter";
type LeaveAction = "leave";
type PressAction = "press";
type ReleaseAction = "release";

type Action =
 | EnterAction
 | LeaveAction
 | PressAction
 | ReleaseAction;

function reducer(action: Action, state: ButtonState): ButtonState {
 switch(action) {
 case "enter": {
 if (state === "idle, up" || state === "idle, down" || state === "held, outside") {
 return enter(state);
 }
 return state;
 }
 ...
 }
}

The other cases have been omitted for brevity, but if you're familiar with the direction that frontend development has been going for the last few years, this might seem incredibly familiar. This is no surprise: finite state machines are an incredibly good model for applications. UI development in general has captured a lot of this intuition by decoupling state representation from rendering.

If we take a step back and look again at the button finite state machine as an example, suppose we only look at the subset of states in the red box.

This is still a finite state machine! In fact, finite state machines are composed of finite state machines. And when we look at the implementation of reducer, it looks suspiciously similar to the code for a transition, because it is, in fact, a transition.

So finite state machines are just composed of finite state machines. Take this to an extreme and look at a computer. A von Neumann architected computer is composed of memory, which stores state, and a processor, which performs transitions. Therefore, a physical computer is nothing but a finite state machine. Functions are nothing but transitions between memory states. And I think this realization is why writing a finite state machine from scratch is so insightful, because viewing code as nothing but state transitions offers a way to dramatically reduce complexity and make code easier to understand.

Finite state machines aren't just good models for computing, they are an exceptionally good model for practical software engineering. The next time you run into a complex problem and it's hard to rack your head around the complexity, try decomposing it into simple states and transitions between them. As we've seen, finite state machines are composed of finite state machines, so not only can you start by modeling a small subset of states and build up from there, but also every problem is guaranteed to be modelable.

After all, if it runs on a computer, it runs on a finite state machine, so it has to be, fundamentally, a transition between states. That is, in my opinion, unreasonably effective.

Want more finite state machines? Read on with Time-Invariant Finite State Machines!

Comment on this post on Twitter or Hacker News

Back to the blog
 
 September 23rd, 2020
 
 Finite state machines are a powerful, but underutilized model for
 computation. They are the most flexible and "perfect" abstraction of how
 practical code behaves that I've run across, even moreso than my love of
 everything-is-a-function. There are more flexible models, such as Turing
 machines, but finite state machines strike a balance between being overly
 abstract and being grounded in application.
 
 
 A finite state machine is a machine that transitions
 between states, of which there are
 finitely many.
 
 A simple example
 
 We'll take this example straight from Wikipedia. A door has two states:
 open and closed. Transitioning between
 them is simply opening a closed door or
 closing an open door.
 
 
 
 With these states and transitions, the door can be modeled as a
 finite state machine. In this example, we don't dictate
 what opens or closes the door, but we can also specify that too. For
 example, in an automatically closing door the
 closing action is simply labeled as an automatic
 transition.
 
 
 Implementing a finite state machine
 
 Now that we have a door, let's implement it. Some tutorials choose to use
 a library like
 XState, but I think
 finite state machines can be implemented more elegantly without any
 libraries. Moreover, bare-bones implementations get to the real power of
 such models.
 
 
 First, declare our states. Each state is represented by a
 type.
 
 type OpenState = "open";
type ClosedState = "closed";
type DoorState = OpenState | ClosedState;
 
 Second, declare our transitions. Each transition is represented by a
 function.
 
 function close(state: OpenState): ClosedState {
 return "closed";
}
function open(state: ClosedState): OpenState {
 return "open";
}
 
 And that's it! To use this implementation, we can maintain our door as a
 variable.
 
 let door: DoorState = "open";
door = close(door); // Close the door.
door = open(door); // Open it again.
door = open(door); // This won't compile, can't open an open door.
 
 Try this on the
 TypeScript playground.
 
 Adding complexity
 
 The door example is simple, however. Practical finite state machines have
 many more states, transitions, and behaviors. But these can still be
 modeled with the above pattern. Let's consider an HTML button, which can
 be modeled with the following diagram.
 
 
 
 Play around with the behavior
 here, taking note of how the transitions work.
 
 Using the above framework, first declare our states.
 type IdleUpState = "idle, up";
type IdleDownState = "idle, down";
type HoverState = "hover";
type HoverPressedState = "hover, pressed";
type HeldOutsideState = "held outside";
type PressedState = "pressed";
type FireState = "fire!";
type ButtonState = 
 | IdleUpState
 | IdleDownState
 | HoverState
 | HoverPressedState
 | HeldOutsideState
 | PressedState
 | FireState;
 
 Then declare our transitions. I've omitted leave,
 press, and release to keep this readable.
 
 function enter(state: IdleUpState | IdleDownState | HeldOutsideState): ButtonState {
 switch (state) {
 case "idle, up":
 return "hover";
 case "idle, down":
 return "hover, pressed";
 case "held, outside":
 return "pressed";
 }
}
...
 
 Unfortunately, when we try using this implementation, something goes
 wrong!
 
 let button: ButtonState = "idle, up";
button = enter(button); // This is ok.
button = leave(button); // This won't compile, but should.
 
 What happened here? Sadly, the compiler doesn't know that
 button is a valid state when it gets passed to
 leave. There are three ways to solve this problem:
 
 
 
 Separate each transition into a different function for every input
 state.
 
 
 Loosen the input state requirements and don't transition if an invalid
 state is passed.
 
 Improve the compiler.
 
 
 #1 results in a lot of boilerplate but will work and #3 is outside the
 scope of this post, so I'll focus on #2. In order to not transition if the
 state is invalid, we adjust the transition so nothing happens if an
 incorrect state is passed in.
 
 function enter(state: ButtonState): ButtonState {
 switch (state) {
 case "idle, up":
 return "hover";
 case "idle, down":
 return "hover, pressed";
 case "held, outside":
 return "pressed";
 default:
 return state;
 }
}
 And with that simple change, our code now works.
 let button = "idle, up";
button = enter(button); // This is ok.
button = leave(button); // Also now works.
 Lastly, we want to bind our transitions to events.
 let state = "idle, up";
button.on("mouseenter", () =&gt; state = enter(state));
button.on("mouseleave", () =&gt; state = enter(state));
document.on("mousedown", () =&gt; state = press(state));
document.on("mouseup", () =&gt; state = release(state));
 
 And we're done! state now tracks the state of the button
 following the behavior of the state machine. Try it below, view the source
 with your browser.
 
 
 
  
 
 
 Why is this significant?
 
 Another approach to solving #2 would be to use a higher order function
 called a reducer. This would remove the need for default
 statements, and might look something like this.
 
 type EnterAction = "enter";
type LeaveAction = "leave";
type PressAction = "press";
type ReleaseAction = "release";
type Action =
 | EnterAction
 | LeaveAction
 | PressAction
 | ReleaseAction;
function reducer(action: Action, state: ButtonState): ButtonState {
 switch(action) {
 case "enter": {
 if (state === "idle, up" || state === "idle, down" || state === "held, outside") {
 return enter(state);
 }
 return state;
 }
 ...
 }
}
 
 The other cases have been omitted for brevity, but if you're familiar with
 the direction that frontend development has been going for the last few
 years, this might seem
 incredibly familiar. This is no
 surprise: finite state machines are an incredibly good model for
 applications. UI development in general has captured a lot of this
 intuition by decoupling state representation from rendering.
 
 
 If we take a step back and look again at the button finite state machine
 as an example, suppose we only look at the subset of states in the red
 box.
 
 
 
 This is still a finite state machine! In fact, finite state machines are
 composed of finite state machines. And when we look at the
 implementation of reducer, it looks suspiciously similar to
 the code for a transition, because it is, in fact, a transition.
 
 
 So finite state machines are just composed of finite state machines. Take
 this to an extreme and look at a computer. A
 von Neumann
 architected computer is composed of memory, which stores state, and a
 processor, which performs transitions. Therefore,
 a physical computer is nothing but a finite state machine.
 Functions are nothing but transitions between memory states. And I think
 this realization is why writing a finite state machine from scratch is so
 insightful, because viewing code as nothing but state transitions offers a
 way to dramatically reduce complexity and make code easier to understand.
 
 
 Finite state machines aren't just good models for computing, they are an
 exceptionally good model for practical software engineering. The
 next time you run into a complex problem and it's hard to rack your head
 around the complexity, try decomposing it into simple states and
 transitions between them. As we've seen, finite state machines are
 composed of finite state machines, so not only can you start by modeling a
 small subset of states and build up from there, but also every problem is
 guaranteed to be modelable.
 
 
 After all, if it runs on a computer, it runs on a finite state machine, so
 it has to be, fundamentally, a transition between states. That is, in my
 opinion, unreasonably effective.
 
 
 
 Want more finite state machines? Read on with
 Time-Invariant Finite State Machines!
 
 
 Comment on this post on
 Twitter
 or Hacker News
On the Unreasonable Effectiveness of Finite State Machines
https://blog.kevmo314.com/finite-state-machines.html
https://blog.allegro.tech/2021/03/state-machines-made-easy.html
Finite-state machines made easy
2021-03-05
Tymon Felski
<div><article>
 <div> <p>Coordinating complex processes, both business and technical, can be a challenging issue in a distributed system.
Especially when the complications associated with them, such as concurrency, idempotency, scalability and hindered
testability, come into play — possibly all at once.
This is definitely something that can keep many programmers awake at night.</p>
<p>While this may sound dramatic, in reality there are many different solutions to this problem.
One group of such solutions extensively uses finite-state machines also known as finite-state automata.
As this article is meant for everyone, let’s start with some background information about what they are and how
they work.</p>
<h2>What are finite-state machines</h2>
<p>For the sake of readability, I will sometimes be referring to finite-state machines as simply “state machines”.
If you’re already familiar with their formal definition, you may skip to the next section.</p>
<p>Formally speaking, a finite-state machine is a mathematical model of computation that describes an abstract system
having a finite number of permissible states. At any given point in time, the system is in exactly one of these
states.</p>
<p>In practical terms such a machine is described by:</p>
<ul>
 <li>an initial (start) state,</li>
 <li>a set of possible states,</li>
 <li>a sequence of possible inputs (events),</li>
 <li>a transition function which based on the last observed input and the machine’s current state, determines the next
state,</li>
 <li>and a set of terminal (end) states — optionally.</li>
</ul>
<p>The machine starts off in the initial state and inspects all inputs in sequence.
Upon observing each input, it uses the transition function to change its state.
The processing ends once all inputs have been observed.
If the machine’s state at the end of processing is one of the terminal states, the machine is assumed to accept the
given input, and to reject it otherwise.</p>
<p>As a side note, the description provided above uses the deterministic model of a finite-state automaton.
There are also alternative, non-deterministic machines, but I won’t be delving into them too much, as they can be
converted to deterministic ones anyway.</p>
<p>If you’ve read this far, hopefully the rest of this post will be easier on the mind!
Despite the intimidating formal definition, finite-state machines are more widespread than you might imagine.
Upon closer inspection you might be able to spot them in your everyday life under the hood of things such as vending
machines, traffic lights and elevators.
In .NET they are used in the implementation of <code>async</code>/<code>await</code> or <code>yield</code> syntax.
Try decompiling your code and see how the language authors implemented those features.
You have probably already used some sort of state machine in your code without even realising it!</p>
<h2>A simple example of a basic state machine implementation</h2>
<p>Let’s take a closer look at the following piece of code:</p>
<div><pre><code>public enum Button
{
 Play,
 Stop
}
public enum State
{
 Playing,
 NotPlaying
}
public class Player
{
 public State CurrentState { get; private set; } = State.NotPlaying;
 public void HandleClick(Button button)
 {
 if (button == Button.Play)
 {
 CurrentState = State.Playing;
 }
 else if (button == Button.Stop)
 {
 CurrentState = State.NotPlaying;
 }
 }
}
</code></pre></div>
<p>This is an example of a very simple state machine with two states and two events.
The player can be either playing or not playing music.
Pressing the Play button will start or keep playing music, while pressing the Stop button will stop the music or keep
it stopped.</p>
<p>The code is manageable at this point, since the example is quite straightforward.
However, I hope you can imagine how adding more media controls that are usually present in music players will quickly
result in complex code that will not be easily maintainable anymore.
For instance, consider a feature request to have the Previous button change to the previous track if the player is less
than 3 seconds into the current track, or stay at the current track but rewind it to the beginning otherwise.</p>
<p>That being said, in many cases this kind of state machine implementation would be more than enough.
You don’t always have to use a state-of-the-art, bleeding-edge solution to achieve desired goals.
However, this is not something I would recommend to someone who is trying to tackle coordination of asynchronous
operations within a distributed system.</p>
<h2>A more maintainable approach</h2>
<p>Let me share our approach to this problem.
We have come up with a framework for building state machines in .NET with very little code that is maintainable and
testable.
State machines built with this framework serve a purpose of orchestrating a wide range of business and technical
processes in our microservice-based distributed architecture.
The framework is equipped with various features dictated by paradigms of distributed systems, which extend the
standard state machine definition.</p>
<p>Before I move on to an example of its real-life application, I want to break down the state machine’s API, so that you
know what you’re looking at later on.
To keep things short, I will be showing you only snippets of simplified code, as implementation details can be
expanded upon in follow-up posts.</p>
<p>The state machine is defined as follows:</p>
<div><pre><code>public class StateMachine&lt;TStateBase, TEventBase&gt;
{
 public StateMachine(TStateBase initialState)
 {
 // ...
 }
}
</code></pre></div>
<p>It is a generic class that requires a definition of base types for all states (<code>TStateBase</code>) and for all transition
triggers (<code>TEventBase</code>).
The constructor accepts an initial state from which the state machine will start its operation.</p>
<p>We can create the state machine like so:</p>
<div><pre><code>var stateMachine = new StateMachine&lt;StateBase, EventBase&gt;(new InitialState());
</code></pre></div>
<p>Having created this object, we may start defining transitions.
Each transtion consists of three required elements:</p>
<ul>
 <li>the state from which the transition can be triggered,</li>
 <li>the transition trigger (an event),</li>
 <li>the state to which the transition leads.</li>
</ul>
<p>Optionally, a transition can have side effects, called simply commands.
A command (or a set of commands) will be run after a successful transition.
Those can be pretty much anything from sending a message to another state machine to calling a specific service and
delegating a task to it.
Commands allow a seamless integration of the state machine with external components.
There is a drawback, however.
Command execution may fail, what requires an idempotent retry policy to be employed.</p>
<p>Transitions and commands can be conditional, meaning that a logical expression can be guarding and preventing the
execution of the whole transition, or just a command in some cases.</p>
<p>A builder pattern which combines all of the aforementioned requirements is used to compose a transition:</p>
<div><pre><code>stateMachine
 .FromState&lt;TStateFrom&gt;()
 .When&lt;TEvent&gt;(/* Expression&lt;Func&lt;TStateFrom, TEvent, bool&gt;&gt; */)
 .ToState&lt;TStateTo&gt;(/* Expression&lt;Func&lt;TStateFrom, TEvent, TStateTo&gt;&gt; */)
 .RunCommand(
 when: /* Expression&lt;Func&lt;TStateFrom, TEvent, TStateTo, bool&gt;&gt; */,
 then: /* Expression&lt;Func&lt;TStateFrom, TEvent, TStateTo, TCommand&gt;&gt; */)
 .RunCommand(/* ... */)
 // ...
 .RunCommand(/* ... */);
</code></pre></div>
<p>This structure may be overwhelming at first, but bear with me.
It will all make sense in a bit.
As for now you need to know that:</p>
<ul>
 <li>if the state machine is in a state of <code>TStateFrom</code> type,</li>
 <li>and an event of type <code>TEvent</code> occured,</li>
 <li>and the condition (if defined) passed to <code>When</code> method is satisfied,</li>
</ul>
<p>the transition will fire up and the state machine will switch to a new state of type <code>TStateTo</code> using the factory method
expression passed to <code>ToState</code> method.</p>
<p>The state machine class exposes a method which can be used to trigger a transition with an event:</p>
<div><pre><code>public TransitionResult Apply(TEventBase @event);
</code></pre></div>
<p>This method applies the supplied event to the current state by trying to find a matching transition in the transition
mapping and changes the internal state of the machine.
The result contains information about the performed transition, along with commands that should be run.</p>
<div><pre><code>public class TransitionResult
{
 public TStateBase StateFrom { get; }
 public TEventBase Event { get; }
 public TStateBase StateTo { get; }
 public IEnumerable&lt;object&gt; Commands { get; }
 public bool IsValid { get; }
}
</code></pre></div>
<p>If no matching transition definition is found given the machine’s current state – or if a transition is found, but its
precondition is not met – the event application will result in an invalid transition.
In that case, the state machine will not react to the event and simply retain its state, possibly logging that an
invalid event was received.</p>
<p>If there are multiple matching transitions, the first one will fire up (in definition order).
We should, however, avoid designing states machines in such way if possible.</p>
<p>As you can see, the representation of the machine is primarily type-based, so we can use polymorphism to our advantage.
By calling <code>FromState&lt;TStateBase&gt;()</code> we can have a transition that can be fired up from any state.
Similarily, if we use a type that only some selected states derive from, we will have a transition applicable to these
states only.
It would be equivalent to duplicating that same transition for every one of these states.
This syntax can be particularly useful, when dealing with processes that have an expiration date and have to be
completed within a specific time frame.
Receiving an event about the process’ expiration may have to be handled regardless of the machine’s current state and
result in termination immediately.</p>
<h2>Implementation manifest</h2>
<p>Now that I’ve described how a state machine can be defined in a declarative way using our framework, let’s talk a bit
about how that definition actually runs.</p>
<p>In general, we use dependency injection to make the state machine definition work with other services and add-ons,
such as specific command runners.
I promised I would not dive deep into implementation details, but nonetheless I want to give you a little taste of
what sits behind the scenes.</p>
<p>If you don’t care for any of that, you may skip to the “Real-life example” section.</p>
<h3>Storage</h3>
<p>The state machine implementation is based on event sourcing.
In summary, this means that we don’t save the state of our application objects, but rather a series of events that
change the object’s state.
These events will give us the current view of the object when aggregated (“replayed”) in the order in which they
occured.
For this purpose we currently use an open-source library,
<a href="https://github.com/SQLStreamStore/SQLStreamStore">SQLStreamStore</a>.</p>
<p>SQLStreamStore offers an atomic and idempotent stream append operation.
Thanks to this, we don’t have to worry about race conditions between multiple instances of the same state machine
handling events in parallel, what is great for scalability.
We can also recognize a situation when the same event is being fed into the state machine again, what can happen
when dealing with retry policies.
That, on the other hand, gives us a pretty good way of achieving deduplication.</p>
<h3>Caching</h3>
<p>Replaying all events every single time is an unnecessary overhead that we decided to mitigate by introducing
a lock-free cache for state snapshots.
After each transition a state machine will cache its current state along with the stream version which it
corresponds to.
This gives us the ability to restore the state machine to the state it was in after handling <code>n-1</code> events, when we want
to handle the <code>n</code>-th event.</p>
<p>It requires, however, the states to be immutable, or else we may end up with an unexpected behaviour.
If states were not immutable, we could not be certain that the state object we cached after handling the <code>n</code>-th event
is still the same and has not been modified in the meantime.
Moreover, immutable objects are inherently thread-safe.</p>
<h3>Instrumentation</h3>
<p>This state machine framework is also equipped with heavy instrumentation.
Every executed transition and state change is logged, so we can easily track any errors.
This data is also useful for our data analytics team to keep track of business processes.
A lot of different metrics come into play to measure performance of each and every state machine deployed.</p>
<p>Finally, every historical transition is saved in a stream — separate from the event stream — and can be viewed with our
internal back-office tools.</p>
<p>Thanks to the finite nature of the state machines, we can always log the current state and trace what made the system
to be in that state, along with all its previous states.
This is something quite invaluable and helps immensely when diagnosing issues, especially in complex and
inter-dependent business scenarios where a lot of things happen at the same time.</p>
<h2>Real-life example</h2>
<p>So finally, the big question: how is this framework used on our platform?
To answer that we don’t need to reach far.</p>
<p>In July of 2020 we launched a new payment method called Allegro Pay, which will eventually be available to every buyer
on Allegro.
This service allows the users to buy items now and pay for them later in a single payment after 30 days, or
in multiple smaller monthly payments, depending on the value of purchased goods.</p>
<p>Repayments can be easily made online using our Allegro Pay dashboard.
Users are even able to pay for multiple purchases at once.</p>
<p></p>
<p>For more detailed information, I encourage you to read our
<a href="https://allegro.pl/pomoc/dla-kupujacych/allegro-pay/">Allegro Pay FAQ</a>.</p>
<p>I suspect you already know where I’m going with this — my goal is to show you how our state machine framework is used
to coordinate an actual business process of handling a repayment within Allegro Pay.</p>
<h3>The process specification</h3>
<p>Before we create our state machine, we need to have a rough idea of which states and events we want to handle.</p>
<p>The repayment process is started in two different ways, but will end alike in both cases.</p>
<ol>
 <li>
 <p>The user can initialize an immediate online repayment via the Allegro Pay dashboard.
 They are redirected to the payment provider’s website and complete the process there.</p>
 </li>
 <li>
 <p>Alternatively, the user can request to repay their purchase with a traditional wire transfer, what we call an
offline repayment.
 We know nothing about it until the money actually arrives at the target bank account.
 That method may take up to a couple of days.</p>
 </li>
</ol>
<p>At this point both repayment paths merge, since from now on we can discard any knowledge about how the money
was received and focus on the fact that we have it.
The money source won’t be useful during the process coordination anymore.</p>
<p>The next step is to register the repayment in our bookkeeping system.
The registration part is important for us, because it will tell us whether the purchase was paid in full or if further
repayments are required.
After receiving information about successful registration, an email is sent to the user confirming that the
repayment was registered.</p>
<p>After repayment is registered, we asynchronously await the feedback from the bookkeeping system on the transaction
being settled and the repayment is to be marked as completed.</p>
<h3>Designing the state machine</h3>
<p>I could show you the finished state machine and just paste a wall of text here, describing what it does, but where’s
the fun in that?
Instead, I’d like to show you the step-by-step process I would go through when designing this kind of state machine.</p>
<p>I’ll start off by simply creating a state machine that handles a repayment of a single purchase and then extend it
to work for all cases.
First we need to define the base types for all states and events:</p>
<div><pre><code>public abstract class StateBase
{
 public string UserId { get; init; }
 public string RepaymentId { get; init; }
 public string PaymentId { get; init; }
 public StateBase(StateBase other)
 {
 if (other == null) return;
 UserId = other.UserId;
 RepaymentId = other.RepaymentId;
 PaymentId = other.PaymentId;
 }
}
public abstract class EventBase { }
</code></pre></div>
<p>As you can see, I already equipped the base state with all the properties we will need later to orchestrate this
process:</p>
<ul>
 <li><code>UserId</code> — tells us who is repaying,</li>
 <li><code>RepaymentId</code> — uniquely identifies the repayment process,</li>
 <li><code>PaymentId</code> — identifies the payment made by the user for this repayment.</li>
</ul>
<p>As for the event base, we don’t need anything special there, so it’s empty.</p>
<p><code>StateBase</code> class features a copy constructor, which rewrites all properties from the previous state when creating
a new one.</p>
<p>None of the states will have any special properties, so there’s no point in painstakingly showing all individual
subclasses representing each state.
Here are just a few sample classes of the initial state (<code>NotStarted</code>) and one of the subsequent states:</p>
<div><pre><code>public class NotStarted : StateBase
{
 public NotStarted() : base(null) { }
}
public class Created : StateBase
{
 public Created(StateBase other) : base(other) { }
}
</code></pre></div>
<p>All of the remaining states are defined similarly to <code>Created</code> state.
Those states are:</p>
<ul>
 <li><code>Paid</code>,</li>
 <li><code>Failed</code>,</li>
 <li><code>Registered</code>,</li>
 <li><code>Completed</code>.</li>
</ul>
<p>Once we have our states, we can focus on events and introduce transitions, one by one.</p>
<p>We’ll start by handling an online repayment.
It begins with the user creating a repayment entity by selecting a payment in the Allegro Pay dashboard.
This generates an event, let’s call it <code>OnlineRepaymentCreated</code>:</p>
<div><pre><code>public class OnlineRepaymentCreated : EventBase
{
 public string UserId { get; init; }
 public string RepaymentId { get; init; }
 public string PaymentId { get; init; }
}
</code></pre></div>
<p>We want this event to trigger a transition from <code>NotStarted</code> to <code>Created</code> state.</p>
<div><pre><code>// 1. NotStarted : OnlineRepaymentCreated -&gt; Created
stateMachine
 .FromState&lt;NotStarted&gt;()
 .When&lt;OnlineRepaymentCreated&gt;()
 .ToState&lt;Created&gt;((from, @event) =&gt; new Created(from)
 {
 UserId = @event.UserId,
 RepaymentId = @event.RepaymentId,
 PaymentId = @event.PaymentId
 });
</code></pre></div>
<p>At this point the user sees the repayment form, where they select a payment method.
They are then redirected to the payment provider’s website to finish the repayment.
In most cases it will succeed, but in some it fails due to a multitude of reasons, such as insufficient
funds on the account or the user providing wrong confirmation code.</p>
<p>Therefore, we need two events:</p>
<div><pre><code>public class OnlineRepaymentPaid : EventBase { }
public class OnlineRepaymentFailed : EventBase { }
</code></pre></div>
<p>We could go about a single event called <code>RepaymentResult</code> with an enum property indicating a success or a failure, but
in my opinion the former approach is cleaner and more open to future changes.
These events will play a role in tranistions no. 2 and 3:</p>
<div><pre><code>// 2. Created : OnlineRepaymentPaid -&gt; Paid
stateMachine
 .FromState&lt;Created&gt;()
 .When&lt;OnlineRepaymentPaid&gt;()
 .ToState&lt;Paid&gt;((from, @event) =&gt; new Paid(from))
 .RunCommand((from, @event, to) =&gt; new RegisterPaymentCommand(to.PaymentId));
// 3. Created : OnlineRepaymentFailed -&gt; Failed
stateMachine
 .FromState&lt;Created&gt;()
 .When&lt;OnlineRepaymentFailed&gt;()
 .ToState&lt;Failed&gt;((from, @event) =&gt; new Failed(from));
</code></pre></div>
<p>You might notice that transition no. 2 has a command that I didn’t mention earlier specified.
It’s a simple class wrapping the payment identifier, which is interpreted by the command runner as a request to our
bookkeeping system for registering this payment.</p>
<p>This concludes the online path and we can move on to the offline repayment.
It is significantly easier, since we are simply notified that the money has been transferred to us.
Because of this, we can completely skip the <code>Created</code> state and go straight to <code>Paid</code>.
We need an event called <code>OfflineRepaymentPaid</code> which contains the same properties as <code>OnlineRepaymentCreated</code>:</p>
<div><pre><code>public class OfflineRepaymentPaid : EventBase
{
 public string UserId { get; init; }
 public string RepaymentId { get; init; }
 public string PaymentId { get; init; }
}
</code></pre></div>
<p>Let’s use it in the fourth transition that will once again call the registration command:</p>
<div><pre><code>// 4. NotStarted : OfflineRepaymentPaid -&gt; Paid
stateMachine
 .FromState&lt;NotStarted&gt;()
 .When&lt;OfflineRepaymentPaid&gt;()
 .ToState&lt;Paid&gt;((from, @event) =&gt; new Paid(from)
 {
 UserId = @event.UserId,
 RepaymentId = @event.RepaymentId,
 PaymentId = @event.PaymentId
 })
 .RunCommand((from, @event, to) =&gt; new RegisterPaymentCommand(@event.PaymentId));
</code></pre></div>
<p>This is where both repayment types merge and we can focus on our bookkeeping system.
As mentioned before, we are mostly interested in registration events published by the bookkeeping system.</p>
<div><pre><code>public class PaymentRegistered : EventBase
{
 public string PaymentId { get; init; }
}
</code></pre></div>
<p>Once we know that a payment is registered in our bookkeeping system, we are able to send an email to the user
summarizing what they paid for and follow-up information whether they still have more payments to be made.
This can be done with a transition like this:</p>
<div><pre><code>// 5. Paid : PaymentRegistered -&gt; Registered
stateMachine
 .FromState&lt;Paid&gt;()
 .When&lt;PaymentRegistered&gt;()
 .ToState&lt;Registered&gt;((from, @event) =&gt; new Registered(from))
 .RunCommand((from, @event, to) =&gt; new SendRepaymentRegisteredEmailCommand(to.UserId, to.RepaymentId));
</code></pre></div>
<p>Again, <code>SendRepaymentRegisteredEmailCommand</code> is a simple class wrapping two properties that will be interpreted by the
command runner and result in an email sent to the user.</p>
<p>The last step is to wait for the final event <code>PaymentCompleted</code>.
Theoretically it could be omitted in this state machine, but it’s useful for auditing purposes.</p>
<div><pre><code>public class PaymentCompleted : EventBase
{
 public string PaymentId { get; init; }
}
</code></pre></div>
<p>It’s time to create the final transition:</p>
<div><pre><code>// 6. Registered : PaymentCompleted -&gt; Completed
stateMachine
 .FromState&lt;Registered&gt;()
 .When&lt;PaymentCompleted&gt;()
 .ToState&lt;Completed&gt;((from, @event) =&gt; new Completed(from));
</code></pre></div>
<p>And we’re done!
Right?
Not really.</p>
<p>As is usually the case with event-sourced systems, the order of incoming events can’t always be relied on.
The bookkeeping system will sometimes — for reasons I don’t want to delve into — publish events <code>PaymentRegistered</code> and
<code>PaymentCompleted</code> at the same time.
Due to a lack of message ordering guarantees, it’s possible for us to receive <code>PaymentCompleted</code> event before
<code>PaymentRegistered</code>.
The current definition of the state machine is not prepared for that.</p>
<p>Fortunately, the fix is quite easy, as we can simply introduce two more transitions:</p>
<div><pre><code>// 7. Paid : PaymentCompleted -&gt; Completed
stateMachine
 .FromState&lt;Paid&gt;()
 .When&lt;PaymentCompleted&gt;()
 .ToState&lt;Completed&gt;((from, @event) =&gt; new Completed(from));
// 8. Completed : PaymentRegistered -&gt; Completed
stateMachine
 .FromState&lt;Completed&gt;()
 .When&lt;PaymentRegistered&gt;()
 .ToState&lt;Completed&gt;((from, @event) =&gt; new Completed(from))
 .RunCommand((from, @event, to) =&gt; new SendRepaymentRegisteredEmailCommand(to.UserId, to.RepaymentId));
</code></pre></div>
<p>In this scenario we completely skip <code>Registered</code> state, but that’s something we just have to deal with.
Events are important from analytics and diagnostics standpoints, states — not so much.</p>
<p>That wraps up the event flow for this scenario.
It can be visualized with a diagram like this:</p>
<p></p>
<p>Our state machine framework is equipped with a pretty useful visualisation extension.
It’s a pretty neat tool that makes it easy to figure out how the state machine works, without having to delve into the
code.
A simple unit test that calls this extension can render a state machine diagram like one above using PlantUML
syntax, which in this case would go as follows:</p>
<div><pre><code>@startuml
hide empty description
[*] --&gt; NotStarted
NotStarted --&gt; Created: &lt;b&gt;OnlineRepaymentCreated&lt;/b&gt;
Created --&gt; Paid: &lt;b&gt;OnlineRepaymentPaid&lt;/b&gt;\n&lt;i&gt;then:&lt;/i&gt; RegisterPaymentCommand
Created --&gt; Failed: &lt;b&gt;OnlineRepaymentFailed&lt;/b&gt;
NotStarted --&gt; Paid: &lt;b&gt;OfflineRepaymentPaid&lt;/b&gt;\n&lt;i&gt;then:&lt;/i&gt; RegisterPaymentCommand
Paid --&gt; Registered: &lt;b&gt;PaymentRegistered&lt;/b&gt;\n&lt;i&gt;then:&lt;/i&gt; SendRepaymentRegisteredEmailCommand
Registered --&gt; Completed: &lt;b&gt;PaymentCompleted&lt;/b&gt;
Paid --&gt; Completed: &lt;b&gt;PaymentCompleted&lt;/b&gt;
Completed --&gt; Completed: &lt;b&gt;PaymentRegistered&lt;/b&gt;\n&lt;i&gt;then:&lt;/i&gt; SendRepaymentRegisteredEmailCommand
@enduml
</code></pre></div>
<p>However, recall that this is just a state machine to handle a repayment for a single payment selected by the user.
I promised we would extend it to handle multiple payments.
Let’s do that.
First we’ll extend the base state definition:</p>
<div><pre><code>public abstract class StateBase
{
 public string UserId { get; init; }
 public string RepaymentId { get; init; }
 public HashSet&lt;string&gt; PaymentIds { get; init; }
 public HashSet&lt;string&gt; RegisteredPaymentIds { get; init; }
 public HashSet&lt;string&gt; CompletedPaymentIds { get; init; }
 public StateBase(StateBase other)
 {
 if (other == null) return;
 UserId = other.UserId;
 RepaymentId = other.RepaymentId;
 PaymentIds = new HashSet&lt;string&gt;(other.PaymentIds);
 RegisteredPaymentIds = new HashSet&lt;string&gt;(other.RegisteredPaymentIds);
 CompletedPaymentIds = new HashSet&lt;string&gt;(other.CompletedPaymentIds);
 }
}
</code></pre></div>
<p>Three things have changed:</p>
<ul>
 <li><code>PaymentId</code> changed from a <code>string</code> to a <code>HashSet&lt;string&gt;</code> and is now called <code>PaymentIds</code> to represent multiple
payments selected by the user,</li>
 <li>a <code>HashSet&lt;string&gt;</code> called <code>RegisteredPaymentIds</code> was added to keep track of registered payments,</li>
 <li>a <code>HashSet&lt;string&gt;</code> called <code>CompletedPaymentIds</code> was added to keep track of completed payments.</li>
</ul>
<p>We also have to modify <code>OnlineRepaymentCreated</code> event by introducing a collection of payment identifiers, like so:</p>
<div><pre><code>public class OnlineRepaymentCreated : EventBase
{
 public string UserId { get; init; }
 public string RepaymentId { get; init; }
 public IEnumerable&lt;string&gt; PaymentIds { get; init; }
}
</code></pre></div>
<p>Because of this, we need to introduce a small change to transition no. 1:</p>
<div><pre><code>// 1'. NotStarted : OnlineRepaymentCreated -&gt; Created
stateMachine
 .FromState&lt;NotStarted&gt;()
 .When&lt;OnlineRepaymentCreated&gt;()
 .ToState&lt;Created&gt;((from, @event) =&gt; new Created(from)
 {
 UserId = @event.UserId,
 RepaymentId = @event.RepaymentId,
 PaymentIds = new HashSet&lt;string&gt;(@event.PaymentIds)
 });
</code></pre></div>
<p>Transition no. 2 is similarly adjusted to these changes by passing the set of payment identifiers to the
bookkeeping system:</p>
<div><pre><code>(from, @event, to) =&gt; new RegisterPaymentsCommand(to.PaymentIds)
</code></pre></div>
<p>Now we can move to the remaining transitions.
Transition no. 5 has to be split into two.
We want to stay in the <code>Paid</code> state until all registration events are received and only then switch to <code>Registered</code>.
This can be achieved with:</p>
<div><pre><code>// 5'. Paid : PaymentRegistered -&gt; Paid
stateMachine
 .FromState&lt;Paid&gt;()
 .When&lt;PaymentRegistered&gt;((from, @event) =&gt; !from.PaymentIds.All(
 id =&gt; from.RegisteredPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState&lt;Paid&gt;((from, @event) =&gt;
 {
 var to = new Paid(from);
 to.RegisteredPaymentIds.Add(@event.PaymentId);
 return to;
 });
// 5''. Paid : PaymentRegistered -&gt; Registered
stateMachine
 .FromState&lt;Paid&gt;()
 .When&lt;PaymentRegistered&gt;((from, @event) =&gt; from.PaymentIds.All(
 id =&gt; from.RegisteredPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState&lt;Registered&gt;((from, @event) =&gt;
 {
 var to = new Registered(from);
 to.RegisteredPaymentIds.Add(@event.PaymentId);
 return to;
 });
 .RunCommand((from, @event, to) =&gt; new SendRepaymentRegisteredEmailCommand(to.UserId, to.RepaymentId));
</code></pre></div>
<p>As you can see, the command is now defined only for the second transition that triggers after all payments are
registered.
We have to slightly adjust transition no. 8, which is responsible for sending the email, too:</p>
<div><pre><code>// 8'. Completed : PaymentRegistered -&gt; Completed
stateMachine
 .FromState&lt;Completed&gt;()
 .When&lt;PaymentRegistered&gt;()
 .ToState&lt;Completed&gt;((from, @event) =&gt;
 {
 var to = new Completed(from);
 to.RegisteredPaymentIds.Add(@event.PaymentId);
 return to;
 })
 .RunCommand(
 when: (from, @event, to) =&gt; to.PaymentIds.SetEquals(to.RegisteredPaymentIds),
 then: (from, @event, to) =&gt; new SendRepaymentRegisteredEmailCommand(to.UserId, to.RepaymentId));
</code></pre></div>
<p>This time the transition is not conditional, however the command is.
It will execute only after we handle final registration event and the payment identifier sets <code>PaymentIds</code> and
<code>RegisteredPaymentIds</code> become equal.</p>
<p>We cannot forget about all transitions that are triggered by <code>PaymentCompleted</code> event.
Those will have to be changed like so:</p>
<div><pre><code>// 6'. Registered : PaymentCompleted -&gt; Registered
stateMachine
 .FromState&lt;Registered&gt;()
 .When&lt;PaymentCompleted&gt;((from, @event) =&gt; !from.PaymentIds.All(
 id =&gt; from.CompletedPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState&lt;Registered&gt;((from, @event) =&gt;
 {
 var to = new Registered(from);
 to.CompletedPaymentIds.Add(@event.PaymentId);
 return to;
 });
// 6''. Registered : PaymentCompleted -&gt; Completed
stateMachine
 .FromState&lt;Registered&gt;()
 .When&lt;PaymentCompleted&gt;((from, @event) =&gt; from.PaymentIds.All(
 id =&gt; from.CompletedPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState&lt;Completed&gt;((from, @event) =&gt;
 {
 var to = new Completed(from);
 to.CompletedPaymentIds.Add(@event.PaymentId);
 return to;
 });
// 7'. Paid : PaymentCompleted -&gt; Paid
stateMachine
 .FromState&lt;Paid&gt;()
 .When&lt;PaymentCompleted&gt;((from, @event) =&gt; !from.PaymentIds.All(
 id =&gt; from.CompletedPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState&lt;Completed&gt;((from, @event) =&gt;
 {
 var to = new Completed(from);
 to.CompletedPaymentIds.Add(@event.PaymentId);
 return to;
 });
// 7''. Paid : PaymentCompleted -&gt; Completed
stateMachine
 .FromState&lt;Paid&gt;()
 .When&lt;PaymentCompleted&gt;((from, @event) =&gt; from.PaymentIds.All(
 id =&gt; from.CompletedPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState&lt;Completed&gt;((from, @event) =&gt;
 {
 var to = new Completed(from);
 to.CompletedPaymentIds.Add(@event.PaymentId);
 return to;
 });
</code></pre></div>
<p>And we’re done!
This time for real.
After all this work we ended up with something that can be visualized with a diagram like this:</p>
<p></p>
<p>I mentioned earlier that we have internal back-office tools to view state machines and their transition history.
Here is an example view of a repayment process that was coordinated with the state machine we defined above (sensitive
data was redacted, since this was taken in production):</p>
<p></p>
<h3>Caveats</h3>
<p>You may have noticed some problems here.
What would happen if we received an event with a wrong payment identifier?
Right now the state machine is not ready for that.
However, that can be remedied relatively easily by adding another condition to <code>When</code> methods of all transitions for
<code>PaymentRegistered</code> and <code>PaymentCompleted</code> events, such as:</p>
<div><pre><code>.When((from, @event) =&gt; from.PaymentIds.Contains(@event.PaymentId) &amp;&amp; /* ... */)
</code></pre></div>
<p>This would essentialy ensure that we react only to events with correct identifiers.</p>
<p>And what about retries and duplicated events?
Transition no. 8’ is especially susceptible to that, because it may result in the email being sent to the user twice.
And we don’t want that!
Fortunately, achieving deduplication is just as easy — we have to ignore events for payments that are already
registered:</p>
<div><pre><code>.When((from, @event) =&gt; !from.RegisteredPaymentIds.Contains(@event.PaymentId) &amp;&amp; /* ... */)
</code></pre></div>
<p>For the sake of readability I also decided to skip a lot of other validations, which would otherwise have to be there,
such as simple null checks.</p>
<h3>How to test</h3>
<p>Testing state machines is important, because, well, ideally all code should be tested.
In this case, apart from the usual assurance that we are able to handle both happy paths and edge cases, we can also
ensure that we keep the transitions backwards compatible.
An event that is once inserted into the stream will stay there forever.</p>
<p>We developed a simple DSL in F# to help us test state machines created with our framework.
This is a perfect solution to test a declarative definition of a state machine like the one we developed in this post.
Here’s an example of a happy path test for our state machine:</p>
<div><pre><code>[&lt;Fact&gt;]
let ``Online repayment - single payment happy path`` () =
 [
 When(Events.OnlineRepaymentCreated ("paymentId1"))
 |&gt; GoToState&lt;Created&gt;
 |&gt; DoNothing
 When(Events.OnlineRepaymentPaid ())
 |&gt; GoToState&lt;Paid&gt;
 |&gt; Do&lt;RegisterPaymentsCommand&gt;
 When(Events.PaymentRegistered ("paymentId1"))
 |&gt; GoToState&lt;Registered&gt;
 |&gt; Do&lt;SendRepaymentRegisteredEmailCommand&gt;
 When(Events.PaymentCompleted ("paymentId1"))
 |&gt; GoToState&lt;Completed&gt;
 |&gt; DoNothing
 ]
 |&gt; apply
</code></pre></div>
<p>A test like this is essentially a list of events with expected side effects that are then aggregated using the <code>apply</code>
function that feeds them one by one into the state machine and verifies the outcome.</p>
<h2>Summary</h2>
<p>I hope I was able to convince you how a framework like this can make designing complex state machines easy.
It provides numerous advantages, such as declarative syntax, effortless testability, thorough instrumentation and
visualization.</p>
<p>This is precisely why state machines became an integral part of pretty much every business process architecture
in our project.
Having such a robust and scalable solution for asynchronous process orchestration and isolation is crucial in
distributed systems.
The presented code, albeit extremely simplified, is from an actual state machine — one of many used in our
applications.</p>
<p>While the presented solution is a strictly internal framework right now, who knows what the future might bring?
Stay tuned!</p>
</div>
 </article>||||I|||| Back to

 * about us
 * 
 * 
 * 
 * 
 * 
 Mar 5 2021

 Finite-state machines made easy

 * allegro.tech blog
 * Tymon Felski
 * Finite-state machines made easy

 Coordinating complex processes, both business and technical, can be a challenging issue in a distributed system. Especially when the complications associated with them, such as concurrency, idempotency, scalability and hindered testability, come into play — possibly all at once. This is definitely something that can keep many programmers awake at night.

 While this may sound dramatic, in reality there are many different solutions to this problem. One group of such solutions extensively uses finite-state machines also known as finite-state automata. As this article is meant for everyone, let’s start with some background information about what they are and how they work.

 What are finite-state machines

 For the sake of readability, I will sometimes be referring to finite-state machines as simply “state machines”. If you’re already familiar with their formal definition, you may skip to the next section.

 Formally speaking, a finite-state machine is a mathematical model of computation that describes an abstract system having a finite number of permissible states. At any given point in time, the system is in exactly one of these states.

 In practical terms such a machine is described by:

 * an initial (start) state,
 * a set of possible states,
 * a sequence of possible inputs (events),
 * a transition function which based on the last observed input and the machine’s current state, determines the next state,
 * and a set of terminal (end) states — optionally.

 The machine starts off in the initial state and inspects all inputs in sequence. Upon observing each input, it uses the transition function to change its state. The processing ends once all inputs have been observed. If the machine’s state at the end of processing is one of the terminal states, the machine is assumed to accept the given input, and to reject it otherwise.

 As a side note, the description provided above uses the deterministic model of a finite-state automaton. There are also alternative, non-deterministic machines, but I won’t be delving into them too much, as they can be converted to deterministic ones anyway.

 If you’ve read this far, hopefully the rest of this post will be easier on the mind! Despite the intimidating formal definition, finite-state machines are more widespread than you might imagine. Upon closer inspection you might be able to spot them in your everyday life under the hood of things such as vending machines, traffic lights and elevators. In .NET they are used in the implementation of async/await or yield syntax. Try decompiling your code and see how the language authors implemented those features. You have probably already used some sort of state machine in your code without even realising it!

 A simple example of a basic state machine implementation

 Let’s take a closer look at the following piece of code:

 public enum Button
 {
 Play,
 Stop
 }
 
 public enum State
 {
 Playing,
 NotPlaying
 }
 
 public class Player
 {
 public State CurrentState { get; private set; } = State.NotPlaying;
 
 public void HandleClick(Button button)
 {
 if (button == Button.Play)
 {
 CurrentState = State.Playing;
 }
 else if (button == Button.Stop)
 {
 CurrentState = State.NotPlaying;
 }
 }
 }
 

 This is an example of a very simple state machine with two states and two events. The player can be either playing or not playing music. Pressing the Play button will start or keep playing music, while pressing the Stop button will stop the music or keep it stopped.

 The code is manageable at this point, since the example is quite straightforward. However, I hope you can imagine how adding more media controls that are usually present in music players will quickly result in complex code that will not be easily maintainable anymore. For instance, consider a feature request to have the Previous button change to the previous track if the player is less than 3 seconds into the current track, or stay at the current track but rewind it to the beginning otherwise.

 That being said, in many cases this kind of state machine implementation would be more than enough. You don’t always have to use a state-of-the-art, bleeding-edge solution to achieve desired goals. However, this is not something I would recommend to someone who is trying to tackle coordination of asynchronous operations within a distributed system.

 A more maintainable approach

 Let me share our approach to this problem. We have come up with a framework for building state machines in .NET with very little code that is maintainable and testable. State machines built with this framework serve a purpose of orchestrating a wide range of business and technical processes in our microservice-based distributed architecture. The framework is equipped with various features dictated by paradigms of distributed systems, which extend the standard state machine definition.

 Before I move on to an example of its real-life application, I want to break down the state machine’s API, so that you know what you’re looking at later on. To keep things short, I will be showing you only snippets of simplified code, as implementation details can be expanded upon in follow-up posts.

 The state machine is defined as follows:

 public class StateMachine<TStateBase, TEventBase>
 {
 public StateMachine(TStateBase initialState)
 {
 // ...
 }
 }
 

 It is a generic class that requires a definition of base types for all states (TStateBase) and for all transition triggers (TEventBase). The constructor accepts an initial state from which the state machine will start its operation.

 We can create the state machine like so:

 var stateMachine = new StateMachine<StateBase, EventBase>(new InitialState());
 

 Having created this object, we may start defining transitions. Each transtion consists of three required elements:

 * the state from which the transition can be triggered,
 * the transition trigger (an event),
 * the state to which the transition leads.

 Optionally, a transition can have side effects, called simply commands. A command (or a set of commands) will be run after a successful transition. Those can be pretty much anything from sending a message to another state machine to calling a specific service and delegating a task to it. Commands allow a seamless integration of the state machine with external components. There is a drawback, however. Command execution may fail, what requires an idempotent retry policy to be employed.

 Transitions and commands can be conditional, meaning that a logical expression can be guarding and preventing the execution of the whole transition, or just a command in some cases.

 A builder pattern which combines all of the aforementioned requirements is used to compose a transition:

 stateMachine
 .FromState<TStateFrom>()
 .When<TEvent>(/* Expression<Func<TStateFrom, TEvent, bool>> */)
 .ToState<TStateTo>(/* Expression<Func<TStateFrom, TEvent, TStateTo>> */)
 .RunCommand(
 when: /* Expression<Func<TStateFrom, TEvent, TStateTo, bool>> */,
 then: /* Expression<Func<TStateFrom, TEvent, TStateTo, TCommand>> */)
 .RunCommand(/* ... */)
 // ...
 .RunCommand(/* ... */);
 

 This structure may be overwhelming at first, but bear with me. It will all make sense in a bit. As for now you need to know that:

 * if the state machine is in a state of TStateFrom type,
 * and an event of type TEvent occured,
 * and the condition (if defined) passed to When method is satisfied,

 the transition will fire up and the state machine will switch to a new state of type TStateTo using the factory method expression passed to ToState method.

 The state machine class exposes a method which can be used to trigger a transition with an event:

 public TransitionResult Apply(TEventBase @event);
 

 This method applies the supplied event to the current state by trying to find a matching transition in the transition mapping and changes the internal state of the machine. The result contains information about the performed transition, along with commands that should be run.

 public class TransitionResult
 {
 public TStateBase StateFrom { get; }
 public TEventBase Event { get; }
 public TStateBase StateTo { get; }
 public IEnumerable<object> Commands { get; }
 public bool IsValid { get; }
 }
 

 If no matching transition definition is found given the machine’s current state – or if a transition is found, but its precondition is not met – the event application will result in an invalid transition. In that case, the state machine will not react to the event and simply retain its state, possibly logging that an invalid event was received.

 If there are multiple matching transitions, the first one will fire up (in definition order). We should, however, avoid designing states machines in such way if possible.

 As you can see, the representation of the machine is primarily type-based, so we can use polymorphism to our advantage. By calling FromState<TStateBase>() we can have a transition that can be fired up from any state. Similarily, if we use a type that only some selected states derive from, we will have a transition applicable to these states only. It would be equivalent to duplicating that same transition for every one of these states. This syntax can be particularly useful, when dealing with processes that have an expiration date and have to be completed within a specific time frame. Receiving an event about the process’ expiration may have to be handled regardless of the machine’s current state and result in termination immediately.

 Implementation manifest

 Now that I’ve described how a state machine can be defined in a declarative way using our framework, let’s talk a bit about how that definition actually runs.

 In general, we use dependency injection to make the state machine definition work with other services and add-ons, such as specific command runners. I promised I would not dive deep into implementation details, but nonetheless I want to give you a little taste of what sits behind the scenes.

 If you don’t care for any of that, you may skip to the “Real-life example” section.

 Storage

 The state machine implementation is based on event sourcing. In summary, this means that we don’t save the state of our application objects, but rather a series of events that change the object’s state. These events will give us the current view of the object when aggregated (“replayed”) in the order in which they occured. For this purpose we currently use an open-source library, SQLStreamStore.

 SQLStreamStore offers an atomic and idempotent stream append operation. Thanks to this, we don’t have to worry about race conditions between multiple instances of the same state machine handling events in parallel, what is great for scalability. We can also recognize a situation when the same event is being fed into the state machine again, what can happen when dealing with retry policies. That, on the other hand, gives us a pretty good way of achieving deduplication.

 Caching

 Replaying all events every single time is an unnecessary overhead that we decided to mitigate by introducing a lock-free cache for state snapshots. After each transition a state machine will cache its current state along with the stream version which it corresponds to. This gives us the ability to restore the state machine to the state it was in after handling n-1 events, when we want to handle the n-th event.

 It requires, however, the states to be immutable, or else we may end up with an unexpected behaviour. If states were not immutable, we could not be certain that the state object we cached after handling the n-th event is still the same and has not been modified in the meantime. Moreover, immutable objects are inherently thread-safe.

 Instrumentation

 This state machine framework is also equipped with heavy instrumentation. Every executed transition and state change is logged, so we can easily track any errors. This data is also useful for our data analytics team to keep track of business processes. A lot of different metrics come into play to measure performance of each and every state machine deployed.

 Finally, every historical transition is saved in a stream — separate from the event stream — and can be viewed with our internal back-office tools.

 Thanks to the finite nature of the state machines, we can always log the current state and trace what made the system to be in that state, along with all its previous states. This is something quite invaluable and helps immensely when diagnosing issues, especially in complex and inter-dependent business scenarios where a lot of things happen at the same time.

 Real-life example

 So finally, the big question: how is this framework used on our platform? To answer that we don’t need to reach far.

 In July of 2020 we launched a new payment method called Allegro Pay, which will eventually be available to every buyer on Allegro. This service allows the users to buy items now and pay for them later in a single payment after 30 days, or in multiple smaller monthly payments, depending on the value of purchased goods.

 Repayments can be easily made online using our Allegro Pay dashboard. Users are even able to pay for multiple purchases at once.

 For more detailed information, I encourage you to read our Allegro Pay FAQ.

 I suspect you already know where I’m going with this — my goal is to show you how our state machine framework is used to coordinate an actual business process of handling a repayment within Allegro Pay.

 The process specification

 Before we create our state machine, we need to have a rough idea of which states and events we want to handle.

 The repayment process is started in two different ways, but will end alike in both cases.

 1. The user can initialize an immediate online repayment via the Allegro Pay dashboard. They are redirected to the payment provider’s website and complete the process there.

 2. Alternatively, the user can request to repay their purchase with a traditional wire transfer, what we call an offline repayment. We know nothing about it until the money actually arrives at the target bank account. That method may take up to a couple of days.

 At this point both repayment paths merge, since from now on we can discard any knowledge about how the money was received and focus on the fact that we have it. The money source won’t be useful during the process coordination anymore.

 The next step is to register the repayment in our bookkeeping system. The registration part is important for us, because it will tell us whether the purchase was paid in full or if further repayments are required. After receiving information about successful registration, an email is sent to the user confirming that the repayment was registered.

 After repayment is registered, we asynchronously await the feedback from the bookkeeping system on the transaction being settled and the repayment is to be marked as completed.

 Designing the state machine

 I could show you the finished state machine and just paste a wall of text here, describing what it does, but where’s the fun in that? Instead, I’d like to show you the step-by-step process I would go through when designing this kind of state machine.

 I’ll start off by simply creating a state machine that handles a repayment of a single purchase and then extend it to work for all cases. First we need to define the base types for all states and events:

 public abstract class StateBase
 {
 public string UserId { get; init; }
 public string RepaymentId { get; init; }
 public string PaymentId { get; init; }
 
 public StateBase(StateBase other)
 {
 if (other == null) return;
 
 UserId = other.UserId;
 RepaymentId = other.RepaymentId;
 PaymentId = other.PaymentId;
 }
 }
 
 public abstract class EventBase { }
 

 As you can see, I already equipped the base state with all the properties we will need later to orchestrate this process:

 * UserId — tells us who is repaying,
 * RepaymentId — uniquely identifies the repayment process,
 * PaymentId — identifies the payment made by the user for this repayment.

 As for the event base, we don’t need anything special there, so it’s empty.

 StateBase class features a copy constructor, which rewrites all properties from the previous state when creating a new one.

 None of the states will have any special properties, so there’s no point in painstakingly showing all individual subclasses representing each state. Here are just a few sample classes of the initial state (NotStarted) and one of the subsequent states:

 public class NotStarted : StateBase
 {
 public NotStarted() : base(null) { }
 }
 
 public class Created : StateBase
 {
 public Created(StateBase other) : base(other) { }
 }
 

 All of the remaining states are defined similarly to Created state. Those states are:

 * Paid,
 * Failed,
 * Registered,
 * Completed.

 Once we have our states, we can focus on events and introduce transitions, one by one.

 We’ll start by handling an online repayment. It begins with the user creating a repayment entity by selecting a payment in the Allegro Pay dashboard. This generates an event, let’s call it OnlineRepaymentCreated:

 public class OnlineRepaymentCreated : EventBase
 {
 public string UserId { get; init; }
 public string RepaymentId { get; init; }
 public string PaymentId { get; init; }
 }
 

 We want this event to trigger a transition from NotStarted to Created state.

 // 1. NotStarted : OnlineRepaymentCreated -> Created
 stateMachine
 .FromState<NotStarted>()
 .When<OnlineRepaymentCreated>()
 .ToState<Created>((from, @event) => new Created(from)
 {
 UserId = @event.UserId,
 RepaymentId = @event.RepaymentId,
 PaymentId = @event.PaymentId
 });
 

 At this point the user sees the repayment form, where they select a payment method. They are then redirected to the payment provider’s website to finish the repayment. In most cases it will succeed, but in some it fails due to a multitude of reasons, such as insufficient funds on the account or the user providing wrong confirmation code.

 Therefore, we need two events:

 public class OnlineRepaymentPaid : EventBase { }
 
 public class OnlineRepaymentFailed : EventBase { }
 

 We could go about a single event called RepaymentResult with an enum property indicating a success or a failure, but in my opinion the former approach is cleaner and more open to future changes. These events will play a role in tranistions no. 2 and 3:

 // 2. Created : OnlineRepaymentPaid -> Paid
 stateMachine
 .FromState<Created>()
 .When<OnlineRepaymentPaid>()
 .ToState<Paid>((from, @event) => new Paid(from))
 .RunCommand((from, @event, to) => new RegisterPaymentCommand(to.PaymentId));
 
 // 3. Created : OnlineRepaymentFailed -> Failed
 stateMachine
 .FromState<Created>()
 .When<OnlineRepaymentFailed>()
 .ToState<Failed>((from, @event) => new Failed(from));
 

 You might notice that transition no. 2 has a command that I didn’t mention earlier specified. It’s a simple class wrapping the payment identifier, which is interpreted by the command runner as a request to our bookkeeping system for registering this payment.

 This concludes the online path and we can move on to the offline repayment. It is significantly easier, since we are simply notified that the money has been transferred to us. Because of this, we can completely skip the Created state and go straight to Paid. We need an event called OfflineRepaymentPaid which contains the same properties as OnlineRepaymentCreated:

 public class OfflineRepaymentPaid : EventBase
 {
 public string UserId { get; init; }
 public string RepaymentId { get; init; }
 public string PaymentId { get; init; }
 }
 

 Let’s use it in the fourth transition that will once again call the registration command:

 // 4. NotStarted : OfflineRepaymentPaid -> Paid
 stateMachine
 .FromState<NotStarted>()
 .When<OfflineRepaymentPaid>()
 .ToState<Paid>((from, @event) => new Paid(from)
 {
 UserId = @event.UserId,
 RepaymentId = @event.RepaymentId,
 PaymentId = @event.PaymentId
 })
 .RunCommand((from, @event, to) => new RegisterPaymentCommand(@event.PaymentId));
 

 This is where both repayment types merge and we can focus on our bookkeeping system. As mentioned before, we are mostly interested in registration events published by the bookkeeping system.

 public class PaymentRegistered : EventBase
 {
 public string PaymentId { get; init; }
 }
 

 Once we know that a payment is registered in our bookkeeping system, we are able to send an email to the user summarizing what they paid for and follow-up information whether they still have more payments to be made. This can be done with a transition like this:

 // 5. Paid : PaymentRegistered -> Registered
 stateMachine
 .FromState<Paid>()
 .When<PaymentRegistered>()
 .ToState<Registered>((from, @event) => new Registered(from))
 .RunCommand((from, @event, to) => new SendRepaymentRegisteredEmailCommand(to.UserId, to.RepaymentId));
 

 Again, SendRepaymentRegisteredEmailCommand is a simple class wrapping two properties that will be interpreted by the command runner and result in an email sent to the user.

 The last step is to wait for the final event PaymentCompleted. Theoretically it could be omitted in this state machine, but it’s useful for auditing purposes.

 public class PaymentCompleted : EventBase
 {
 public string PaymentId { get; init; }
 }
 

 It’s time to create the final transition:

 // 6. Registered : PaymentCompleted -> Completed
 stateMachine
 .FromState<Registered>()
 .When<PaymentCompleted>()
 .ToState<Completed>((from, @event) => new Completed(from));
 

 And we’re done! Right? Not really.

 As is usually the case with event-sourced systems, the order of incoming events can’t always be relied on. The bookkeeping system will sometimes — for reasons I don’t want to delve into — publish events PaymentRegistered and PaymentCompleted at the same time. Due to a lack of message ordering guarantees, it’s possible for us to receive PaymentCompleted event before PaymentRegistered. The current definition of the state machine is not prepared for that.

 Fortunately, the fix is quite easy, as we can simply introduce two more transitions:

 // 7. Paid : PaymentCompleted -> Completed
 stateMachine
 .FromState<Paid>()
 .When<PaymentCompleted>()
 .ToState<Completed>((from, @event) => new Completed(from));
 
 // 8. Completed : PaymentRegistered -> Completed
 stateMachine
 .FromState<Completed>()
 .When<PaymentRegistered>()
 .ToState<Completed>((from, @event) => new Completed(from))
 .RunCommand((from, @event, to) => new SendRepaymentRegisteredEmailCommand(to.UserId, to.RepaymentId));
 

 In this scenario we completely skip Registered state, but that’s something we just have to deal with. Events are important from analytics and diagnostics standpoints, states — not so much.

 That wraps up the event flow for this scenario. It can be visualized with a diagram like this:

 Our state machine framework is equipped with a pretty useful visualisation extension. It’s a pretty neat tool that makes it easy to figure out how the state machine works, without having to delve into the code. A simple unit test that calls this extension can render a state machine diagram like one above using PlantUML syntax, which in this case would go as follows:

 @startuml
 hide empty description
 [*] --> NotStarted
 NotStarted --> Created: <b>OnlineRepaymentCreated</b>
 Created --> Paid: <b>OnlineRepaymentPaid</b>\n<i>then:</i> RegisterPaymentCommand
 Created --> Failed: <b>OnlineRepaymentFailed</b>
 NotStarted --> Paid: <b>OfflineRepaymentPaid</b>\n<i>then:</i> RegisterPaymentCommand
 Paid --> Registered: <b>PaymentRegistered</b>\n<i>then:</i> SendRepaymentRegisteredEmailCommand
 Registered --> Completed: <b>PaymentCompleted</b>
 Paid --> Completed: <b>PaymentCompleted</b>
 Completed --> Completed: <b>PaymentRegistered</b>\n<i>then:</i> SendRepaymentRegisteredEmailCommand
 @enduml
 

 However, recall that this is just a state machine to handle a repayment for a single payment selected by the user. I promised we would extend it to handle multiple payments. Let’s do that. First we’ll extend the base state definition:

 public abstract class StateBase
 {
 public string UserId { get; init; }
 public string RepaymentId { get; init; }
 public HashSet<string> PaymentIds { get; init; }
 public HashSet<string> RegisteredPaymentIds { get; init; }
 public HashSet<string> CompletedPaymentIds { get; init; }
 
 public StateBase(StateBase other)
 {
 if (other == null) return;
 
 UserId = other.UserId;
 RepaymentId = other.RepaymentId;
 PaymentIds = new HashSet<string>(other.PaymentIds);
 RegisteredPaymentIds = new HashSet<string>(other.RegisteredPaymentIds);
 CompletedPaymentIds = new HashSet<string>(other.CompletedPaymentIds);
 }
 }
 

 Three things have changed:

 * PaymentId changed from a string to a HashSet<string> and is now called PaymentIds to represent multiple payments selected by the user,
 * a HashSet<string> called RegisteredPaymentIds was added to keep track of registered payments,
 * a HashSet<string> called CompletedPaymentIds was added to keep track of completed payments.

 We also have to modify OnlineRepaymentCreated event by introducing a collection of payment identifiers, like so:

 public class OnlineRepaymentCreated : EventBase
 {
 public string UserId { get; init; }
 public string RepaymentId { get; init; }
 public IEnumerable<string> PaymentIds { get; init; }
 }
 

 Because of this, we need to introduce a small change to transition no. 1:

 // 1'. NotStarted : OnlineRepaymentCreated -> Created
 stateMachine
 .FromState<NotStarted>()
 .When<OnlineRepaymentCreated>()
 .ToState<Created>((from, @event) => new Created(from)
 {
 UserId = @event.UserId,
 RepaymentId = @event.RepaymentId,
 PaymentIds = new HashSet<string>(@event.PaymentIds)
 });
 

 Transition no. 2 is similarly adjusted to these changes by passing the set of payment identifiers to the bookkeeping system:

 (from, @event, to) => new RegisterPaymentsCommand(to.PaymentIds)
 

 Now we can move to the remaining transitions. Transition no. 5 has to be split into two. We want to stay in the Paid state until all registration events are received and only then switch to Registered. This can be achieved with:

 // 5'. Paid : PaymentRegistered -> Paid
 stateMachine
 .FromState<Paid>()
 .When<PaymentRegistered>((from, @event) => !from.PaymentIds.All(
 id => from.RegisteredPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState<Paid>((from, @event) =>
 {
 var to = new Paid(from);
 to.RegisteredPaymentIds.Add(@event.PaymentId);
 return to;
 });
 
 // 5''. Paid : PaymentRegistered -> Registered
 stateMachine
 .FromState<Paid>()
 .When<PaymentRegistered>((from, @event) => from.PaymentIds.All(
 id => from.RegisteredPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState<Registered>((from, @event) =>
 {
 var to = new Registered(from);
 to.RegisteredPaymentIds.Add(@event.PaymentId);
 return to;
 });
 .RunCommand((from, @event, to) => new SendRepaymentRegisteredEmailCommand(to.UserId, to.RepaymentId));
 

 As you can see, the command is now defined only for the second transition that triggers after all payments are registered. We have to slightly adjust transition no. 8, which is responsible for sending the email, too:

 // 8'. Completed : PaymentRegistered -> Completed
 stateMachine
 .FromState<Completed>()
 .When<PaymentRegistered>()
 .ToState<Completed>((from, @event) =>
 {
 var to = new Completed(from);
 to.RegisteredPaymentIds.Add(@event.PaymentId);
 return to;
 })
 .RunCommand(
 when: (from, @event, to) => to.PaymentIds.SetEquals(to.RegisteredPaymentIds),
 then: (from, @event, to) => new SendRepaymentRegisteredEmailCommand(to.UserId, to.RepaymentId));
 

 This time the transition is not conditional, however the command is. It will execute only after we handle final registration event and the payment identifier sets PaymentIds and RegisteredPaymentIds become equal.

 We cannot forget about all transitions that are triggered by PaymentCompleted event. Those will have to be changed like so:

 // 6'. Registered : PaymentCompleted -> Registered
 stateMachine
 .FromState<Registered>()
 .When<PaymentCompleted>((from, @event) => !from.PaymentIds.All(
 id => from.CompletedPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState<Registered>((from, @event) =>
 {
 var to = new Registered(from);
 to.CompletedPaymentIds.Add(@event.PaymentId);
 return to;
 });
 
 // 6''. Registered : PaymentCompleted -> Completed
 stateMachine
 .FromState<Registered>()
 .When<PaymentCompleted>((from, @event) => from.PaymentIds.All(
 id => from.CompletedPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState<Completed>((from, @event) =>
 {
 var to = new Completed(from);
 to.CompletedPaymentIds.Add(@event.PaymentId);
 return to;
 });
 
 // 7'. Paid : PaymentCompleted -> Paid
 stateMachine
 .FromState<Paid>()
 .When<PaymentCompleted>((from, @event) => !from.PaymentIds.All(
 id => from.CompletedPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState<Completed>((from, @event) =>
 {
 var to = new Completed(from);
 to.CompletedPaymentIds.Add(@event.PaymentId);
 return to;
 });
 
 // 7''. Paid : PaymentCompleted -> Completed
 stateMachine
 .FromState<Paid>()
 .When<PaymentCompleted>((from, @event) => from.PaymentIds.All(
 id => from.CompletedPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState<Completed>((from, @event) =>
 {
 var to = new Completed(from);
 to.CompletedPaymentIds.Add(@event.PaymentId);
 return to;
 });
 

 And we’re done! This time for real. After all this work we ended up with something that can be visualized with a diagram like this:

 I mentioned earlier that we have internal back-office tools to view state machines and their transition history. Here is an example view of a repayment process that was coordinated with the state machine we defined above (sensitive data was redacted, since this was taken in production):

 Caveats

 You may have noticed some problems here. What would happen if we received an event with a wrong payment identifier? Right now the state machine is not ready for that. However, that can be remedied relatively easily by adding another condition to When methods of all transitions for PaymentRegistered and PaymentCompleted events, such as:

 .When((from, @event) => from.PaymentIds.Contains(@event.PaymentId) && /* ... */)
 

 This would essentialy ensure that we react only to events with correct identifiers.

 And what about retries and duplicated events? Transition no. 8’ is especially susceptible to that, because it may result in the email being sent to the user twice. And we don’t want that! Fortunately, achieving deduplication is just as easy — we have to ignore events for payments that are already registered:

 .When((from, @event) => !from.RegisteredPaymentIds.Contains(@event.PaymentId) && /* ... */)
 

 For the sake of readability I also decided to skip a lot of other validations, which would otherwise have to be there, such as simple null checks.

 How to test

 Testing state machines is important, because, well, ideally all code should be tested. In this case, apart from the usual assurance that we are able to handle both happy paths and edge cases, we can also ensure that we keep the transitions backwards compatible. An event that is once inserted into the stream will stay there forever.

 We developed a simple DSL in F# to help us test state machines created with our framework. This is a perfect solution to test a declarative definition of a state machine like the one we developed in this post. Here’s an example of a happy path test for our state machine:

 [<Fact>]
 let ``Online repayment - single payment happy path`` () =
 [
 When(Events.OnlineRepaymentCreated ("paymentId1"))
 |> GoToState<Created>
 |> DoNothing
 When(Events.OnlineRepaymentPaid ())
 |> GoToState<Paid>
 |> Do<RegisterPaymentsCommand>
 When(Events.PaymentRegistered ("paymentId1"))
 |> GoToState<Registered>
 |> Do<SendRepaymentRegisteredEmailCommand>
 When(Events.PaymentCompleted ("paymentId1"))
 |> GoToState<Completed>
 |> DoNothing
 ]
 |> apply
 

 A test like this is essentially a list of events with expected side effects that are then aggregated using the apply function that feeds them one by one into the state machine and verifies the outcome.

 Summary

 I hope I was able to convince you how a framework like this can make designing complex state machines easy. It provides numerous advantages, such as declarative syntax, effortless testability, thorough instrumentation and visualization.

 This is precisely why state machines became an integral part of pretty much every business process architecture in our project. Having such a robust and scalable solution for asynchronous process orchestration and isolation is crucial in distributed systems. The presented code, albeit extremely simplified, is from an actual state machine — one of many used in our applications.

 While the presented solution is a strictly internal framework right now, who knows what the future might bring? Stay tuned!

 Tymon Felski

 Tymon is a software developer who likes doing back-end stuff. He is highly skilled at applying screen protectors.

 see post by Tymon Felski

 * Share this post
 * 
 * 
Please enable JavaScript to view the comments powered by Disqus. comments powered by Disqus
 Proudly built by engineers
Coordinating complex processes, both business and technical, can be a challenging issue in a distributed system.
Especially when the complications associated with them, such as concurrency, idempotency, scalability and hindered
testability, come into play — possibly all at once.
This is definitely something that can keep many programmers awake at night.
While this may sound dramatic, in reality there are many different solutions to this problem.
One group of such solutions extensively uses finite-state machines also known as finite-state automata.
As this article is meant for everyone, let’s start with some background information about what they are and how
they work.
What are finite-state machines
For the sake of readability, I will sometimes be referring to finite-state machines as simply “state machines”.
If you’re already familiar with their formal definition, you may skip to the next section.
Formally speaking, a finite-state machine is a mathematical model of computation that describes an abstract system
having a finite number of permissible states. At any given point in time, the system is in exactly one of these
states.
In practical terms such a machine is described by:
 an initial (start) state,
 a set of possible states,
 a sequence of possible inputs (events),
 a transition function which based on the last observed input and the machine’s current state, determines the next
state,
 and a set of terminal (end) states — optionally.
The machine starts off in the initial state and inspects all inputs in sequence.
Upon observing each input, it uses the transition function to change its state.
The processing ends once all inputs have been observed.
If the machine’s state at the end of processing is one of the terminal states, the machine is assumed to accept the
given input, and to reject it otherwise.
As a side note, the description provided above uses the deterministic model of a finite-state automaton.
There are also alternative, non-deterministic machines, but I won’t be delving into them too much, as they can be
converted to deterministic ones anyway.
If you’ve read this far, hopefully the rest of this post will be easier on the mind!
Despite the intimidating formal definition, finite-state machines are more widespread than you might imagine.
Upon closer inspection you might be able to spot them in your everyday life under the hood of things such as vending
machines, traffic lights and elevators.
In .NET they are used in the implementation of async/await or yield syntax.
Try decompiling your code and see how the language authors implemented those features.
You have probably already used some sort of state machine in your code without even realising it!
A simple example of a basic state machine implementation
Let’s take a closer look at the following piece of code:
  public enum Button
{
 Play,
 Stop
}
public enum State
{
 Playing,
 NotPlaying
}
public class Player
{
 public State CurrentState { get; private set; } = State.NotPlaying;
 public void HandleClick(Button button)
 {
 if (button == Button.Play)
 {
 CurrentState = State.Playing;
 }
 else if (button == Button.Stop)
 {
 CurrentState = State.NotPlaying;
 }
 }
}
  
This is an example of a very simple state machine with two states and two events.
The player can be either playing or not playing music.
Pressing the Play button will start or keep playing music, while pressing the Stop button will stop the music or keep
it stopped.
The code is manageable at this point, since the example is quite straightforward.
However, I hope you can imagine how adding more media controls that are usually present in music players will quickly
result in complex code that will not be easily maintainable anymore.
For instance, consider a feature request to have the Previous button change to the previous track if the player is less
than 3 seconds into the current track, or stay at the current track but rewind it to the beginning otherwise.
That being said, in many cases this kind of state machine implementation would be more than enough.
You don’t always have to use a state-of-the-art, bleeding-edge solution to achieve desired goals.
However, this is not something I would recommend to someone who is trying to tackle coordination of asynchronous
operations within a distributed system.
A more maintainable approach
Let me share our approach to this problem.
We have come up with a framework for building state machines in .NET with very little code that is maintainable and
testable.
State machines built with this framework serve a purpose of orchestrating a wide range of business and technical
processes in our microservice-based distributed architecture.
The framework is equipped with various features dictated by paradigms of distributed systems, which extend the
standard state machine definition.
Before I move on to an example of its real-life application, I want to break down the state machine’s API, so that you
know what you’re looking at later on.
To keep things short, I will be showing you only snippets of simplified code, as implementation details can be
expanded upon in follow-up posts.
The state machine is defined as follows:
  public class StateMachine&lt;TStateBase, TEventBase&gt;
{
 public StateMachine(TStateBase initialState)
 {
 // ...
 }
}
  
It is a generic class that requires a definition of base types for all states (TStateBase) and for all transition
triggers (TEventBase).
The constructor accepts an initial state from which the state machine will start its operation.
We can create the state machine like so:
  var stateMachine = new StateMachine&lt;StateBase, EventBase&gt;(new InitialState());
  
Having created this object, we may start defining transitions.
Each transtion consists of three required elements:
 the state from which the transition can be triggered,
 the transition trigger (an event),
 the state to which the transition leads.
Optionally, a transition can have side effects, called simply commands.
A command (or a set of commands) will be run after a successful transition.
Those can be pretty much anything from sending a message to another state machine to calling a specific service and
delegating a task to it.
Commands allow a seamless integration of the state machine with external components.
There is a drawback, however.
Command execution may fail, what requires an idempotent retry policy to be employed.
Transitions and commands can be conditional, meaning that a logical expression can be guarding and preventing the
execution of the whole transition, or just a command in some cases.
A builder pattern which combines all of the aforementioned requirements is used to compose a transition:
  stateMachine
 .FromState&lt;TStateFrom&gt;()
 .When&lt;TEvent&gt;(/* Expression&lt;Func&lt;TStateFrom, TEvent, bool&gt;&gt; */)
 .ToState&lt;TStateTo&gt;(/* Expression&lt;Func&lt;TStateFrom, TEvent, TStateTo&gt;&gt; */)
 .RunCommand(
 when: /* Expression&lt;Func&lt;TStateFrom, TEvent, TStateTo, bool&gt;&gt; */,
 then: /* Expression&lt;Func&lt;TStateFrom, TEvent, TStateTo, TCommand&gt;&gt; */)
 .RunCommand(/* ... */)
 // ...
 .RunCommand(/* ... */);
  
This structure may be overwhelming at first, but bear with me.
It will all make sense in a bit.
As for now you need to know that:
 if the state machine is in a state of TStateFrom type,
 and an event of type TEvent occured,
 and the condition (if defined) passed to When method is satisfied,
the transition will fire up and the state machine will switch to a new state of type TStateTo using the factory method
expression passed to ToState method.
The state machine class exposes a method which can be used to trigger a transition with an event:
  public TransitionResult Apply(TEventBase @event);
  
This method applies the supplied event to the current state by trying to find a matching transition in the transition
mapping and changes the internal state of the machine.
The result contains information about the performed transition, along with commands that should be run.
  public class TransitionResult
{
 public TStateBase StateFrom { get; }
 public TEventBase Event { get; }
 public TStateBase StateTo { get; }
 public IEnumerable&lt;object&gt; Commands { get; }
 public bool IsValid { get; }
}
  
If no matching transition definition is found given the machine’s current state – or if a transition is found, but its
precondition is not met – the event application will result in an invalid transition.
In that case, the state machine will not react to the event and simply retain its state, possibly logging that an
invalid event was received.
If there are multiple matching transitions, the first one will fire up (in definition order).
We should, however, avoid designing states machines in such way if possible.
As you can see, the representation of the machine is primarily type-based, so we can use polymorphism to our advantage.
By calling FromState&lt;TStateBase&gt;() we can have a transition that can be fired up from any state.
Similarily, if we use a type that only some selected states derive from, we will have a transition applicable to these
states only.
It would be equivalent to duplicating that same transition for every one of these states.
This syntax can be particularly useful, when dealing with processes that have an expiration date and have to be
completed within a specific time frame.
Receiving an event about the process’ expiration may have to be handled regardless of the machine’s current state and
result in termination immediately.
Implementation manifest
Now that I’ve described how a state machine can be defined in a declarative way using our framework, let’s talk a bit
about how that definition actually runs.
In general, we use dependency injection to make the state machine definition work with other services and add-ons,
such as specific command runners.
I promised I would not dive deep into implementation details, but nonetheless I want to give you a little taste of
what sits behind the scenes.
If you don’t care for any of that, you may skip to the “Real-life example” section.
Storage
The state machine implementation is based on event sourcing.
In summary, this means that we don’t save the state of our application objects, but rather a series of events that
change the object’s state.
These events will give us the current view of the object when aggregated (“replayed”) in the order in which they
occured.
For this purpose we currently use an open-source library,
SQLStreamStore.
SQLStreamStore offers an atomic and idempotent stream append operation.
Thanks to this, we don’t have to worry about race conditions between multiple instances of the same state machine
handling events in parallel, what is great for scalability.
We can also recognize a situation when the same event is being fed into the state machine again, what can happen
when dealing with retry policies.
That, on the other hand, gives us a pretty good way of achieving deduplication.
Caching
Replaying all events every single time is an unnecessary overhead that we decided to mitigate by introducing
a lock-free cache for state snapshots.
After each transition a state machine will cache its current state along with the stream version which it
corresponds to.
This gives us the ability to restore the state machine to the state it was in after handling n-1 events, when we want
to handle the n-th event.
It requires, however, the states to be immutable, or else we may end up with an unexpected behaviour.
If states were not immutable, we could not be certain that the state object we cached after handling the n-th event
is still the same and has not been modified in the meantime.
Moreover, immutable objects are inherently thread-safe.
Instrumentation
This state machine framework is also equipped with heavy instrumentation.
Every executed transition and state change is logged, so we can easily track any errors.
This data is also useful for our data analytics team to keep track of business processes.
A lot of different metrics come into play to measure performance of each and every state machine deployed.
Finally, every historical transition is saved in a stream — separate from the event stream — and can be viewed with our
internal back-office tools.
Thanks to the finite nature of the state machines, we can always log the current state and trace what made the system
to be in that state, along with all its previous states.
This is something quite invaluable and helps immensely when diagnosing issues, especially in complex and
inter-dependent business scenarios where a lot of things happen at the same time.
Real-life example
So finally, the big question: how is this framework used on our platform?
To answer that we don’t need to reach far.
In July of 2020 we launched a new payment method called Allegro Pay, which will eventually be available to every buyer
on Allegro.
This service allows the users to buy items now and pay for them later in a single payment after 30 days, or
in multiple smaller monthly payments, depending on the value of purchased goods.
Repayments can be easily made online using our Allegro Pay dashboard.
Users are even able to pay for multiple purchases at once.
 
For more detailed information, I encourage you to read our
Allegro Pay FAQ.
I suspect you already know where I’m going with this — my goal is to show you how our state machine framework is used
to coordinate an actual business process of handling a repayment within Allegro Pay.
The process specification
Before we create our state machine, we need to have a rough idea of which states and events we want to handle.
The repayment process is started in two different ways, but will end alike in both cases.
 
 The user can initialize an immediate online repayment via the Allegro Pay dashboard.
 They are redirected to the payment provider’s website and complete the process there.
 
 
 Alternatively, the user can request to repay their purchase with a traditional wire transfer, what we call an
offline repayment.
 We know nothing about it until the money actually arrives at the target bank account.
 That method may take up to a couple of days.
 
At this point both repayment paths merge, since from now on we can discard any knowledge about how the money
was received and focus on the fact that we have it.
The money source won’t be useful during the process coordination anymore.
The next step is to register the repayment in our bookkeeping system.
The registration part is important for us, because it will tell us whether the purchase was paid in full or if further
repayments are required.
After receiving information about successful registration, an email is sent to the user confirming that the
repayment was registered.
After repayment is registered, we asynchronously await the feedback from the bookkeeping system on the transaction
being settled and the repayment is to be marked as completed.
Designing the state machine
I could show you the finished state machine and just paste a wall of text here, describing what it does, but where’s
the fun in that?
Instead, I’d like to show you the step-by-step process I would go through when designing this kind of state machine.
I’ll start off by simply creating a state machine that handles a repayment of a single purchase and then extend it
to work for all cases.
First we need to define the base types for all states and events:
  public abstract class StateBase
{
 public string UserId { get; init; }
 public string RepaymentId { get; init; }
 public string PaymentId { get; init; }
 public StateBase(StateBase other)
 {
 if (other == null) return;
 UserId = other.UserId;
 RepaymentId = other.RepaymentId;
 PaymentId = other.PaymentId;
 }
}
public abstract class EventBase { }
  
As you can see, I already equipped the base state with all the properties we will need later to orchestrate this
process:
  UserId — tells us who is repaying,
  RepaymentId — uniquely identifies the repayment process,
  PaymentId — identifies the payment made by the user for this repayment.
As for the event base, we don’t need anything special there, so it’s empty.
 StateBase class features a copy constructor, which rewrites all properties from the previous state when creating
a new one.
None of the states will have any special properties, so there’s no point in painstakingly showing all individual
subclasses representing each state.
Here are just a few sample classes of the initial state (NotStarted) and one of the subsequent states:
  public class NotStarted : StateBase
{
 public NotStarted() : base(null) { }
}
public class Created : StateBase
{
 public Created(StateBase other) : base(other) { }
}
  
All of the remaining states are defined similarly to Created state.
Those states are:
  Paid,
  Failed,
  Registered,
  Completed.
Once we have our states, we can focus on events and introduce transitions, one by one.
We’ll start by handling an online repayment.
It begins with the user creating a repayment entity by selecting a payment in the Allegro Pay dashboard.
This generates an event, let’s call it OnlineRepaymentCreated:
  public class OnlineRepaymentCreated : EventBase
{
 public string UserId { get; init; }
 public string RepaymentId { get; init; }
 public string PaymentId { get; init; }
}
  
We want this event to trigger a transition from NotStarted to Created state.
  // 1. NotStarted : OnlineRepaymentCreated -&gt; Created
stateMachine
 .FromState&lt;NotStarted&gt;()
 .When&lt;OnlineRepaymentCreated&gt;()
 .ToState&lt;Created&gt;((from, @event) =&gt; new Created(from)
 {
 UserId = @event.UserId,
 RepaymentId = @event.RepaymentId,
 PaymentId = @event.PaymentId
 });
  
At this point the user sees the repayment form, where they select a payment method.
They are then redirected to the payment provider’s website to finish the repayment.
In most cases it will succeed, but in some it fails due to a multitude of reasons, such as insufficient
funds on the account or the user providing wrong confirmation code.
Therefore, we need two events:
  public class OnlineRepaymentPaid : EventBase { }
public class OnlineRepaymentFailed : EventBase { }
  
We could go about a single event called RepaymentResult with an enum property indicating a success or a failure, but
in my opinion the former approach is cleaner and more open to future changes.
These events will play a role in tranistions no. 2 and 3:
  // 2. Created : OnlineRepaymentPaid -&gt; Paid
stateMachine
 .FromState&lt;Created&gt;()
 .When&lt;OnlineRepaymentPaid&gt;()
 .ToState&lt;Paid&gt;((from, @event) =&gt; new Paid(from))
 .RunCommand((from, @event, to) =&gt; new RegisterPaymentCommand(to.PaymentId));
// 3. Created : OnlineRepaymentFailed -&gt; Failed
stateMachine
 .FromState&lt;Created&gt;()
 .When&lt;OnlineRepaymentFailed&gt;()
 .ToState&lt;Failed&gt;((from, @event) =&gt; new Failed(from));
  
You might notice that transition no. 2 has a command that I didn’t mention earlier specified.
It’s a simple class wrapping the payment identifier, which is interpreted by the command runner as a request to our
bookkeeping system for registering this payment.
This concludes the online path and we can move on to the offline repayment.
It is significantly easier, since we are simply notified that the money has been transferred to us.
Because of this, we can completely skip the Created state and go straight to Paid.
We need an event called OfflineRepaymentPaid which contains the same properties as OnlineRepaymentCreated:
  public class OfflineRepaymentPaid : EventBase
{
 public string UserId { get; init; }
 public string RepaymentId { get; init; }
 public string PaymentId { get; init; }
}
  
Let’s use it in the fourth transition that will once again call the registration command:
  // 4. NotStarted : OfflineRepaymentPaid -&gt; Paid
stateMachine
 .FromState&lt;NotStarted&gt;()
 .When&lt;OfflineRepaymentPaid&gt;()
 .ToState&lt;Paid&gt;((from, @event) =&gt; new Paid(from)
 {
 UserId = @event.UserId,
 RepaymentId = @event.RepaymentId,
 PaymentId = @event.PaymentId
 })
 .RunCommand((from, @event, to) =&gt; new RegisterPaymentCommand(@event.PaymentId));
  
This is where both repayment types merge and we can focus on our bookkeeping system.
As mentioned before, we are mostly interested in registration events published by the bookkeeping system.
  public class PaymentRegistered : EventBase
{
 public string PaymentId { get; init; }
}
  
Once we know that a payment is registered in our bookkeeping system, we are able to send an email to the user
summarizing what they paid for and follow-up information whether they still have more payments to be made.
This can be done with a transition like this:
  // 5. Paid : PaymentRegistered -&gt; Registered
stateMachine
 .FromState&lt;Paid&gt;()
 .When&lt;PaymentRegistered&gt;()
 .ToState&lt;Registered&gt;((from, @event) =&gt; new Registered(from))
 .RunCommand((from, @event, to) =&gt; new SendRepaymentRegisteredEmailCommand(to.UserId, to.RepaymentId));
  
Again, SendRepaymentRegisteredEmailCommand is a simple class wrapping two properties that will be interpreted by the
command runner and result in an email sent to the user.
The last step is to wait for the final event PaymentCompleted.
Theoretically it could be omitted in this state machine, but it’s useful for auditing purposes.
  public class PaymentCompleted : EventBase
{
 public string PaymentId { get; init; }
}
  
It’s time to create the final transition:
  // 6. Registered : PaymentCompleted -&gt; Completed
stateMachine
 .FromState&lt;Registered&gt;()
 .When&lt;PaymentCompleted&gt;()
 .ToState&lt;Completed&gt;((from, @event) =&gt; new Completed(from));
  
And we’re done!
Right?
Not really.
As is usually the case with event-sourced systems, the order of incoming events can’t always be relied on.
The bookkeeping system will sometimes — for reasons I don’t want to delve into — publish events PaymentRegistered and
PaymentCompleted at the same time.
Due to a lack of message ordering guarantees, it’s possible for us to receive PaymentCompleted event before
PaymentRegistered.
The current definition of the state machine is not prepared for that.
Fortunately, the fix is quite easy, as we can simply introduce two more transitions:
  // 7. Paid : PaymentCompleted -&gt; Completed
stateMachine
 .FromState&lt;Paid&gt;()
 .When&lt;PaymentCompleted&gt;()
 .ToState&lt;Completed&gt;((from, @event) =&gt; new Completed(from));
// 8. Completed : PaymentRegistered -&gt; Completed
stateMachine
 .FromState&lt;Completed&gt;()
 .When&lt;PaymentRegistered&gt;()
 .ToState&lt;Completed&gt;((from, @event) =&gt; new Completed(from))
 .RunCommand((from, @event, to) =&gt; new SendRepaymentRegisteredEmailCommand(to.UserId, to.RepaymentId));
  
In this scenario we completely skip Registered state, but that’s something we just have to deal with.
Events are important from analytics and diagnostics standpoints, states — not so much.
That wraps up the event flow for this scenario.
It can be visualized with a diagram like this:
 
Our state machine framework is equipped with a pretty useful visualisation extension.
It’s a pretty neat tool that makes it easy to figure out how the state machine works, without having to delve into the
code.
A simple unit test that calls this extension can render a state machine diagram like one above using PlantUML
syntax, which in this case would go as follows:
  @startuml
hide empty description
[*] --&gt; NotStarted
NotStarted --&gt; Created: &lt;b&gt;OnlineRepaymentCreated&lt;/b&gt;
Created --&gt; Paid: &lt;b&gt;OnlineRepaymentPaid&lt;/b&gt;\n&lt;i&gt;then:&lt;/i&gt; RegisterPaymentCommand
Created --&gt; Failed: &lt;b&gt;OnlineRepaymentFailed&lt;/b&gt;
NotStarted --&gt; Paid: &lt;b&gt;OfflineRepaymentPaid&lt;/b&gt;\n&lt;i&gt;then:&lt;/i&gt; RegisterPaymentCommand
Paid --&gt; Registered: &lt;b&gt;PaymentRegistered&lt;/b&gt;\n&lt;i&gt;then:&lt;/i&gt; SendRepaymentRegisteredEmailCommand
Registered --&gt; Completed: &lt;b&gt;PaymentCompleted&lt;/b&gt;
Paid --&gt; Completed: &lt;b&gt;PaymentCompleted&lt;/b&gt;
Completed --&gt; Completed: &lt;b&gt;PaymentRegistered&lt;/b&gt;\n&lt;i&gt;then:&lt;/i&gt; SendRepaymentRegisteredEmailCommand
@enduml
  
However, recall that this is just a state machine to handle a repayment for a single payment selected by the user.
I promised we would extend it to handle multiple payments.
Let’s do that.
First we’ll extend the base state definition:
  public abstract class StateBase
{
 public string UserId { get; init; }
 public string RepaymentId { get; init; }
 public HashSet&lt;string&gt; PaymentIds { get; init; }
 public HashSet&lt;string&gt; RegisteredPaymentIds { get; init; }
 public HashSet&lt;string&gt; CompletedPaymentIds { get; init; }
 public StateBase(StateBase other)
 {
 if (other == null) return;
 UserId = other.UserId;
 RepaymentId = other.RepaymentId;
 PaymentIds = new HashSet&lt;string&gt;(other.PaymentIds);
 RegisteredPaymentIds = new HashSet&lt;string&gt;(other.RegisteredPaymentIds);
 CompletedPaymentIds = new HashSet&lt;string&gt;(other.CompletedPaymentIds);
 }
}
  
Three things have changed:
  PaymentId changed from a string to a HashSet&lt;string&gt; and is now called PaymentIds to represent multiple
payments selected by the user,
 a HashSet&lt;string&gt; called RegisteredPaymentIds was added to keep track of registered payments,
 a HashSet&lt;string&gt; called CompletedPaymentIds was added to keep track of completed payments.
We also have to modify OnlineRepaymentCreated event by introducing a collection of payment identifiers, like so:
  public class OnlineRepaymentCreated : EventBase
{
 public string UserId { get; init; }
 public string RepaymentId { get; init; }
 public IEnumerable&lt;string&gt; PaymentIds { get; init; }
}
  
Because of this, we need to introduce a small change to transition no. 1:
  // 1'. NotStarted : OnlineRepaymentCreated -&gt; Created
stateMachine
 .FromState&lt;NotStarted&gt;()
 .When&lt;OnlineRepaymentCreated&gt;()
 .ToState&lt;Created&gt;((from, @event) =&gt; new Created(from)
 {
 UserId = @event.UserId,
 RepaymentId = @event.RepaymentId,
 PaymentIds = new HashSet&lt;string&gt;(@event.PaymentIds)
 });
  
Transition no. 2 is similarly adjusted to these changes by passing the set of payment identifiers to the
bookkeeping system:
  (from, @event, to) =&gt; new RegisterPaymentsCommand(to.PaymentIds)
  
Now we can move to the remaining transitions.
Transition no. 5 has to be split into two.
We want to stay in the Paid state until all registration events are received and only then switch to Registered.
This can be achieved with:
  // 5'. Paid : PaymentRegistered -&gt; Paid
stateMachine
 .FromState&lt;Paid&gt;()
 .When&lt;PaymentRegistered&gt;((from, @event) =&gt; !from.PaymentIds.All(
 id =&gt; from.RegisteredPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState&lt;Paid&gt;((from, @event) =&gt;
 {
 var to = new Paid(from);
 to.RegisteredPaymentIds.Add(@event.PaymentId);
 return to;
 });
// 5''. Paid : PaymentRegistered -&gt; Registered
stateMachine
 .FromState&lt;Paid&gt;()
 .When&lt;PaymentRegistered&gt;((from, @event) =&gt; from.PaymentIds.All(
 id =&gt; from.RegisteredPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState&lt;Registered&gt;((from, @event) =&gt;
 {
 var to = new Registered(from);
 to.RegisteredPaymentIds.Add(@event.PaymentId);
 return to;
 });
 .RunCommand((from, @event, to) =&gt; new SendRepaymentRegisteredEmailCommand(to.UserId, to.RepaymentId));
  
As you can see, the command is now defined only for the second transition that triggers after all payments are
registered.
We have to slightly adjust transition no. 8, which is responsible for sending the email, too:
  // 8'. Completed : PaymentRegistered -&gt; Completed
stateMachine
 .FromState&lt;Completed&gt;()
 .When&lt;PaymentRegistered&gt;()
 .ToState&lt;Completed&gt;((from, @event) =&gt;
 {
 var to = new Completed(from);
 to.RegisteredPaymentIds.Add(@event.PaymentId);
 return to;
 })
 .RunCommand(
 when: (from, @event, to) =&gt; to.PaymentIds.SetEquals(to.RegisteredPaymentIds),
 then: (from, @event, to) =&gt; new SendRepaymentRegisteredEmailCommand(to.UserId, to.RepaymentId));
  
This time the transition is not conditional, however the command is.
It will execute only after we handle final registration event and the payment identifier sets PaymentIds and
RegisteredPaymentIds become equal.
We cannot forget about all transitions that are triggered by PaymentCompleted event.
Those will have to be changed like so:
  // 6'. Registered : PaymentCompleted -&gt; Registered
stateMachine
 .FromState&lt;Registered&gt;()
 .When&lt;PaymentCompleted&gt;((from, @event) =&gt; !from.PaymentIds.All(
 id =&gt; from.CompletedPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState&lt;Registered&gt;((from, @event) =&gt;
 {
 var to = new Registered(from);
 to.CompletedPaymentIds.Add(@event.PaymentId);
 return to;
 });
// 6''. Registered : PaymentCompleted -&gt; Completed
stateMachine
 .FromState&lt;Registered&gt;()
 .When&lt;PaymentCompleted&gt;((from, @event) =&gt; from.PaymentIds.All(
 id =&gt; from.CompletedPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState&lt;Completed&gt;((from, @event) =&gt;
 {
 var to = new Completed(from);
 to.CompletedPaymentIds.Add(@event.PaymentId);
 return to;
 });
// 7'. Paid : PaymentCompleted -&gt; Paid
stateMachine
 .FromState&lt;Paid&gt;()
 .When&lt;PaymentCompleted&gt;((from, @event) =&gt; !from.PaymentIds.All(
 id =&gt; from.CompletedPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState&lt;Completed&gt;((from, @event) =&gt;
 {
 var to = new Completed(from);
 to.CompletedPaymentIds.Add(@event.PaymentId);
 return to;
 });
// 7''. Paid : PaymentCompleted -&gt; Completed
stateMachine
 .FromState&lt;Paid&gt;()
 .When&lt;PaymentCompleted&gt;((from, @event) =&gt; from.PaymentIds.All(
 id =&gt; from.CompletedPaymentIds.Contains(id) || id == @event.PaymentId))
 .ToState&lt;Completed&gt;((from, @event) =&gt;
 {
 var to = new Completed(from);
 to.CompletedPaymentIds.Add(@event.PaymentId);
 return to;
 });
  
And we’re done!
This time for real.
After all this work we ended up with something that can be visualized with a diagram like this:
 
I mentioned earlier that we have internal back-office tools to view state machines and their transition history.
Here is an example view of a repayment process that was coordinated with the state machine we defined above (sensitive
data was redacted, since this was taken in production):
 
Caveats
You may have noticed some problems here.
What would happen if we received an event with a wrong payment identifier?
Right now the state machine is not ready for that.
However, that can be remedied relatively easily by adding another condition to When methods of all transitions for
PaymentRegistered and PaymentCompleted events, such as:
  .When((from, @event) =&gt; from.PaymentIds.Contains(@event.PaymentId) &amp;&amp; /* ... */)
  
This would essentialy ensure that we react only to events with correct identifiers.
And what about retries and duplicated events?
Transition no. 8’ is especially susceptible to that, because it may result in the email being sent to the user twice.
And we don’t want that!
Fortunately, achieving deduplication is just as easy — we have to ignore events for payments that are already
registered:
  .When((from, @event) =&gt; !from.RegisteredPaymentIds.Contains(@event.PaymentId) &amp;&amp; /* ... */)
  
For the sake of readability I also decided to skip a lot of other validations, which would otherwise have to be there,
such as simple null checks.
How to test
Testing state machines is important, because, well, ideally all code should be tested.
In this case, apart from the usual assurance that we are able to handle both happy paths and edge cases, we can also
ensure that we keep the transitions backwards compatible.
An event that is once inserted into the stream will stay there forever.
We developed a simple DSL in F# to help us test state machines created with our framework.
This is a perfect solution to test a declarative definition of a state machine like the one we developed in this post.
Here’s an example of a happy path test for our state machine:
  [&lt;Fact&gt;]
let ``Online repayment - single payment happy path`` () =
 [
 When(Events.OnlineRepaymentCreated ("paymentId1"))
 |&gt; GoToState&lt;Created&gt;
 |&gt; DoNothing
 When(Events.OnlineRepaymentPaid ())
 |&gt; GoToState&lt;Paid&gt;
 |&gt; Do&lt;RegisterPaymentsCommand&gt;
 When(Events.PaymentRegistered ("paymentId1"))
 |&gt; GoToState&lt;Registered&gt;
 |&gt; Do&lt;SendRepaymentRegisteredEmailCommand&gt;
 When(Events.PaymentCompleted ("paymentId1"))
 |&gt; GoToState&lt;Completed&gt;
 |&gt; DoNothing
 ]
 |&gt; apply
  
A test like this is essentially a list of events with expected side effects that are then aggregated using the apply
function that feeds them one by one into the state machine and verifies the outcome.
Summary
I hope I was able to convince you how a framework like this can make designing complex state machines easy.
It provides numerous advantages, such as declarative syntax, effortless testability, thorough instrumentation and
visualization.
This is precisely why state machines became an integral part of pretty much every business process architecture
in our project.
Having such a robust and scalable solution for asynchronous process orchestration and isolation is crucial in
distributed systems.
The presented code, albeit extremely simplified, is from an actual state machine — one of many used in our
applications.
While the presented solution is a strictly internal framework right now, who knows what the future might bring?
Stay tuned!
Finite-state machines made easy
https://blog.allegro.tech/2021/03/state-machines-made-easy.html