# Coherent Memory Structures in Neural Fields

This is the official repository for the 2025 Neurips submission [Coherent Memory Structures in Neural Fields](README.md).

# Install

1. Retrieve the repository

```
git clone repository.git
```

2. Install dependencies

```
pip3 install -e .
```

# Running an Experiment

For ease of use and review - our provided architecture is written in a standalone notebook `MZCopy.py`.
You can run the notebook in Google Colab or Jupyter Notebook.


For the baseline architectures, we provide a script `run_training.py` that can be run from the command line.

```
!python3 run_training.py --model {model} --dataset {task} --loss {loss} --mem_len {num_tokens} --max_steps {num_batches} --lr {lr} --n_hid {n_hid}
```
model: 'wavernn', 'mamba', 'nope', 'rope', 'alibi'

dataset: 'copy', 'selective'

loss: 'crossentropy' or 'mse'

men_len: number of tokens to memorize, default 10

max_steps: maximum number of batches, default 1000

lr: learning rate, default 1e-3

n_hid: latent size/model dim, default 100

Other useful args: batch/batch_test (batch size for train/test), grad_clip (gradient clipping value, default None), n_layers (number of layers, default 2 for mamba/transformers). 

# Citing our Work


# Acknowledgments

Some of our code was adapted from the repos for [Traveling Waves Encode the Recent Past and Enhance Sequence Learning](https://github.com/akandykeller/Wave_RNNs) and [Repeat After Me: Transformers are Better than State Space Models at Copying](https://github.com/sjelassi/transformers_ssm_copy/tree/main). We thank the authors of those works!
