A hierarchy of evidence for compositional competence reveals systematic limits in compositionality of meta-learned neural networks

Code accompanying the manuscript "A hierarchy of evidence for compositional competence reveals systematic limits in compositionality of meta-learned neural networks".

Credits

This work builds on Lake and Baroni's MLC implementation for grammar interpretation, episode generation, training and evaluation.
Lake, B. M. & Baroni, M. (2023). Human-like systematic generalization through a meta-learning neural network. Nature, 623, 115–121. Article, Code

The episodic-memory LSTM and parts of the SCAN episode sampling code are adapted from Lake's meta-seq2seq implementation.
Lake, B. M. (2019). Compositional generalization through meta sequence-to-sequence learning. Advances in Neural Information Processing Systems, 32. Article, Code

Repository contents

sweeps_and_probes_demo.ipynb      examples of probes and sweeps
configs/                          model, training and dataset settings
scripts/                          training, evaluation and plotting
src/                              models and shared functions
datasets/                         training and validation data
out_models/                       model checkpoints and training logs
results/                          evaluation data
figures/                          overview image and plotting instructions

Getting started

Install the dependencies with uv from the repository root:

uv sync --locked

Tested with Python 3.12.8.

CUDA is used when available; otherwise commands run on CPU. Use --device cpu or --device cuda to select it explicitly for training and evaluation.

Evaluation data and model checkpoints

The evaluation data, model checkpoints and training logs will be made available on Zenodo.

Once available, extract data.zip and models.zip into the repository root:

• Evaluation data (results/): episodes, predictions and summaries. See the results README.
• Model checkpoints and training logs (out_models/). See the model checkpoints README.

Download Lake and Baroni's algebraic dataset and place its train/ and val/ folders in datasets/data_algebraic/.

For commands to reproduce the paper's figures, see the figures README.

Evaluation

Benchmark evidence (Level 1)

Evaluate the 200 original validation episodes (2,000 queries) using greedy decoding:

uv run mlc-score validation --models transformer lstm --out-dir results_new/validation

This scores datasets/data_algebraic/val/ by default. Use --validation path/to/val to score another split. To select other model checkpoints, list their names from configs/models.json after --models.

Analyse the training data:

uv run python -m scripts.eval.benchmark_analysis structure \
    --out-dir results_new/benchmark_analysis

To analyse a different dataset, pass --dataset path/to/dataset.

Behavioural evidence (Level 2)

Probes and sweeps generate new episodes and exclude grammars found in datasets/data_algebraic/train/. Use --training-data path/to/dataset to exclude grammars from another dataset's train/ folder.

Probes, sweeps and data generation use the fast sampler by default. Pass --sampler original to use Lake and Baroni's original sampling implementation.

Parameter sweeps

A sweep varies task parameters across conditions. This example varies query input length, sampling 100 episodes per length:

uv run mlc-sweep --sweep length_x --max-len 12 --episodes 100 \
    --models transformer lstm --out-dir results_new/input_length

For both sweeps and probes, --models selects the checkpoints to score. The transformer checkpoint supplies the vocabulary for sampling by default. If you omit it from --models, use --reference-model to name a checkpoint from your selection.

Rule and primitive probes

A probe replaces selected grammar rules or primitive mappings. A battery runs a selection of probes.

Run selected probes:

uv run mlc-battery --probes swap_span holdout_prim_w3 --number-samples 100 \
    --models transformer lstm --out-dir results_new/injections

Summarize accuracy:

uv run python -m scripts.eval.battery_accuracy --dir results_new/injections

Extending probes and sweeps

The sweeps and probes demo shows how to define probes and sweeps and visualize their episodes. It needs no downloaded data or checkpoints.

Parameters

Support and query examples share a grammar and can have different sequence constraints. The table lists the parameters you can use to configure probes and sweeps.

What varies	Parameters
Number of support and query pairs	nsupport, nquery, min_nquery
Grammar size and rules	nprims, nrules, minlen_rules, maxlen_rules, n_unary_rules, n_binary_rules
Input and output lengths	minlen, maxlen, minlen_input, maxlen_input, minlen_output, maxlen_output
Rule applications	func_count, concat_count
Distinct rules applied	distinct_func_count, distinct_unary_count, distinct_binary_count
Nesting	min_nesting, max_nesting

By default, sampling requires nquery queries per episode. Set min_nquery to accept fewer. See src/data/datasets_grammar.py for the sampling defaults and the examples below for how to set parameters in a probe or sweep.

Probes

Add a dictionary with a unique name to PROBES in scripts/eval/battery.py. Use entries to specify which rules or primitive mappings to replace and options to set the episode parameters.

Rules use u1, u2 for single input tokens and x1, x2 for non-empty input sequences. f1 is an operator; brackets such as [u1] indicate the corresponding output sequence.

This probe returns the right token's output followed by the left token's output:

{
    "name": "primitive_swap",
    "entries": [["u1 f1 u2", "[u2] [u1]"]],
    "options": {
        "query": {"mode": "filter", "templates": ["u1 f1 u2"], "nquery": 10, "min_nquery": 1}
    },
}

Primitive entries such as ["p1", "c1 c2 c3"] map a primitive to three distinct colours sampled for each episode.

Use options to configure the probe.

• Episode sampling. params sets shared parameters for support and queries. support overrides them for the support set.
• Query forms. Set query.mode to choose how templates are used.
  • "construct" (default) builds queries from the templates. Use query.depth to control where the injected rule applies.
  • "filter" keeps sampled queries that match the templates.
• Rule replacement. host_arity specifies how many arguments the rule being replaced takes.
• Primitive holdout. hold_out_bare=True requires the first injected primitive to appear on its own and in at least two longer support examples. The sampler removes the standalone support example and tests the primitive on its own as a query.
• Decoding. decoder_cap sets the maximum decoded output length.

Evaluate the new probe and summarize accuracy:

uv run mlc-battery --probes primitive_swap --number-samples 20 \
    --models transformer lstm --out-dir results_new/primitive_swap
uv run python -m scripts.eval.battery_accuracy --dir results_new/primitive_swap

Sweeps

Add a named entry after the SWEEPS dictionary in scripts/eval/sweep.py. Set the parameters shared by all conditions in support and query. Each entry in cells sets the parameters for one condition and labels it with value. For example, vary query input length:

SWEEPS["query_length"] = {
    "query": {"nquery": 4, "min_nquery": 4},
    "cells": [
        {"value": 2, "query": {"minlen_input": 2, "maxlen_input": 2}},
        {"value": 3, "query": {"minlen_input": 3, "maxlen_input": 3}},
    ],
}

uv run mlc-sweep --sweep query_length --episodes 20 \
    --models transformer lstm --out-dir results_new/query_length

Training

In mlc-train, use --model_name to choose between the Transformer and the episodic-memory LSTM. The table lists the default configuration files for each model:

--model_name	Model configuration	Training configuration
mlc_transformer	MLC_TRANSFORMER_config.json	train_config.json
mlc_lstm	MLC_LSTM_config.json	train_config_lstm.json

To change a run, copy the relevant JSON file, edit its settings and select the copy:

• --model_config_file: layers, hidden size, dropout and other architecture settings.
• --train_config_file: epochs, batch size, learning rate and checkpoint selection.
• --dataset_config_file: dataset directory and training condition; defaults to dataset_config.json.

uv run mlc-train --model_name mlc_transformer --additional-context \
    --dir_model out_models/my_transformer --fn_out_model transformer.pt --seed 0

uv run mlc-train --model_name mlc_lstm --additional-context \
    --dir_model out_models/my_lstm --fn_out_model lstm.pt --seed 0

--additional-context includes support examples in the encoder input; --no-additional-context leaves them out.

Model checkpoints and training logs are saved in --dir_model. Add --resume to the same command to continue an interrupted run.

The model checkpoints README explains how to evaluate your own checkpoints or add a model architecture.

Generating training data

Generate a new dataset with the original MLC defaults: 100,000 training episodes and 200 validation episodes.

uv run mlc-generate --data_dir datasets/my_algebraic --seed 0

Use --nsamp_train and --nsamp_val to change the number of episodes. Training and validation settings can be changed separately. For example, --nsupport_grammar_train 20 sets 20 support pairs per training episode; --nsupport_grammar_val 20 does the same for validation. Run uv run mlc-generate --help to see all options.

Set data_dir to my_algebraic in your dataset config and pass it to mlc-train with --dataset_config_file.

Licence

The original code is licensed under MIT. Credits and licence terms for adapted code are listed in NOTICE.

The evaluation data, model checkpoints and training logs in the Zenodo release are licensed under CC BY 4.0.