---
license: mit
language:
  - en
datasets:
  - wikimedia/wikipedia
  - bctnry/CosmopediaV2
  - nvidia/Nemotron-Pretraining-Specialized-v1
  - nvidia/Nemotron-Pretraining-Specialized-v1.2
  - nvidia/Nemotron-SFT-Science-v2
  - nvidia/Nemotron-Math-v2
  - OctoThinker/MegaMath-Web-Pro-Max
pipeline_tag: text-generation
---

# Walker

Walker is a decoder-only language model that adds **trajectory modeling** to the Transformer.

> **Attention looks at other tokens. Walker looks at how representations move.**

## Results

### Evaluation

![MMLU](MMLU.png)

### Validation cross-entropy

![Validation cross-entropy so far](<Val CE so far.png>)

### Training curve

![Training plot](<Training Plot.png>)

## The idea

Attention asks *"what should I look at?"* by comparing queries with keys:

$$
\text{Attention}(Q,K,V)=\text{softmax}\!\left(\frac{QK^\top}{\sqrt{d}}\right)V
$$

Walker keeps attention and asks a second question: *"where is my representation going?"*

Treat the hidden states x₁, x₂, x₃, … as a path. Walker measures how that path moves:

$$
d_t = x_t - x_{t-1}
$$

This is a discrete derivative of the hidden state. Attention averages over content. The difference signal captures direction and change, which averaging does not.

Attention can **look around the path**. Walker also **walks along it**.

![The trajectory signals: velocity, multi-step displacement and the turning angle](walker_trajectory.png)

**Figure 2.** The hidden states form a path. Walker reads its velocity d⁽¹⁾, its displacement over longer strides d⁽²⁾ and d⁽⁴⁾, and how sharply it bends (the angle θ).

## Architecture

![Walker architecture: (a) the full model, (b) one Walker block](walker_architecture.png)

**Figure 1.** (a) Walker is a stack of 16 blocks between a token embedding and a language-model head. (b) Inside a block, the Walk pathway (green) runs in every block and attention (blue) runs in every fourth block, both added to the residual stream. A feed-forward sub-layer follows.

### What the Walk pathway computes

| Signal | Definition | What it captures |
|---|---|---|
| Velocity (stride 1) | d⁽¹⁾ₜ = xₜ − xₜ₋₁ | Where the state just moved |
| Velocity (stride 2, 4) | d⁽²⁾ₜ = xₜ − xₜ₋₂,  d⁽⁴⁾ₜ = xₜ − xₜ₋₄ | Longer-range motion |
| Turn | τₜ = d⁽¹⁾ₜ − d⁽¹⁾ₜ₋₁ | Change of direction |
| Curvature | κₜ = 1 − cos(d⁽¹⁾ₜ, d⁽¹⁾ₜ₋₁) | How sharply the path bends; gates the update |
| Global dynamics | Path-level reductions (see Figure 1) | Whole-path structure that attention does not compute |

The local update is

$$
\Delta x_t=\Big[\sum_{k\in\{1,2,4\}}\tanh(w_k)\odot d^{(k)}_t+\tanh(w_\tau)\odot\tau_t\Big]\odot\big(1+\tanh(\gamma)\,\kappa_t\big),
\qquad x_t \leftarrow x_t+\Delta x_t
$$

Two design decisions:

- **Exact no-op at initialisation.** With every Walk gain and the curvature-gate parameter set to zero, the update Δxₜ is exactly zero. The model starts as an ordinary Transformer and learns how much trajectory to use.
- **Higher learning rate for the Walk parameters** (three times the backbone's).

Walker adds trajectory features alongside attention. It does not replace it.

## Model details

| | |
|---|---|
| Parameters | **~0.95 B** |
| Layers | 16 |
| Hidden size `D` | 2048 |
| Attention heads | 16 |
| Vocabulary (GPT-2 BPE) | 50257 |
| Context length `max_t` | 16384 |
| Precision | bfloat16 |
| Architecture | custom `WalkGPT` (attention + trajectory). **Not** a `transformers` `AutoModel` |

## Training

Trained on a **~15 B-token** mixture of open corpora (see `datasets` above). That is roughly **1000× less data** than SmolLM2/OLMo-class models and is the main limit on factual precision.

## Checkpoints

| File | Notes |
|---|---|
| `checkpoints/ckpt_ongoing_step29750.pt` | Training checkpoint at **step 29750** (model + optimizer state) |

## Usage

> ⚠️ Walker is a custom architecture, so `transformers` `AutoModel` / `pipeline` **cannot load it**. Use the trainer's own inference loader, which rebuilds the network from the saved config.

The quickest route is `try_walker.py`, placed next to `train_walker_1b_continue.py`:

```bash
python try_walker.py --ckpt checkpoints/ckpt_ongoing.pt --prompt "The migration of the Arctic tern"
```

Or from Python:

```python
import importlib.util

spec = importlib.util.spec_from_file_location("walker", "train_walker_1b_continue.py")
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)

model, tokenizer, cfg = mod.build_model_for_inference(
    "checkpoints/ckpt_ongoing.pt", device="cuda")
model.eval()

ids = tokenizer("The migration of the Arctic tern", return_tensors="pt")["input_ids"].cuda()
out = model.generate(ids, max_new=64, temperature=0.8, top_p=0.9)
print(tokenizer.decode(out[0]))
```

A companion **Gradio Space / MoLab portal** visualises the hidden-state trajectory in 3-D while tokens are generated; see the Space linked from this repository.

## Design history

Walker is the end of a longer search: twelve designs, most of them falsified. The full story, with a verdict for each design, is in **[HISTORY.md](HISTORY.md)**.

<details>
<summary><b>Show the lineage</b></summary>

<br>

![Cladogram of the designs that led to Walker](walker_history_cladogram.png)

Stored objects, richer relations and learned interaction operators did not beat attention. What survived was modelling how the representation moves, which became the Walk. See [HISTORY.md](HISTORY.md) for the evidence behind each branch.

</details>

## Limitations

- **Base model.** It continues text and does **not** answer questions. Instruction following needs SFT.
- **Factual precision** is bounded by the ~15 B-token budget.
- **Math** is unreliable for base language models; route arithmetic to a tool.

## Summary

$$
\boxed{\text{Walker}=\text{Attention}+\text{Trajectory Modeling}}
$$

**Attention:** *what information is relevant?*  **Walker:** *how is the representation moving?*