tem_tf2 · field guide
What this model assumes, what happens on every one of its two million training steps, and how to read every number it prints at you.
Imagine you walk into a room you have never seen before. You take four steps: north, east, south, west. You are back where you started, and you knew that before you got there. Nobody had to teach you this room. You already knew the rule, because the rule belongs to space itself, not to any particular room.
That is the whole idea. TEM separates two kinds of knowledge that are usually tangled together:
The claim in the paper is that the brain keeps these in different places. Grid cells in medial entorhinal cortex carry the reusable structure. Place cells in hippocampus bind that structure to whatever you are currently seeing. Because TEM stores structure separately, a new environment can reuse it immediately. That is what lets you predict correctly in a room you have never visited.
The model has three representations. Almost everything else in the codebase exists to keep them honest. Every diagram below uses this colour coding.
TEM never sees pixels or coordinates. Its world is a graph: a set of nodes connected by labelled edges. At each node sits one sensory observation drawn from 45 possibilities. The agent takes a random walk, and at every step it must answer one question: what am I about to see next?
That question is the entire training signal. There are no labels, no rewards driving the loss, no supervision beyond "predict the next observation". Everything the model learns, grid cells included, falls out of getting better at that one task.
parameters.py is not a plain grid but
frank2000, a W‑maze alternation task with reward sites. The shipped settings give
it 52 nodes.The name comes from Frank, Brown and Wilson (2000), a rat experiment on continuous spatial alternation. The animal runs a W-shaped track with three arms. From the centre arm it must go to one side arm, return, then go to the other side arm. Reward arrives only when the alternation is correct, so the right answer at the choice point depends on where you came from, not on where you are. A purely spatial map cannot solve it.
Frank2000 in environments.py builds this as a graph rather than a track.
With the shipped settings (height=5, width=1) it has
4*(h-1+h+w) + 2*(w+h) = 52 nodes: a 40-node main circuit that runs both correct
alternation loops, plus two 6-node error branches for the wrong turns. Four nodes are rewarded
(0, 10, 20, 30) and two are the unrewarded ends of the error branches (45, 51).
reward_val, which each environment sets
to n_states/2, so 26 here. Those nodes dominate the loss on purpose.Drawn out, that graph is a single directed 40‑node cycle with two 6‑node cul‑de‑sacs hanging off it. Nothing in the code knows about arms or corridors: the W‑shape only shows up in which physical stretch each run of node indices happens to correspond to, and in the fact that the centre arm is traversed four times per full circuit while each side arm is traversed twice.
Frank2000.world() actually builds. adj is a
52×52 matrix: a chain i → i+1 for the first 40 nodes with 39 → 0
closing the loop, two chains for the error branches, and six hand‑written edges joining them.
tran is adj row‑normalised, with the branch entries weighted
p/(1-p) before normalisation so each choice point leaks 10% of its probability onto the
wrong turn. Node 0 is always the start state.Both words invite the wrong picture. The walk is not produced by the model: env.walk()
is called once when the environment is created, and it generates the whole trajectory up front by
sampling nothing but the environment's own tran matrix. The model is then told
which action was taken, as the input d. It never chooses one. There is no policy, no
value function and no reinforcement learning anywhere in the code. Wrong turns onto the error
branches are the environment's dice roll at error_prob, not the model's mistakes.
Reward is a sensory class like the other forty‑four, drawn at random for each environment.
The agent sees it; it does not collect it. What makes those six nodes special is a single line in
compute_losses, where cross‑entropy is multiplied by reward_val. That is a
weight on the prediction objective, not a reinforcement.
So the alternation task here is not the behavioural one. It is not “choose the correct arm” but “on arriving at the top of the centre arm, predict whether you are about to see the reward sense or the no‑reward one”. Those two nodes are the same physical place, so the only way to call it is to know which loop you are on, which is exactly the history that a place‑only code throws away.
In most networks a batch is a stack of independent examples. Here it is a stack of 16 independent
worlds, run in parallel for the whole of training. params.n_envs = params.batch_size
= 16, and the two numbers are the same thing by construction.
Each of the 16 gets its own graph, its own random assignment of 45 sensory classes to nodes, its own walk, and its own Hebbian memory matrix. Nothing is shared between them except the weights. That is what forces the weights to hold what is common to all worlds: anything specific to one of the 16 belongs in that world's memory matrix, not in a weight.
The walks are not synchronised (asyncrounous_envs = True). Each environment draws its
own walk length, and when one ends, a fresh environment takes that slot in the middle of a batch
while the other 15 carry on. The model therefore always sees a mix of familiar and brand-new worlds.
So the model never learns one map. It learns what all maps have in common.
Three different things in this code are called a step, and the log prints all three. Getting them mixed up makes every runtime estimate wrong by a factor of 16 or 75.
| Name | In code | What it is | Size |
|---|---|---|---|
| Environment step | one iteration of the seq_len loop | One move along one edge, in one environment. The six stages in section 04 happen once. | 1 |
| Sequence | seq_len | The truncation window for backpropagation through time. Gradients flow across these 75 steps and stop at the boundary. The grid code and both memory matrices carry over to the next sequence, detached from the graph. | 75 steps |
| Gradient update | train_i | One call to train_step: run the sequence forward for all 16 environments,
sum the losses, take one Adam step. This is the unit every schedule in section 07 counts in. |
16 × 75 = 1,200 steps |
total_steps in the log | train_i * seq_len | Steps taken by one environment. It ignores the batch, so the model has really seen 16 times this many transitions. | 75 per update |
So the configured maximum of 2,000,000 gradient updates is 2.4 billion environment steps. Nobody runs that. The paper's figure is 50,000 updates, and useful structure appears well before then.
Adam, with the learning rate set fresh on every iteration from the schedule
(9.4e-4 - 8e-5) * 0.5^(train_i/4000) + 8e-5. That halves the rate every 4,000 updates
and floors it at 8e-5. Every gradient is clipped to norm 2 before it is applied, which
is what keeps the Hebbian matrix from blowing the whole run up early on.
An environment's walk is n_restart laps of its graph, where n_restart
starts at restart_max = 40 and decays toward restart_min = 10 as training
proceeds, plus a jitter of up to 30 laps. Early on the agent gets roughly 2,000 steps to learn a
52-node world. Later it gets a quarter of that. The task quietly gets harder.
This loop is the model. It runs 75 times per gradient update, and it is worth understanding completely, because every metric you read later is measured at one of these stages.
Path integration is dead reckoning: update your estimate of where you are from the movement you just made, with no reference to what you can see. TEM does it with a matrix built from the action.
# tem_model.py, g2g() and get_g2g_update()
T = reshape(t_vec(d), [16, 120, 120]) * mask_g # action d becomes a transition matrix
g_gen = clip(g_prev + T @ g_prev, -1, 1) # the whole of path integrationd is the action just taken, one of up, down,
left, right in this world. A small network turns it into a 120 × 120
matrix, so the model learns a separate linear operator per action rather than being handed one.
mask_g then zeroes most of that matrix: modules may only be driven by modules at their
own timescale or slower, which is the hierarchy of section 06.
Two properties matter. The update is a function of the action alone, so g_gen cannot
smuggle in the current observation. And it is applied over and over, so error accumulates: after
enough steps the estimate has drifted, exactly as an animal's does.
Correcting that drift is the job of p2g, and it works by pattern completion. Take
what you can see right now, use it to query memory, and see which structural slot the memory that
comes back was stored in.
x2p.p_x.p_x over the 10 sensory components of each phase strips the sensory part
back out, leaving a structural estimate. A per-module layer maps it to mu_p2g.sigma.combine2 fuses the two estimates by precision, the standard Gaussian product:
mu = sigma² (mu_path/sigma_path² + mu_mem/sigma_mem²). Whichever source
is more confident wins.That fusion is the whole of drift correction. Path integration proposes, memory corrects, and the
weighting is learned rather than fixed. The result is g, as opposed to
g_gen.
Memory is nonsense on update 1, so trusting it would poison g. The code adds
(1 - p2g_use) * 10000 to sigma_p2g, an uncertainty so large that
combine2 ignores memory entirely. As p2g_use rises from 0 to 1 over the
first few hundred updates, memory is let in. This is the single most important thing the
curriculum does.
The same underlying question, "what will I see next", gets asked three times per step from three different states. The shorthand for each:
| Phrase | Prediction | Knows | Does not know |
|---|---|---|---|
| Honest prediction, or real forecast | x_gt = f_x(gen_p(g_gen)) | Every action taken so far, and everything in memory | The observation it is being scored against. It is a forecast in the ordinary sense: made before the fact, from a state that cannot contain the answer. |
| Corrected prediction | x_g = f_x(gen_p(g)) | The same, plus the drift correction that used the current observation to fix
g |
The answer directly, though the correction has touched it. This is why aGinf sits above aGt. |
| Reading the answer off the page | x_p = f_x(p) | Everything, including the current observation, which went into p through
x2p |
Nothing. Predicting x from a code built out of x is an autoencoder round trip, so a high score here proves that the encoder and decoder are wired up, and nothing about structure. |
All three are trained, and all three should be. The autoencoder path has to work before anything else can. The mistake is quoting aP as the model's performance.
TEM does not keep a list of past events. It keeps one matrix, a_rnn in the code, and
adds every new memory into it by Hebbian learning. Retrieval is not a lookup. It is an attractor
process that settles onto the stored pattern closest to what you asked for.
a_rnn and a_rnn_inv, one
for each direction of retrieval.The textbook Hebbian write is M += p pᵀ. TEM does not do that. It writes the
difference between what it inferred and what it predicted, which stops it re-storing things it could
already retrieve.
# tem_model.py, hebbian() collects these each timestep
a = p - p_g # inferred place code minus the one generated from memory
b = p + p_g
# final_hebbian(), once per 75-step sequence
M ← clip( M * (forget * λ)^75 + b @ (a * forget_vec)ᵀ * mask_p, -1, +1 )Read the update term as a product: b aᵀ = p pᵀ - p_g p_gᵀ + p_g pᵀ -
p p_gᵀ. The first part stores the new pattern. The rest subtracts off what memory
already generated. If the model predicted the place code perfectly then p = p_g,
a = 0, and nothing is written at all. This is the fast-weights rule from Ba et al.
(2016), and it is why memory does not saturate on repeated visits to the same node.
| Term | Value | Job |
|---|---|---|
| λ | 0.9999 | Passive decay per step. Old memories fade whatever else happens. |
| forget | 0 → 1 over 200 updates | Curriculum multiplier on that decay. At update 1 it is 0.005, so
(forget * λ)^75 is effectively zero and each sequence starts from a blank
matrix. |
| forget_vec | (forget λ)^(74...0) | Within one sequence, memories laid down earlier are already more faded than the ones at the end. |
| η | 0.5 | Write strength. Enters as sqrt(η * h_l) on each stored vector, so the
product of two of them carries the full η * h_l. |
| h_l | 0 → 1 over 16,000 updates | The curriculum on write strength. The slowest ramp in the model. |
| mask_p | 0/1 matrix | Blocks connections between place modules that the hierarchy does not allow. |
| hebb_lim | ±1 | Hard clip on every entry. MaxHebb in the log is the largest absolute value, so
it can never exceed 1. If it pins at 1, memory has saturated. |
Two matrices are built this way. a_rnn uses a, b from
p_g and retrieves a place code from a structural query. a_rnn_inv uses
p - p_x, p + p_x and retrieves in the other direction, from sensory input, which is what
p2g needs to correct drift.
Retrieval is iterative. You start from a partial or noisy pattern, multiply by the memory matrix, damp, clip, and repeat. Each pass pulls the state further toward whichever stored pattern it most resembles.
p = f_p(init)
for i in 0..4: # n_recurs = n_freq = 5
update = M @ p # as scalar products, not a real matmul
for f in modules still active at i:
p[f] = f_p(0.8 * p[f] + update[f]) # kappa = 0.8kappa = 0.8 is how much of the current state survives each pass. Below 1 the loop
contracts toward a fixed point instead of ringing. f_p is a leaky ReLU on values clipped
to [-1, 1], so activity cannot run away.
The loop is not uniform across modules. max_attractor_its = [5, 4, 3, 2, 1] stops
each module after a set number of passes, fastest module first to finish last:
max_attractor_its = [n_recurs - f], straight out of
parameters.py, and it is why retrieval costs less than five full matrix products.
The 120 g cells are not one undifferentiated pool. They are split into
five modules that are forced to update at different rates, from fast to slow. This is a structural
prior, hard-coded in parameters.py, and it is the reason grid-like cells of several
different scales appear after training.
params.freqs, and it is trainable, so the model can move these scales.A structural prior is a constraint built into the architecture rather than learned from data. It
does not say what the answer is. It restricts the shape of the answers the model can represent, so
that the ones the world actually supports are easy to reach. TEM has five, all in
parameters.py, and grid cells do not appear without them.
| Prior | Where | What it forces |
|---|---|---|
| Split into 5 modules | n_grids_all = [30,30,24,18,18] | The structural code cannot be one undifferentiated pool. |
| A timescale per module | freqs, gamma | Module f smooths its sensory stream as
xᵗ = a xᵗ₋₁ + (1-a) x with its own a. Fast
modules track single steps, slow ones average over many. Trainable, but it starts split. |
| Hierarchical transitions | mask_g, R_G_F_f | A module can be driven by its own timescale and slower, never by a faster one. Coarse structure informs fine, not the reverse. |
| Place cells are a product | n_place_all = n_phases * 10 | p is the outer product of a structural phase with a sensory component, then
flattened. It cannot represent a location without representing what is there. |
| Two-hot sensory code | two_hot = True, s_size_comp = 10 | Each of the 45 stimuli is a pair of active units out of 10, and
C(10,2) = 45 exactly. Stimuli share units, so the model cannot treat them as 45
unrelated symbols. |
Nothing here mentions hexagons, spacing, or space at all. That is the claim worth taking seriously: give a predictive model a factorised code, a hierarchy of timescales, and a fast Hebbian memory, and grid cells fall out of prediction on graphs.
You cannot train TEM by switching everything on at once. The memory would be garbage on step
one, and inference that trusted it would never recover. Instead several quantities ramp up over the
first few thousand steps. These are the scalings printed in the log,
and knowing their schedule tells you whether a run is behaving.
forget reaches full strength almost
immediately; temp (the weight on the consistency losses) takes 2,000 updates;
hebb_learn takes 16,000. One quantity runs the opposite way: p_cell_reg
decays from 1 to 0 over 4,000 updates, so the sparsity penalty on p constrains early
representations and then lets go. Its counterpart g_cell_reg is scheduled over
40,000,000 updates, which means it never meaningfully moves."Several quantities ramp up" means five scalars, recomputed from the update counter on every
iteration by get_scaling_parameters and passed into the model as
scalings. They are not learned. They are a fixed function of train_i.
| Scaling | Formula | Full by | What it gates |
|---|---|---|---|
| forget | min((i+1)/200, 1) | 200 | Multiplies λ in every memory decay term. At 0.005 on update 1 the
matrix is wiped each sequence, so the model is forced to learn path integration before it can
lean on memory at all. |
| p2g_use | sigmoid(i/200) | ~600 | Adds (1 - p2g_use) * 10000 to the uncertainty of the memory-based estimate of
g. While it is near 0, drift correction is switched off. It also scales the
lp_x loss. Note it starts at 0.5, not 0, because the sigmoid is centred on
p2g_use_it = 0. |
| temp | min((i+1)/2000, 1) | 2,000 | Weight on the consistency losses lp, lg and lp_x.
Demanding that inference and generation agree is meaningless while both are noise, so it is
phased in. |
| h_l | min((i+1)/16000, 1) | 16,000 | Write strength, as sqrt(η * h_l) on every stored vector. The slowest ramp,
and the reason accuracies before 16,000 updates do not mean much. |
| p_cell_reg | 1 - min((i+1)/4000, 1) | 0 by 4,000 | Runs the other way. The L1 sparsity penalty on p is strong early, then
released once place cells have taken shape. |
| g_cell_reg | 1 - min((i+1)/4e7, 1) | never, in practice | g_reg_it is 40,000,000, so after 50,000 updates this is still 0.9988. The L2
penalty on g is effectively constant. Treat it as a fixed regulariser rather than
a schedule. |
The order is the point, and it is the order of dependency. Forgetting comes first so nothing trusts an empty memory. Then memory is allowed into inference. Then the consistency terms start to bite. Write strength arrives last, once there is something worth writing. Reverse any two and the run fails in a way that looks like a bug in the model.
Accuracy figures before roughly 16,000 updates are not meaningful. The memory system is still being phased in. A run that looks mediocre at step 2,000 is simply a run at step 2,000.
run_tem.py appends two lines to summaries.log every 100 updates. These
are real numbers from a short run of this code:
MaxHebb=0.10516, T=0.10, F=1.00, L=0.01, train_i=200, total_steps=15000
it=200, lxP=2.25, lxG=6.02, lxGt=6.02, lp=0.07, lg=0.00, aP=89.08, aGinf=15.75, aGt=15.75A loss is one number that the optimiser tries to make smaller. It is not an error rate and it has
no upper bound. TEM's is a plain sum of eight terms, assembled in compute_losses, and
they are not interchangeable. They fall into three kinds:
lx_p, lx_g, lx_gt.
Cross-entropy between the true observation and a predicted distribution over the 45 stimuli,
measured in nats. These are the ones that mean "was the model right".lp, lg, lp_x.
Squared error between two of the model's own internal codes. Nothing external is involved. They
say "the route you inferred and the route you generated should be the same route", which is what
stops the two halves of the model drifting apart into separate solutions.lg_reg, lp_reg. Penalties on raw
activity, with no reference to the data at all. They shape what the codes look like: L2 keeps
g small, L1 makes p sparse.A useful anchor for the cross-entropy terms: predicting all 45 stimuli as equally likely gives
ln(45) = 3.81. So lxGt = 6.02 in the sample log above is worse than
guessing, which is what an untrained network that is confidently wrong looks like. Below 3.81 the
model is doing something. The consistency and regularisation terms have no such reference point, and
only their trend is meaningful.
# per timestep, then summed over the 75 steps of the sequence
norm = 1 / (n_visited_in_batch * 75)
lx_gt += sum( CE(x, x_gt) * s_visited * x_mult ) * norm
lp += sum( 0.5*||p - p_g||² * s_visited ) * temp * norm
lg_reg+= sum( ||g||² * s_visited ) * 0.01 * g_cell_reg * norm
train_loss = lx_p + lx_g + lx_gt + lp + lg + lp_x + lg_reg + lp_regThree things to read off that. Every term is masked by s_visited. Every term is
normalised by how many states in the batch counted, so the numbers stay comparable as that varies.
And the curriculum scalars multiply the losses directly, so lp and lg are
near zero early on because temp is near zero, not because the model is doing well.
x_mult is the reward weighting: at the six reward and no-reward nodes of the W-maze,
cross-entropy is multiplied by 26. A model that predicts everything except the reward nodes will show
a bad loss and a decent accuracy at the same time.
All three ask "did the model predict the right one of 45 stimuli?" They differ only in what the model was allowed to know when it guessed. Chance is 1/45 ≈ 2.2%.
| Log name | In code | Predicts from | What it tells you |
|---|---|---|---|
| aP | accuracies.p | Inferred place cells p, which have already absorbed the current observation | A sanity check, not a result. It should climb fast and go high. It mostly says the autoencoder path works. High aP with low aGt means the model is copying, not understanding. |
| aGinf | accuracies.g | Inferred g, after memory evidence has corrected path integration | How good memory retrieval is once you are confident about where you are. Sits between the other two. |
| aGt | accuracies.gt | Path-integrated g_gen only. Actions, no peeking | This is the headline number. A genuine forecast made before observing. This is where generalisation shows up, and the figure to quote. |
The expected ordering is aP ≥ aGinf ≥ aGt, because that is the order of decreasing information. Watch the gap between aP and aGt. It shows how much the model leans on seeing the answer rather than knowing the structure.
Each accuracy has a matching cross-entropy loss, plus the consistency and regularisation terms.
All are printed with the same prefixes. The which_costs list in
parameters.py decides which ones are summed into train_loss; the default
includes all eight below. weight_reg, an L2 penalty on the network weights, is
implemented but left out of that list.
| Term | Type | What it pushes toward |
|---|---|---|
| lx_p | cross-entropy | Prediction from inferred p should match the observation. Partner of aP. |
| lx_g | cross-entropy | Prediction from inferred g via memory. Partner of aGinf. |
| lx_gt | cross-entropy | Prediction from path integration alone. Partner of aGt, the one that matters. |
| lp | squared error | The place code you inferred and the one you generated from memory should agree. Keeps memory honest. |
| lg | squared error | Where you think you are and where path integration says you are should agree. This is the term that makes g a consistent map rather than drifting noise. |
| lp_x | squared error | Place code inferred from sensory input should match the full inference. Only active when
infer_g_type contains p, which it does by default. |
| lg_reg | L2 penalty | Keeps g activity small. Weight decays to zero over training. |
| lp_reg | L1 penalty | Keeps p activity sparse, with L1 rather than L2. Sparsity is what makes place cells look like place cells. |
s_visited masks the losses, so the model is scored only on states it has already
visited in the current walk. Predicting a node you have never been to is information-theoretically
impossible, so those steps drop out. Every accuracy here therefore means "accuracy where prediction
was possible at all", not accuracy over every step.
| Log | Name | Meaning |
|---|---|---|
| T | temp | Weight on the consistency losses (lp, lg). Ramps 0→1 over 2,000 updates. |
| F | forget | Memory decay strength. Reaches 1 within 200 updates. |
| L | h_l | Hebbian learning rate scaling. Ramps over 16,000 updates, the slowest schedule. |
| MaxHebb | · | Largest absolute value in the Hebbian matrix. It should grow and then settle, not explode. |
Training writes into a dated run folder outside the source directory, and copies the scripts it ran alongside the results, so you can always trace an old run back to the code that produced it.
conda activate tem
cd tem_tf2
python run_tem.py
# watch it, from the repo root
tensorboard --logdir=Summariesint(50000/75) = 666 updates. The script/ copy means a run remains
reproducible even after you edit the source.| Milestone | Updates | Notes |
|---|---|---|
| Graph compilation | before step 1 | Several minutes of apparent hang with graph_mode=True. Expected, not a crash. |
| First checkpoint | 0 | Untrained. Loading it gives noise-like rate maps. |
| Curriculum complete | ~16,000 | Only now are the accuracies worth reading. |
| Paper-scale training | up to 50,000 | The figure quoted in the TEM paper for full training. |
| Configured maximum | 2,000,000 | train_iters. Stop it long before this. |
Loss curves tell you the model trained. They do not tell you what it learned. Only the cells answer that, and the notebook exists to draw them.
A rate map is the simplest possible object: for one cell, colour every location in the environment by how active that cell was there, averaged over the walk. What pattern appears tells you what kind of cell the network invented on its own.
square_plot output. The paper's claim is that
nobody designed these cell types in. They emerge from the single objective of predicting the next
observation, given the split between structure and sense.There is no smoothing over space and no notion of a pixel. rate_map_from_timeseries
does one thing: bin cell activity by the node the agent was standing on, then average.
# plotting_functions.py, one environment
for each timestep t:
bucket[ position[t] ].append( activity[:, t] )
rate_map[node] = mean( bucket[node] ) # zeros if never visitedThe result is a table of n_states × n_cells, so 52 × 120 for the grid
cells in the default world. square_plot then places those 52 numbers into the shape of
the environment for display. The optional smoothing argument runs a Savitzky-Golay
filter along time before binning, not across space.
Two consequences worth keeping in mind. A node the walk never reached reads as exactly zero, which looks like silence rather than missing data. And each node's value is an average over every visit, including visits arriving from different directions, so a cell that fires only on one approach shows up washed out.
Cell 2 reports, per environment, the proportion of visited states predicted correctly:
0 Num_states : 52 Approx proportion : [0.051] Approx num : [3.]Read this as: of the states this agent visited, about 5.1% were predicted correctly, roughly 3 states out of 52. Compare against chance of 2.2%. The example above comes from an untrained checkpoint, which is exactly what near-chance looks like.
Loading iteration 0 and concluding the model does not work. The first checkpoint is written
before any training. If the rate maps look like static and accuracy sits near 2%, check which
index the notebook loaded before drawing any conclusion.
| What you see | What it usually means |
|---|---|
| aP high, aGt at chance | Memorising the current observation, no structure learned. Check lg. If it is not falling, nothing is constraining path integration. |
| MaxHebb growing without limit | Memory instability. Check the forget schedule and λ. |
| All accuracies flat at chance after 20k | Learning rate or curriculum problem, not a modelling one. Confirm the scalings reached 1.0. |
| Long silence before step 1 | Normal. Graph compilation with graph_mode=True. |
| Rate maps look like noise | Almost always an early checkpoint. Confirm the loaded iteration first. |
| aGinf equals aGt exactly | Common early on, while inference has nothing to add over path integration. The two should separate as memory becomes useful. |
Everything above describes how TEM works. This section describes what the implementation is for, which is a different question, and the one that matters if you are deciding whether to reuse any of it elsewhere.
tem_tf2 is the reference implementation for the Tolman‑Eichenbaum Machine, and it
exists to support one claim: that training a single generative model to predict the next sensory
observation during random walks on graphs makes grid‑like and place‑like codes appear, and makes
zero‑shot inference on unvisited edges possible. It is a scientific artifact packaged as runnable
code. It is not a library, not a product, and not an agent.
One objective, cross‑entropy on the next observation, computed three ways. The full list is
which_costs in parameters.py, and the arithmetic is in
compute_losses.
| Loss | Prediction it scores |
|---|---|
| lx_p | From the full inferred place code |
| lx_g | From path integration plus memory |
| lx_gt | From path integration alone, before anything at this node has been seen |
| lp, lp_x, lg | Agreement between the inferred and the generated codes |
| lg_reg, lp_reg | Activity regularisation, decayed away by the curriculum |
lx_gt is the scientific payload. Everything else exists to make it work.
The network sees six things, assembled by inputs_2_tf:
x one-hot observation (seq_len, batch, s_size)
d one-hot action taken (seq_len, batch, n_actions)
g grid state carried over (batch, 120)
x_ filtered sensory state (batch, 5 x 10)
hebb_mat per-world memory (batch, 400, 400)
hebb_mat_inv per-world memory (batch, 400, 400)It returns the three predictions, plus an updated g, x_s,
a_rnn and a_rnn_inv to carry into the next chunk.
The remaining fields in the input dictionary, positions, reward_pos,
reward_val and s_visited, never reach the network. They only weight and
mask the loss. The model has no notion of a node, a graph, a world or a reward. It sees an
observation and an action. Anything that can produce those two streams can drive it.
| Layer | Files | How reusable |
|---|---|---|
| Model | tem_model.py | Environment‑agnostic, 1040 lines. Reusable unchanged. |
| Environments | environments.py | 1987 lines, nine world types. Each supplies adj, tran,
states_mat, relation(), walk() and
get_node_positions(). A new world is a subclass plus one if in
initialise_environments. |
| Plumbing | data_utils.py, run_tem.py | Batch‑slot lifecycle, memory zeroing, walk chunking, checkpoints. Tied to the design where a batch slot is a world. |
| Analysis | behaviour_analyses.py, plotting_functions.py, the notebook | Rate maps, autocorrelograms, link‑inference curves. Regenerates the paper's figures, and is useful only if you are asking the paper's questions. |
These, rather than the code structure, decide whether the model ports to another domain.
g is trained so that path integrating through a
sequence of actions matches the destination reached. Your action set needs a consistent algebra. A
fully connected graph, or actions whose meaning changes with context, degrade g to
noise.s_size = 45 by default,
decoded by a softmax. Continuous or high‑dimensional input needs a different MLP_c
and a different lx.b holds a live Python object
and a memory matrix, and both are wiped when its walk ends. There are no world ids, and nothing is
persisted during training.env.walk() generates all of
it from tran.p2g_use, forget,
temp and hebb_learn all ramp over the first few hundred to few thousand
updates. Setting them to 1 from the start does not train.No policy, no value function, no reinforcement learning. No classifier and no supervised head of any kind. No action selection. No inference API for a trained model beyond running a forward pass over a fresh walk. No persistence of environments during training. Reward is a sensory class with a loss multiplier, not a signal.
| If you want | Take |
|---|---|
| The representation | tem_model.py unchanged, your own environment in place of
environments.py, and your own readout on the rate map built by
cell_norm_online or on a_rnn. The cheapest port. |
| The memory mechanism alone | final_hebbian, attractor and hebb_scal_prod are
self‑contained. Holding a and b through the chunk and folding them
in at the end is reusable on its own. |
| A sequence model for a structured domain | The whole architecture. Settle first whether your domain has composable actions. If it does not, TEM's central mechanism has nothing to do. |
| To reproduce or extend the paper | All four layers unchanged, and the budget for the full run. |