T RAINING N EURAL M ACHINES WITH P ARTIAL T RACES
Anonymous authors
Paper under double-blind review
A BSTRACT
We present a novel approach for training neural machines which incorporates additional supervision on the machine’s interpretable components (e g., neural memory).
To cleanly capture the kind of neural machines to which our method applies, we introduce the concept of a differential neural computational machine ( ∂ NCM) and show that several existing architectures (e g., NTMs, NRAMs) can be instantiated as a ∂ NCM and can thus beneﬁt from any amount of additional supervision over their interpretable components.
Based on our method, we performed a detailed experimental evaluation with NTM and NRAM machines, showing the approach leads to signiﬁcantly better convergence and generalization capabilities of the learning phase than standard training using only input-output examples
I NTRODUCTION
Recently, there has been substantial interest in neural abstract machines that can induce programs from examples Feser et al.
(2015; 2016); Frankle et al.
; Gaunt et al.
; Graves et al.
(2014; 2016); Kaiser & Sutskever (2016); Kurach et al.
Reed & de Freitas (2016); Vinyals et al.
; Zaremba & Sutskever (2015); Zhang et al.
While signiﬁcant progress has been made towards learning interesting algorithms Graves et al.
(2016), ensuring the training of these machines converges to the desired solution can be very challenging.
Interestingly however, even though these machines differ architecturally, they tend to rely on components (e.g., neural memory) that are more interpretable than a typical neural network (e.g., an LSTM).
A key question then is:
Can we somehow provide additional amounts of supervision for these interpretable components during training so to bias the learning towards the desired solution?
In this work we investigate this question in depth.
We refer to the type of supervision mentioned above as partial trace supervision, capturing the intuition that more detailed information, beyond inputoutput examples, is provided during learning.
To study the question systematically, we introduce the notion of a differential neural computational machine ( ∂ NCM), a formalism which allows for clean characterization of the neural abstract machines that fall inside our class and that can beneﬁt from any amount of partial trace information.
We show that common architectures such as Neural Turing
Machines (NTMs) and Neural Random Access Machines (NRAMs) can be phrased as ∂ NCMs, useful also because these architectures form the basis for many recent extensions, e.g., Graves et al.
Grefenstette et al.
; Kaiser & Sutskever (2016).
We also explain why other machines such as the Neural Program Interpreter (NPI)
Reed & de Freitas (2016) or its recent extensions (e.g., the
Neural Program Lattice Li et al.
(2017)) cannot be instantiated as an ∂ NCM and are thus restricted to require large (and potentially prohibitive) amounts of supervision.
We believe the ∂ NCM abstraction is a useful step in better understanding how different neural abstract machines compare when it comes to additional supervision.
We then present ∂ NCM loss functions which abstractly capture the concept of partial trace information and show how to instantiate these for both the NTM and the
NRAM.
We also performed an extensive evaluation for how partial trace information affects training in both architectures.
Overall, our experimental results indicate that the additional supervision can substantially improve convergence while leading to better generalization and interpretability.
To provide an intuition for the problem we study in this work, consider the simple task of training an NTM to ﬂip the third bit in a bit stream (called Flip3rd ) – such bitstream tasks have been extensively studied in the area of program synthesis (e.g., Jha et al.
; Raychev et al.
An example input-output pair for this task could be [0 , 1 , 0 , 0] → [0 , 1 , 1 , 0] .
Given a set of such
[Figure]
examples, our goal is to train an NTM that solves this task.
An example NTM that generalizes well and is understandable is shown in Figure 1c.
Here, the y-axis is time (descending), the x-axis is the accessed memory location, the white squares represent the write head of the NTM, and the orange squares represent the read head.
As we can see, the model writes the input sequence to the tape and then reads from the tape in the same order.
However, in the absence of richer supervision, the NTM (and other neural architectures) can easily overﬁt to the training set – an example of an overﬁtting
NTM is shown in Figure 1a.
Here, the traces are chaotic and difﬁcult to interpret.
Further, even if the NTM generalizes, it can do so with erratic reads and writes, an example of which is shown in
Figure 1b.
Here, the NTM learns to read from the third bit (circled) with a smaller weight than from other locations, and also reads and writes erratically near the end of the sequence.
This model is less interpretable than the one in Figure 1c because it is unclear how the model knows which the third bit actually is, or why a different read weight would help ﬂip that bit.
In this work we will develop principled ways for guiding the training of a neural abstract machine towards the behavior shown in Figure 1c.
For instance, for Flip3rd , providing partial trace information on the NTM’s read heads for 10% of the input-output examples is sufﬁcient to bias the learning towards the NTM shown in Figure 1c 100% of the time.
N EURAL C OMPUTATIONAL M ACHINES
To capture the essence of our method and illustrate its applicability, we now deﬁne the abstract notion of a neural computational machine (NCM).
NCMs mimic classic computational machines with a controller and a memory, and generalize multiple existing architectures.
Our approach for supervision with partial trace information applies to all neural architectures expressible as NCMs.
A useful feature of the NCM abstraction is that it clearly delineates end-to-end differentiable architectures (Graves et al.
(2014)’s NTM, Kurach et al.
(2016)’s NRAM), which can train with little to no trace supervision, from architectures that are not end-to-end differentiable (Reed & de Freitas (2016)’s
NPI) and hence require a certain minimum amount of trace information.
In the follow-up section, we show how to phrase two existing neural architectures (NTMs and NRAMs) as an NCM.
An NCM is a triple of functions: a processor, a controller, and a loss:
Processor The processor π : W × C × M → B × M performs a pre-deﬁned set of commands C , which might involve manipulating memories in M .
The commands may produce additional feedback in B .
Also, the processor’s operation may depend on parameters in W .
Controller
The controller κ : W × B × Q × I → C × Q × O decides which operations the machine performs at each step.
It receives external inputs from I and returns external outputs in O .
It can also receive feedback from the processor and command it to do certain operations (e.g., memory read).
The decisions the controller takes may depend on its internal state (from Q ).
The controller can also depend on parameters in W .
For instance, if the controller is a neural network, then the network’s weights will range over W .
I
O
C
B κ π
M
Q (a) NCM w, r address m r read w ea write controller x y
M t (b) NTM
READ
ADD
SUB
LT r t.
WRITE r t, r t+1, r t+1,
M t controller f t c t a t b t (c) NRAM
Loss Function
The loss function L e : Trace × E
→ R indicates how close a trace τ ∈ Trace of an execution of the machine (deﬁned below) is to a behavior from a set E .
The loss function provides a criterion for training a machine to follow a prescribed set of behaviors, and hence we impose certain differentiability conditions.
We require that the loss surface is continuous and piecewise differentiable with respect to the weights w ∈ W for all examples e and inputs x with traces τ ( w, x ) :
[Equation]
Execution The execution of the machine begins with an input sequence x = { x t } n 1 and initial values of the controller state q 0 , memory m 0 , and processor feedback b 0 .
At each time step t = 1 . . .
n , controller and processor take turns executing according to the following equations:
[Equation]
A trace
[Equation]
, m 0 , q 0 )
= { ( c t , b t , q t , y t , m t ) } n 1 records these quantities’ values at each time step.
We will occasionally write τ C , τ B , . . .
for the trace projected onto one of its components c , b , . . . .
∂ NCMs Note that the differentiability conditions that we impose on the loss do not imply that any of the NCM functions π , κ and L e are continuous or differentiable.
They indeed can be highly discontinuous as in NCMs like Weston et al.
(2014)’s memory networks with a hard attention mechanism, or as in Reed & de Freitas (2016)’s neural programmer-interpreters.
In order to ﬁx these discontinuities and recover a differentiable loss surface, these architectures train with strong supervision only: the training examples e ∈ E must provide a value for every traced quantity that comes from a discontinuous parameter.
In contrast, what we call differentiable neural computational machines ( ∂ NCM), have κ , π and L e continuous and piecewise differentiable.
In this case, the loss surface is differentiable with respect to every parameter.
Thus, there is no need to specify corresponding values in the examples, and so we can train with as much trace information as available.
NTM S AND NRAM S AS NCM S
We now show how NTMs and NRAMs can be instantiated as ∂ NCMs.
NTM as ∂ NCM An Neural Turing Machine (NTM) Graves et al.
(Figure 2b) has access to a memory M ∈ R c × n of c cells of n real numbers each.
We suppose the machine has one read head and one write head, whose addresses are, respectively, the probability vectors r, w ∈ [0 , 1] { 1 ...c } .
At every time step, the read head computes the expected value m ∈ R n of a random cell at index i ∼ r .
This value together with the current input are fed into a controller neural network, which then decides on several commands.
It decides what fraction e ∈ R n to erase and how much a ∈ R n to add to the cells underneath the write head.
The write head stores the tape expected after a random modiﬁcation at index i ∼ w .
Then the controller indicates the head movement with two probability vectors ∆ r, ∆ w ∈ [0 , 1] {− 1 , 0 , +1 } which are convolved with the respective head addresses (the actual addressing mechanism is more involved, but we omit it for brevity)
Finally, the controller produces the current output value.
In terms of NCMs, the NTM’s variables fall into the following classes:
I/O
State
Communication
[Equation]
Each of these variables change over time according to certain equations (see Appendix A for details).
The processor π and the controller κ functions for each time step satisfy:
[Equation]
The standard loss function L e for the NTM simply includes a term, such as cross-entropy or Ldistance, for the machine output at every time step.
Each of these compare the machine output to the respective values contained in the examples e ∈ E .
NRAM as ∂ NCM A Neural Random Access Machine (NRAM)
Kurach et al.
(2016) is a neural machine designed for ease of pointer (de-)referencing.
An NRAM has a variable sized memory
M ∈ R c × c whose size varies between runs.
It also has access to a register ﬁle r ∈ R n × c with a constant number n of registers.
Both the memory and the registers store probability vectors over { 1 . . .
c } .
The controller receives no external inputs, but at each time step reads the probability that a register assigns to 0 .
It also produces no external output, except a probability f ∈ [0 , 1] for termination at the current time step.
The output of the run is considered to be the ﬁnal memory state.
Unlike the NTM, computation in the NRAM is performed by a ﬁxed sequence of modules.
Each module implements a simple integer operation/memory manipulation lifted to probability vectors.
For example, addition lifts to convolution, while memory access is like that of the NTM.
At every time step the controller organizes the sequence of modules into a circuit, which is then executed.
The circuit is encoded by a pair of probability distributions per module, as shown in Figure 2c.
These distributions specify respectively which previous modules or registers will provide a given module ﬁrst/second arguments.
The distributions are stacked in the matrices a and b .
A similar matrix c is responsible for specifying what values should be written to the registers at the end of the time step.
The NCM instantiation of an NRAM is the following:
[Equation]
The equations that determine these quantities can be found in Appendix B.
The processor function π and the controller function κ expressed in terms of these quantities are:
[Equation]
The loss of the NRAM is more complex than the NTM loss: it is an expectation with respect to the probability distribution p of termination time, as determined by the termination probabilities f t (see Appendix B).
For every t = 1 . . .
k , the loss considers the negative log likelihood that the i -th memory cell at that time step equals the value e i provided in the example, independently for each i :
[Equation]
S UBTRACE S UPERVISION OF NCM S
Incorporating supervision during NCM training can be helpful with: (i) convergence : additional bias may steer the minimization of the NCM’s loss function L e , as much as possible, away from local minima that do not correspond to good solutions, (ii) interpretability : the bias can also be useful in guiding the NCM towards learning a model that is more intuitive/explainable to a user (especially if the user already has an intuition on what it is that parts of the model should do), and (iii) generalization : the bias can steer the NCM towards solutions which minimize not just the loss on example of difﬁculties it has seen, but on signiﬁcantly more difﬁcult examples.
The way we provide additional supervision to NCMs is, by encoding, for example, speciﬁc commands issued to the processor, into extra loss terms.
Let us illustrate how we can bias the learning with an
NTM.
Consider the task of copying the ﬁrst half of an input sequence { x t } 2 l 1 into the second half of the machine’s output { y t } 2 l 1 , where the last input x l from the ﬁrst half is a special value indicating that the ﬁrst half ended.
Starting with both heads at position 1 , the most direct solution is to consecutively store the input to the tape during the ﬁrst half of the execution, and then recall the stored values during the second half.
In such a solution, we expect the head positions to be:
[Equation]
one-hot ( t ) if t = 1 . . .
l one
-hot ( l
[Equation]
-hot
[Equation]
To incorporate this information into the training, we add loss terms that measure the cross-entropy ( H ) between p ( t ) and w t as well as between q ( t ) and r t .
Importantly, we need not add terms for every time-step, but instead we can consider only the corner cases where heads change direction:
[Equation]
4. G ENERIC SUBTRACE L OSS FOR NCM S
We now describe the general shape of the extra loss terms for arbitrary NCMs.
Since, typically, we can interpret only the memory and the processor in terms of well-understood operations, we will consider loss terms only for the memory state and the communication ﬂow between the controller and the processor.
We leave the controller’s hidden state unconstrained – this also permits us to use the same training procedure with different controllers.
The generic loss is expressed with four loss functions for the different components of an NCM trace:
[Equation]
For each part α ∈ { C, B, O, M } , we provide hints ( t, v, µ ) ∈
σ α that indicate a time step t at which the hint applies, an example v ∈ E α for the relevant component, and a weight w ∈ R of the hint.
The weight is included to account for hints having a different importance at different time-steps, but also to express our conﬁdence in the hint, e.g., hints coming from noisy sources would get less weight.
A subtrace σ is a collection of hints used for a particular input-output example e .
We call it a subtrace because, typically, it contains hints for a proper subset of the states traced by the NCM during execution.
The net loss for a given input-output example and subtrace equals the original loss L e added to the weighted losses for all the hints, scaled by a constant factor λ :
[Equation]
rr
READ READ
INC
ONE
ADD
WRITE rr
4. S UBTRACES FOR NTM
For NTMs, we allow hints on the output y , the addresses r and w , and the tape M .
We include extra loss terms for the memory state only (all other loss terms are zero):
[Equation]
Unlike the output and addresses, values on the tape are interpreted according to an encoding internal to the controller (which emerges only during training).
Forcing the controller to use a speciﬁc encoding for the tape, as we do with NTM output, can have a negative effect on training (in our experiments, training diverged consistently).
To remedy this, we do not apply loss to the tape directly but to a decoded version of a cell on the tape.
While a decoder might ﬁnd multiple representations and overﬁt, we found that it forced just enough consistency to improve the convergence rate.
The decoder itself is an auxiliary network φ trained together with the NTM, which takes a single cell from memory as input.
The output of the decoder is compared against the expected value which should be in that cell:
[Equation]
[Equation] For all subtraces we provide in our experiments with NTMs, the hints have the same unit weight.
4. S UBTRACES FOR NRAM
For NRAMs, we hint which connections should be present in the circuit the controller constructs at each step, including the ones for register updates.
An example circuit is shown in Figure 3.
In terms of an NCM, this amounts to providing loss for commands and no loss for anything else.
We set the loss to the negative log likelihood of the controller choosing speciﬁc connections revealed in the hint:
[Equation]
In our experiments, we observed that assigning higher weight to hints at earlier timesteps is crucial for convergence of the training process.
For a hint at time-step t , we use the weight µ = ( t + 1) − 2 .
A possible reason for why this helps is that the machine’s behavior at later time-steps is highly dependent on its behavior at the early time-steps.
Thus, the machine cannot reach a later behavior that is right before it ﬁxes its early behavior.
Unless the behavior is correct early on, the loss feedback from later time-steps will be mostly noise, masking the feedback from early time-steps.
Other Architectures The NCM can be instantiated to architectures as diverse as a common
LSTM network or End-To-End Differentiable Memory Networks.
Any programming inducing neural network with at least partially interpretable intermediate states for which the dataset contains additional hints could be considered a good candidate for application of this abstraction.
E XPERIMENTAL E VALUATION
We evaluated our NCM supervision method on the NTM and NRAM architectures.
For each of the two architectures we implemented a variety of tasks and experimented with different setups of trace supervision.
The main questions that we address are: (i) does trace supervision help convergence, interpretability, and generalization?
(ii) how much supervision is needed to train such models?
Below, we summarize our ﬁndings – further details are provided in the appendix.
[Figure]
Subtraces For NTMs We measured how often we successfully trained an NTM that achieves strong generalization.
We consider a model to generalize if relative to the training size limit n , it achieves perfect accuracy on all of tests of size ≤ 1 .
5 n
[Figure]
and perfect accuracy on 90% of the tests of size ≤ 2 n .
Figure 4 reports the average improvement compared to a baseline using only I/O examples.
We ran experiments with four different tasks and various types of hints (cf.
Appendices C,
E).
Some of the hint types are: read and write specify respective head addresses for all time steps; address combines the previous two; corner reveals the head addresses, but only when the heads change direction; value gives value for a single cell.
Except for three cases, trace supervision helped improve generalization.
Here, RepeatFlip3d is most challenging, with baseline generalizing only 5% of the time (cf.
Appendix I).
Here we have the largest improvement with extra supervision: corner type of hints achieve eight-fold increase in success rate, reaching 40%.
Another task with an even larger ratio is RepeatCopyTwice (cf. Appendix), where success increases from 15.5% to 100%.
In addition to this experiment, we performed an extensive evaluation of different setups, varying the global λ parameter of the loss Eq. 8, and providing hints for just a fraction of the examples.
The full results are in Appendix I; here we provide those for RepeatFlip3d in Table 5.
The table reveals that the efﬁcacy of our method heavily depends on these two parameters.
The best results in this case are for the read/corner type of hints 12 / 110 of the time, with λ ∈ { 0 . 1 , 1 } .
The best results for other tasks are achieved with different setups.
Generally, our conclusion is that training with traces 50% of the time usually improves performance (or does not lower it much) when compared to the best method.
This observation raises the interesting question of what the best type and amount of hints are for a given task.
Finally, we observed that in all cases where training with trace supervision converged, it successfully learned the head movements/tape values we had intended.
This show that trace supervision can bias the architecture towards more interpretable behaviors.
In those cases, the NTM learned consistently sharper head positions/tape values than the baseline, as Figure 6 shows for Flip3rd .
Effect of Subtraces For NRAMs Ease of generalization for the NRAM is a known issue, with
Neelakantan et al.
(2015) reporting that ListK for example generalizes poorly, even when trained with noise in the gradient, curriculum learning, and an entropy bonus.
We observed that when run on an indeﬁnite number of examples with the correct number of timesteps and a correct module
[Figure] [Figure]
sequence, Swap and Increment would in fact occasionally generalize perfectly, but did not have the resources to run such indeﬁnite tests with Permute , ListK , and Merge .
Figure 7a demonstrates that when training had ﬁnished, either because it had ended early or had reached 5000 training examples (our upper bound), generalization would in fact be on average signiﬁcantly better than the baseline the more hints that were used for all tasks.
Here, number of hints used seemed to be a sufﬁcient predictor for the quality of the trained model.
Robustness to Noise The effect of increasing supervision on the quality of the trained model was so strong that not even noise in the input was able to signiﬁcantly hinder generalization.
In
Figure 7b, we corrupted a single character in the output examples for the Permute problem in 10% of the examples.
We found that without any extra hints, no convergence was seen after training was complete, whereas with just corner subtraces, the generalization was nearly optimal.
Furthermore, we found that noise in the trace does not seriously harm performance.
We corrupted a single hint for 20% of the traces of the Increment task using otherwise full supervision, as can be seen in the NoisyFull line of Figure 14.
C ONCLUSION
We presented a method for incorporating (any amount of) additional supervision into the training of neural abstract machines.
The basic idea was to provide this supervision (called partial trace information) over the interpretable components of the machine and to thus more effectively guide the learning towards the desired solution.
We introduced the ∂ NCM architecture in order to precisely capture the neural abstract machines to which our method applies.
We showed how to formulate partial trace information as abstract loss functions, how to instantiate common neural architectures such as NTMs and NRAMs as ∂ NCMs and concretize the ∂ NCM loss functions.
Our experimental results indicate that partial trace information is effective in biasing the learning of both NTM’s and
NRAM’s towards better converge, generalization and interpretability of the resulting models.
A NTM E QUATIONS
The controller for the NTM consists of the networks ϕ , ψ y , ψ e , ψ a , χ r , χ w , which operate on the variables: x – in
[Equation]
y –
[Equation]
value
[Equation]
write address
[Equation]
w
[Equation] [Equation]
The equations that describe NTM executions are:
[Equation]
B NRAM E QUATIONS
The controller of the NRAM consists of the networks ϕ , ψ a , ψ b , ψ c , ψ f , which operate on the variables:
[Equation]
f –
[Equation]
probability.
The equations that describe the NRAM execution are:
[Equation]
C S ETUP FOR NTM
For all of our NTM experiments we use a densely connected feed-forward controller.
There are two architectural differences from the original NTM Graves et al.
(2014) that helped our baseline performance: (1) the feed-forward controller, the erase and the add gates use tanh activation; (2) the output layer uses softmax .
In the original architecture these are all logistic sigmoids.
For the newly introduced tape decoder (active only during training) we used two alternative implementations: a tanh - softmax network, and a single afﬁne transformation.
We tested the NTM’s learning ability on ﬁve different tasks for sequence manipulation, two of which have not been previously investigated in this domain.
These tasks can be found in Appendix E.
We performed experiments using several combination of losses as summarized in
Appendix F.
The observed training performance per task is shown in Appendix I, with rows corresponding to the different loss setups.
The corner setup differs from the address setup in that the example subtraces were deﬁned only for a few important corner cases.
For example in RepeatCopyTwice , the write head was provided once at the beginning of the input sequence, and once at the end.
Similarly, the read head was revealed at the beginning and at the end of every output repetition.
In all other setups we provide full subtraces (deﬁned for all time steps).
The supervision amount can be tuned by adjusting the λ weight from Equation 8.
Further, we can also control the fraction of examples which get extra subtrace supervision (the density row in Figure I).
The performance metric we use is the percentage of runs that do generalize after 100k iterations for the given task and supervision type.
By generalize we mean that the NTM has perfect accuracy on all testing examples up to 1 .
5 × the size of the max training length, and also perfect accuracy on 90% of the testing examples up to 2 × the maximum training length.
We used a feed-forwad controller with 2 × 50 units, except for RepeatCopyTwice , which uses 2 × 100 units.
For training we used the Adam optimizer Kingma & Ba (2014), a learning rate of 10 − 3 for all tasks except RepeatFlip3d and Flip3rd which use 5 · 10 − 4 .
The lengths of the training sequences for the ﬁrst four tasks are from 1 to 5 , whereas the generalization of the model was tested with sequences of lengths up to 20 .
For Flip3rd and RepeatFlip3d , the training sequence length was up to 16 , whereas the testing sequences have maximum length of 32 .
D S ETUP FOR NRAM
Like in the NTM, we use a densely connected two layer feed forward controller for our experiments, and use ReLU as the activation function.
We make no modiﬁcations to the original architecture, and use noise with the parameter η = 0 .
3 as suggested by Neelakantan et al.
(2015), and curriculum learning as described by Zaremba & Sutskever (2014).
We stop training once we get to a difﬁculty speciﬁed by the task, and increase the difﬁculty once 0 errors were found on a new testing batch ofsamples.
Each training iteration trains with 50 examples of the currently randomly sampled difﬁculty.
Regardless of whether the model had converged, training is stopped after 5000 samples were used.
Such a low number is used to replicate the potential conditions under which such a model might be used.
As with the NTM, the Adam optimizer was used.
The speciﬁc tasks we use are described in
Appendix G, and the speciﬁc kinds of supervision we give are described in Appendix H.
The λ we used here was 40.
The system was implemented using PyTorch.
E NTM T ASKS
Every input sequence ends with a special delimiter x E not occurring elsewhere in the sequence
Copy – The input consists of generic elements, x 1 . . .
x n x E .
The desired output is x 1 . . .
x n x E .
RepeatCopyTwice – The input is again a sequence of generic elements, x 1 . . .
x n x E .
The desired output is the input copied twice x 1 . . .
x n x 1 . . .
x n x E .
Placing the delimiter only at the end of the output ensures that the machine learns to keep track of the number of copies.
Otherwise, it could simply learn to cycle through the tape reproducing the given input indeﬁnitely.
We kept the number of repetitions ﬁxed in order to increase baseline task performance for the beneﬁt of comparison.
DyckWords – The input is a sequence of open and closed parentheses, x 1 . . .
x n x E .
The desired output is a sequence of bits y 1 . . .
y n x E such that y i
=
1 iff the preﬁx x 1 . . .
x i is a balanced string of parentheses (a Dyck word).
Both positive and negative examples were given.
Flip3rd – The input is a sequence of bits, x 1 x 2 x 3 . . .
x n x E .
The desired output is the same sequence of bits but with the 3rd bit ﬂipped: x 1 x 2 ¯ x 3 . . .
x n x E .
Such a task with a speciﬁc index to be updated (e.g., 3rd) still requires handling data dependence on the contents of the index (unlike say the Copy task).
RepeatFlip3d – The input is a sequence of bits, x 1 x 2 x 3 x 4 x 5 x 5 . . .
x E .
The desired output is the same sequence of bits but with every 3rd bit ﬂipped: x 1 x 2 ¯ x 3 x 4 x 5 ¯ x 6 . . .
x E .
F
[Equation]
addr+val value address/corner write read value traces
[Table]
read – provides a hint for the address of the read head at every timestep.
write – provides a hint for the address of the write head at every timestep.
address – provides hints for the address of both the read and the write head at every timestep.
addr+val – provides value, read and write hints for every timestep.
corner – provides hints for the address of both the read and the write head at every “important” timestep - we decided what important means here depends on which task we are referring to.
In general, we consider the ﬁrst and last timesteps important, and also any timestep where a head should change direction.
For example, in RepeatCopyTwice for an example of size n with e repeats, we’d provide the heads at timesteps 0 , n, 2 n, 3 n . . .
, en .
G NRAM T ASKS
Below we describe all the tasks we experimented with.
We predominantly picked tasks that the
NRAM is known to have trouble generalizing on.
We did not introduce any new tasks, and more detailed descriptions of these tasks can be found in Kurach et al.
Swap – Provided two numbers, a and b and an array p , swap p [ a ] and p [ b ] .
All elements but that in the last memory cell are not zero.
Increment –
Given
an array p , return the array with one added to each element.
All elements but that in the last cell for the input are not zero.
Elements can be zero in the output
Permute – Given two arrays p and q return a new array s
such that s
[ i ]
=
q [ p [ i ]] .
The
arrays p and
q
are preceded by
a
pointer, a
, to array
q .
The output is expected to be
a
, s
s [ n ] q [0] , q [ n ] .
ListK – Given a linked list in array form, and an index k return the value at node k .
Merge – given arrays p and q
, and
three pointers a b
, c to array p , q , and the
output
sequence (given as zeros initially),
place
the
sorted
combination of
p and q into the output location.
The following table describes the speciﬁc NRAM instantiation used for each task.
The default sequence (def) is the one described by Kurach et al.
The number of timesteps is usually dependent on the length of the problem instance, M (equivalently the word size or difﬁculty), and in the case of ListK was given with respect to the argument k .
The difﬁculty (D) was simply the length of the sequence used.
Task
No.
Regs
Module Sequence
Timesteps
Learn Rate
Start D
End D
Swap def 0.
Increment def + R
M +0.
Permute
R + def + R + W
M +0.
[Table]
ListK
[Table]
H NRAM S UBTRACES
For each of the tasks listed Appendix G, we hand coded a complete circuit for every module and every timestep we would provide.
The following subtraces types describe how we provide hints based on this circuit.
None – provides no hints.
Full – provides the entire circuit.
SingleHint – provides a random hint at a random timestep.
SingleTimestep – provides the entire circuit at a random timestep.
Corners – provides the entire circuit at the ﬁrst and last timesteps.
Registers – provides hints for the registers at every timestep.
Modules
– provides hints for the modules at every timestep
I NTM R ESULTS
[Figure]
λ
[Figure]
baseline address
[Figure]
read
[Figure]
corner
[Figure]
λ
[Figure]
baseline addr+val
[Figure]
address
[Figure]
value
[Figure]
read
[Figure]
write
[Figure]
density
[Figure]
write
[Figure] [Figure]
Which Details to Reveal for NTM?
The ﬁrst dimension listed in the rows of the tables of Figure I controls the execution details revealed in a Subtrace.
We use subtraces showing either the addresses without the tape values , only the read heads or the write heads , or even weaker supervision in a few corner cases.
In tasks Copy Figure 9a), RepeatCopyTwice (Figure 9b) and DyckWords (Figure 9c), it is frequently the case that when the NTM generalizes without supervision, it converges to an algorithm which we are able to interpret.
For them, we designed the addr+val traces to match this algorithm, and saw increases in generalization frequency of at least 45%.
It can be concluded that when additionally provided supervision reﬂects the interpretable “natural” behavior of the NTM, the learning becomes signiﬁcantly more robust to changes in initial weights.
Additionally, for tasks
Flip3rd (Figure 9d) and RepeatFlip3d (Figure 9e), both the baseline and other supervision types are outperformed by training with read supervision.
It is also notable that corner supervision in
RepeatFlip3d achieves highest improvement over the baseline, 60% over 5%.
In essence, this means that providing only a small part of the trace can diminish the occurrence of local minima in the loss function.
How Often to Reveal for NTM?
The second dimension controls the proportion of examples that receive extra subtrace supervision (the density columns in Figure I).
For Flip3rd ,
RepeatCopyTwice and DyckWords we observed that having only a small number of examples with extra supervision leads to models which are more robust to initial weight changes than the baseline, although not necessarily always as robust as providing supervision all the time.
A couple of interesting cases stand out.
For Flip3rd with 10% corner subtraces and λ = 1 , we ﬁnd a surprisingly high rate of generalization.
Providing address traces 10% of the time when training
RepeatCopyTwice leads to better performance all the time.
For RepeatFlip3d , write traces at 1% frequency and λ = 0 .
1 generalize 30% of the time vs. 5% for baseline.
While the type of trace which works best varies per task, for each task there exists a trace which can be provided only 1% of the time and still greatly improve the performance over the baseline.
This suggests that a small amount of extra supervision can improve performance signiﬁcantly, but the kind of supervision may differ.
It is an interesting research question to ﬁnd out how the task at hand relates to the optimal kind of supervision.
[Figure]
Subtrace Type \Task
Permute
Swap
Increment
ListK
Merge
PermuteNoise
None
SingleHint
Corners
SingleTimestep
Registers    - -
Modules    - -
Full
NoisyFull - - - - -
Subtrace Type \Task
Permute
Swap
Increment
ListK
None 6290.
SingleHint 5565.
Corners 4468.
SingleTimestep 6259.
Registers 6618.
Modules 6523.
Full 4919.
Subtrace Type \Task
ListK
Swap
Permute
Increment
Merge
PermuteNoise
None 95.
SingleHint 93.
Corners 94.
SingleTimestep 36.
Full 12.
Registers 93. Modules 93.
K P ROGRAMMING NRAM S
The NRAM is parametrized by one or more straight-line partial programs, i.e., programs with no branching and no loops, chosen by register states.
The machine runs in a loop, repeatedly selecting the program for that register state then executing it.
The programs are expressend in a simple single-assignment imperative language.
Each program statement i invokes one of the modules of the architecture and assigns the result of the invocation to a local variable x i .
That variable cannot be changed later.
The ﬁnal program statement is a parallel-asignment that modiﬁes the machine registers r 1 . . .
r k .
The values that appear in assignments/invocations can be: variables in scope, machine registers, or holes ? .
These values are not used directly during execution: the actual values needs to be supplied by the NRAM controller.
The values are only used as hints for the controller during training, with the whole ?
denoting no hint.
We can describe the language in an EBNF-style grammar:
F n ::= modules of arity n
[Equation]
An example program for the Increment task would be the following:
[Equation]
x
[Equation]
← READ ( r 1 )
[Equation]
x 3 ← ADD ( x 2 , x 1 )
[Equation]
x 4 ← WRITE ( r 1 , x 3 ); x 5 ← ADD ( r 1 , x 1 ); r 1 ← x
Here, the controller is encouraged to read the memory at the location stored in the ﬁrst register r 1 , add one to it, then store it back, and then increment the the ﬁrst register.
An alternative to the trace-based approach is to make the controller produce values only for the holes, and use directly the speciﬁed variable/register arguments.
This way, only the unspeciﬁed parts of the program are learned.
This is, for example, the approach taken by ∂ Forth Bošnjak et al.
There, programs are expressed in a suitably adapted variant of the Forth programming language, which is as expressive as the language discussed above, but less syntactically constrained.
The drawback of this alternative is that whenever an argument other than a whole is speciﬁed, one must also specify the time steps to which it applies in all possible executions and not just the training ones.
That is why, typically, these values are speciﬁed either for all or for none of the time steps.
In the following examples, we will describe the register states using “0”, “!0” and “-” meaning respectively that a register has 0, that it contains anything but zero, or that it can contain anything.
L NRAM P ERMUTATION P ROGRAM
For any register pattern.
x 1 ← READ ( r 0 ); x
; x 3 ←
; x
; x
; x
M NRAM L IST K P ROGRAM
When the registers are [0, !0, !0, -, -]: x 1 ← READ ( r 0 ); x 2 ← ADD ( x 1 , 2); x 3 ← WRITE (0 , x 1 ); r 0 ← 1; r 1 ← 1; r 2 ← 1; r 3 ← x 2 ;
When the registers are [!0, !0, !0, -, -]: x 1 ← READ ( r 1 ); x 2 ← ADD ( x 1 , 2); x 3 ← WRITE ( r 1 , x 1 ); r 0 ← 1; r 1 ← 0; r 2 ← 1; r 3 ← r 3 ; r 4 ← x 2 ;
When the registers are [!0, 0, !0, -, -]: x 1 ← READ ( r 3 ); x 2 ← WRITE ( r 3 , x 1 ); r 0 ← 1; r 1 ← 0; r 2 ← 0; r 3 ← x 1 ; r 4 ← 4;
When the registers are [!0, 0, 0, -, -]: x 1 ← READ ( r 4 ); x 2 ← WRITE ( r 4 , r 3 ); r 0 ← 1; r 1 ← 1; r 2 ← 0; r 3 ← x 1 ;
When the registers are [0, !0, 0, -, -]: x 1 ← READ ( r 2 ); x 2 ← ADD ( x 1 , 2); x 3 ← WRITE ( x 2 ); r 0 ← 0; r 1 ← 1; r 2 ← 0;
N NRAM L IST K P ROGRAM
Timestep 0: x 1 ← READ ( r 0 );
x 2 ← INC ( x 1 ); x 3 ← 0; x 4 ← WRITE ( x 3 , x 1 ); r 0 ← r 1 ; r 1 ← r 1 ; r 2 ← r 2 ; r 3 ← x 2 ; r 4 ← r 4 ; r 5 ← r 5 ;
Timestep 1: x 1 ← READ ( r 1 ); x 2 ← WRITE ( r 1 , x 1 ); x 3 ← 0; r 0 ← r 0 ; r 1 ← x 1 ; r 2 ← r 2 ; r 3 ← r 3 ; r 4 ← x 3 ; r 5 ← r 5 ;
Timestep 2: x 1 ← READ ( r 2 ); x 2 ← WRITE ( r 2 , x 1 ); x 3 ← 0; r 0 ← r 0 ; r 1 ← r 1 ; r 2 ← x 1 ; r 3 ← r 3 ; r 4 ← x 3 ; r 5 ← x 3 ;
Timestep 3 to 3 + k - 1: x 1 ← READ ( r 0 );
x 2 ← INC ( x 1 ); x 3 ← 0; x 4 ← DEC ( x 1 ); x 5 ← WRITE ( r 0 , x 1 ); r 0 ← x 1 ; r 1 ← x 4 ; r 2 ← r 2 ; r 3 ← x 2 ; r 4 ← x 3 ; r 5 ← x 3 ;
Timestep 3 + k: x 1 ← READ ( r 3 ); x 2 ← WRITE ( r 2 , x 1 ); x 3 ← 0; x 4 ← 1; r 0 ← r 0 ; r 1 ← r 1 ; r 2 ← r 2 ; r 3 ← x 1 ; r 4 ← x 4 ; r 5 ← x 3 ;
Rest: x 1 ← WRITE ( r 2 , r 3 ); x 2 ← 1; x 3 ← 0; r 0 ← r 0 ; r 1 ← r 1 ; r 2 ← r 2 ; r 3 ← r 3 ; r 4 ← x 2 ; r 5 ← x 3 ;