Zoology: Research Playground for Language Model Architectures

Experiments are controlled via nested Pydantic configuration objects. The primary entry point is TrainConfig.

Key configuration components include:

• TrainConfig: The top-level configuration.
• DataConfig: Configures training and testing datasets.
• ModelConfig: Configures the model architecture and vocabulary.
• ModuleConfig: Configures PyTorch modules by providing a class name and kwargs.
• FunctionConfig: Configures partial functions by providing a function name and kwargs.

ModuleConfig and FunctionConfig both implement an .instantiate() method to create the actual object or function.

from zoology.config import TrainConfig, ModelConfig, DataConfig, ModuleConfig, FunctionConfig
from zoology.data.associative_recall import MQARConfig

input_seq_len = 10

config = TrainConfig(
    max_epochs=20,
    data=DataConfig(
        train_configs=[MQARConfig(num_examples=10_000, vocab_size=128, input_seq_len=input_seq_len, )],
        test_configs=[MQARConfig(num_examples=1_000, vocab_size=128, input_seq_len=input_seq_len, )],
    ),
    model=ModelConfig(
        vocab_size=128,
        sequence_mixer=ModuleConfig(name= "zoology.mixers.attention.MHA") 
    ),
)

# Example of using FunctionConfig
fn_config = FunctionConfig(name="torch.sort", kwargs={"descending": True})
fn = fn_config.instantiate()
# fn(torch.tensor([2,4,3])) -> [4, 3, 2]

To reproduce the synthetic data results shown in Figure 2 of the Zoology: Measuring and Improving Recall in Efficient Language Models paper, follow these steps:

1. Setup WandB: Ensure Weights & Biases (WandB) is configured to log all experiment results.
2. Run the experiment sweep: Execute the launch command using the zoology.launch module. Use the -p flag to launch configurations in parallel (recommended for large sweeps).
3. Scale down (Optional): The full sweep contains 448 model/data configurations and takes a long time. To run a smaller subset, modify the loops within zoology/experiments/paper/figure2.py to limit the models, sequence lengths, or learning rates.
4. Generate the plot: Once the experiments are complete, use the plotting code located in zoology/analysis/paper/figure2.py to produce the final visualization.

python -m zoology.launch zoology/experiments/paper/figure2.py -p

You can launch experiments using the zoology.launch module.

1. Single Experiment: Create a python file that defines a global configs list containing a TrainConfig object, then run: python -m zoology.launch <path_to_script>.py

2. Sweep: Create a python file that populates the configs list with multiple TrainConfig objects (e.g., iterating over learning rates), then run the same launch command.

3. Parallel Sweep: If you have multiple GPUs and have installed the [extra] dependencies (specifically Ray), add the -p flag to run configurations in parallel across your devices.

# Run a single basic experiment
python -m zoology.launch zoology/experiments/basic_examples/basic.py

# Run a sweep over learning rates
python -m zoology.launch zoology/experiments/basic_examples/basic_sweep.py

# Run a sweep in parallel across multiple GPUs
python -m zoology.launch zoology/experiments/mqar_example_configs/original_mqar_configs.py -p

To define a new synthetic task in Zoology, you must follow a two-step process: subclassing zoology.config.DataSegmentConfig to define the task parameters and implementing a build method that returns a zoology.data.utils.DataSegment object.

1. Define the Configuration

Subclass zoology.config.DataSegmentConfig (which inherits from BaseConfig). You should define parameters like vocab_size, num_examples, and input_seq_len within this class.

2. Implement the build method

The build method must return a DataSegment dataclass. The inputs and labels fields in the DataSegment must be torch.Tensor objects containing integers in the range [0, vocab_size).

3. Register the task in your experiment config

Once your subclass is defined in an importable file, include it in your TrainConfig by adding instances of your new config class to the train_configs and test_configs lists within a DataConfig object.

# 1. Define the config
class MyDataSegmentConfig(DataSegmentConfig):
    def build(self, **kwargs) -> DataSegment:
        # Implementation returns a DataSegment
        return DataSegment(inputs=..., labels=...)

# 2. Use it in TrainConfig
from zoology.config import TrainConfig, DataConfig

config = TrainConfig(
    DataConfig(
        train_configs=[MyDataSegmentConfig(num_examples=10_000, vocab_size=128, input_seq_len=64)],
        test_configs=[MyDataSegmentConfig(num_examples=1_000, vocab_size=128, input_seq_len=64)],
    ),
)

This notebook provides a scaffolding for manually constructing the projection matrices ($Q, K, V$) for a two-layer Transformer designed to solve associative recall. The embeddings are constructed by concatenating positional encodings and one-hot token embeddings.

Layer 1: Shifting and Associating

• Goal of K1: Zero out token embeddings and keep only position embeddings.
• Goal of Q1: Shift position embeddings by one position.
• Goal of V1: Isolate token embeddings.
• Result (z1): A combination of the original embeddings and the shifted token embeddings (where token $e_i$ is now at position $i+1$).

Layer 2: Extracting the Value

• Goal of K2: Grab the shifted token embeddings from the previous layer.
• Goal of Q2: Isolate the original token embeddings from $z1$.
• Goal of V2: Extract the original token embedding that corresponds to the associative match.
• Result (O2): The final output containing the associative recall tokens.