GenSPP: interlocking-free rationalization through genetic search#

Ruggeri and Signorelli, 2025, Interlocking-free Selective Rationalization Through Genetic-based Learning, ACL 2025. Paper: https://aclanthology.org/2025.acl-long.59/. Reference implementation: nlp-unibo/gen-spp.

Problem#

Every architecture on the preceding pages mitigates interlocking without removing its cause. The cause is that one loss trains both modules at once, so the selector’s gradient is computed through a predictor that the selector’s own past choices shaped, and no amount of freezing, staging or regularising changes that the two are optimised together. Interlocking is therefore a property of joint gradient descent rather than of any particular objective.

Method#

GenSPP stops training the generator by gradient descent. A genetic search runs over generator parameters, and each candidate is scored by training a fresh predictor on the selection it produces, so a candidate is judged by how well a predictor that owes it nothing can classify from its highlight. Since no predictor is carried between candidates, there is no cooperative equilibrium to fall into, and the interlocking loop is not mitigated but absent.

        %%{init: {"theme": "base", "themeVariables": {"fontSize": "17px", "fontFamily": "Lato, sans-serif", "lineColor": "#37474f", "primaryTextColor": "#102027", "edgeLabelBackground": "#ffffff"}, "flowchart": {"nodeSpacing": 55, "rankSpacing": 70, "padding": 14, "curve": "basis"}}%%
flowchart TD
    POP["population of generators"] --> C["one candidate"]
    C --> H["its highlight of the corpus"]
    H --> FP["a fresh predictor,<br/>trained from scratch"]
    FP --> F["fitness:<br/>selection rate against task loss"]
    F --> S["survivors, crossover, mutation"]
    S --> POP

    classDef shared fill:#cfe3ff,stroke:#1a4f9c,stroke-width:2px,color:#0b2545
    classDef head fill:#ffffff,stroke:#37474f,stroke-width:2px,color:#102027
    classDef value fill:#d7f0dc,stroke:#1e6b34,stroke-width:2px,color:#0d3018
    classDef term fill:#ffe9c7,stroke:#a35c00,stroke-width:2px,color:#3d2100
    class C,FP head
    class POP,H,S value
    class F term

The fitness trades what a candidate selected against what its predictor achieved.

\[\begin{split}\text{fitness} = \begin{cases} 1.0, & \text{if } \mathcal{L}_{\text{task}} > \texttt{task\_loss\_limit} \\[4pt] \dfrac{1}{1 - \sqrt{(1 - r)(1 - \mathcal{L}_{\text{task}})}}, & \text{otherwise} \end{cases}\end{split}\]

Here \(r\) is the selection rate of the candidate’s highlight and \(\mathcal{L}_{\text{task}}\) is the cross entropy its freshly trained predictor reached. The limit is what stops the search from buying a sparse selection with a model that has stopped classifying: above it a candidate scores the floor whatever it selected. Below it both terms have to be small for the fitness to be large, so a candidate cannot win on sparsity alone. The paper sets the limit per corpus, 0.1 on the toy corpus, which is nearly solved, and 0.6 on HateXplain, which is not.

Training#

The search is a trainer rather than a training step, and it owns the loop.

A founding population of population_size generators is drawn at random, each carrying its parameters as a chromosome.
Each candidate is scored on one device: its generator is frozen, a fresh predictor is trained on its selection for predictor_epochs, and the fitness above is computed.
Survivors are chosen by half elitism, with the better half kept outright and the rest drawn in proportion to fitness.
Couples are crossed one-point and the children mutated by Gaussian noise on a share of their genes.
The generation is repeated until n_generations or until the best objective stops improving by stop_threshold.

Gradient descent still happens inside step 2, and it only ever reaches the predictor. Training a GenSPP model on its own therefore fits a predictor to whatever selection its untrained generator makes, which is why the search is the thing a study runs.

Implementation#

Two classes rather than one: the model, and the search that fits it.

What	Where
The model	`GenSPP`
The search	`GenSPPTrainer`
Selection without the empty repair	`GenSPP.select`, which scores an empty selection rather than repairing it
Only the predictor is optimised	`GenSPP.configure_optimizers`
The generator held still while a predictor is fitted	`GenSPP.on_train_epoch_start`
The fitness	`GenSPPTrainer.compute_fitness`
Survivors, crossover, mutation	`GenSPPTrainer._select_survivors`, `_crossover`, `_mutate`

GenSPP refuses a configuration whose generator and predictor share parameters, since a search over generator parameters that also moved the predictor would be neither a search nor a training run. The generator’s modules are put in evaluation mode for the whole of a candidate’s predictor fitting, because dropout inside a frozen generator would score the same candidate differently from one epoch to the next. Candidates are evaluated one per device, and devices is the same knob for a pool of CPU workers and for a node’s GPUs.

CPU workers are processes, CUDA workers are threads. A candidate is a small model, so its cost is the training loop stepping from Python rather than the arithmetic inside torch, and that loop holds the GIL: eight threads on eight cores were measured at 240% of a possible 800%. Processes lift that. One generation of the Toy search, 100 candidates at a population of fifty, ran in 142.6 s sequentially, 91.0 s on eight threads and 22.6 s on eight processes: 1426, 910 and 226 ms a candidate, or 2.00, 1.28 and 0.32 hours for the 5050 candidates of a seed. All three settled on the same chromosome. CUDA is the other way round, since its kernels do release the GIL and a process per device would pay for a context each. A search falls back to threads where fork is unavailable, where autograd has already run in the calling process, which torch refuses to combine with fork, or where the calling process has initialised CUDA. On a machine with a GPU, one optimizer step initialises CUDA even for parameters on the CPU, so a search that runs after another model trained in the same process scores on threads. The search logs which of these kept it on threads.

Why the pool is forked once, before the first candidate, rather than per generation:

RuntimeError: Unable to handle autograd's threading in combination with
fork-based multiprocessing.

Torch raises that once autograd has run threads in the parent, and it raises it in the child when the pass is attempted rather than at the fork. So a pool is asked to train something trivial before it is trusted with a candidate, and one that cannot, or that does not answer within the probe timeout, is closed for threads. forkserver would be the start method that avoids forking a threaded process, and it cannot be used here: the registry is process-global state built once by the caller, and a worker that did not inherit it cannot build the model a chromosome is for:

NotExpandedException: The registration graph has yet to be expanded!
Configuration retrieval is not allowed.

Rebuilding it per worker would cost seconds each and register a second copy of every configuration. Inheriting memory is what makes these workers correct, not merely cheap, so the fork risk is bounded rather than removed. A search opens its pool with one Python thread running, torch’s being native, which is why CPython’s own warning about forking a multi-threaded process does not fire outside a test runner that adds threads of its own.

One shared initial state is given to every candidate, so that a fitness is a property of a chromosome rather than of the predictor initialisation drawn alongside it. The reference implementation reaches the same place from the other side: it keeps a pool of models and resets each reused one to that slot’s initial weights, which makes a candidate’s predictor depend on the slot it was given.

Configuration#

Two keys per backbone here, since the model and its search are registered separately.

Key	Configuration
`GRU_GENSPP`	`GRUGenSPPConfig`
`GRU_GENSPP_TRAINER`	`GRUGenSPPTrainerConfig`
`TRANSFORMER_GENSPP`	`TransformerGenSPPConfig`
`TRANSFORMER_GENSPP_TRAINER`	`TransformerGenSPPTrainerConfig`

The defaults are the reference implementation’s: 100 generations over a population of 50, 3 predictor epochs per candidate, a mutation standard deviation of 0.05, and one CPU worker.

One deviation is deliberate. The reference implementation mutates the selector’s output bias at 0.10 while every other gene takes 0.05, which its paper does not report: the paper gives a single N(0.0, 0.05). The default here follows the paper, and threshold_mutation_std sets the threshold’s own deviation for a reproduction that needs the reference implementation’s behaviour. It names the deviation of the threshold rather than of one gene, since a head emitting two logits decides on their difference and two genes perturbed at s give that difference s * sqrt(2).

Which genes those are is asked of the selector rather than read off the end of the chromosome. threshold_parameters() returns them, and a selector that declares none is searched with a single deviation throughout. MLPSelector declares the bias of its output layer; only parameters that land at the end of the flattened chromosome are counted, because the deviation is given to a trailing slice of it.

from pyhighlights.configurations.keys import GRU_GENSPP_TRAINER

search = Registry.from_key(
    GRU_GENSPP_TRAINER,
    n_generations=50,
    population_size=20,
    devices=["cuda:0", "cuda:1"],
)

A search of that size is expensive, and ToyGenSPPTrainerConfig exists for the case where the wiring rather than the result is what is being checked: two candidates and one generation, which finishes in seconds and means nothing. The model’s losses list holds the classification term alone, since the sparsity of a candidate is part of its fitness rather than a penalty on a gradient, and its backbones freeze their embeddings by default.

Running a search#

GenSPP’s generator is not trained: it is searched. A population of generators is evolved, and a candidate is scored by fitting a predictor on the selections it makes, with the generator frozen, so the predictor never teaches the selector what to select, which is the cooperative equilibrium the other models have to fight.

from pyhighlights.configurations.keys import TOY_GENSPP_TASK

Registry.from_key(TOY_GENSPP_TASK, seeds=[42]).run()

GenSPPTask is an SPPTask in every other respect: same corpus, preprocessing, metrics, seeds and output files. Two things differ:

It names a search, not a model. The model key is the search’s own; naming it twice is a way for the two to disagree about which model was evolved. The search builds its own candidates, so a vector file named through embeddings reaches them through the search rather than through build_model, which a searched model never goes through.
A validation split is required. Fitness is task loss traded against selection rate, and both are measured there.

A generation draws int(selection_rate * population_size) couples by roulette wheel and crosses each into two children, so the default 0.5 adds one child per member: 25 couples and 50 children against a population of 50. The children compete with their parents rather than replacing them, and survival is half elitism, the best half kept outright and the other half drawn from what is left, fitness-proportional and without replacement. The population that comes out of a generation is the size that went in.

Each candidate’s predictor is fitted by a throwaway Lightning trainer, so the inner training is the same code path every other model trains through, logging, checkpointing and sanity checks off, since a hundred generations build one trainer per candidate. Gradients reach the predictor only: GenSPP.configure_optimizers hands over the predictor’s parameters, and the generator is put back in evaluation mode at the start of every epoch so its dropout cannot score the same candidate two different ways.

Every candidate trains on the same batches in the same order. The training loader shuffles, so the search draws one permutation from the seed before it starts and hands every candidate that: re-iterating the loader instead would give each candidate its own order, and a chromosome would score differently depending on how many candidates preceded it, or, with several devices, on how the workers interleaved. The whole training split is held in memory for the duration of the search as a result.

Alongside the usual per-seed files, a GenSPP run writes best.ckpt, the weights the search settled on, and search.json, one entry per generation under training_progress. The entry is the best objective the search reached in that generation, 1 / fitness, so it falls as the search improves and it is what stop_threshold is compared against. A search that stopped improving in its tenth generation and one still descending when the budget ran out report the same metrics otherwise.

API#

class pyhighlights.components.models.spp.genspp.GenSPP(**kwargs)[source]#

Bases: SPP

Single-generator SPP evaluated through external genetic search.

The generator is not trained by gradient descent: a genetic search over generator parameters scores each candidate by training a fresh predictor on it, which removes the cooperative equilibrium the other SPP models have to fight. See GenSPPTrainer for the search itself.

Ruggeri and Signorelli, 2025, Interlocking-free Selective Rationalization Through Genetic-based Learning, ACL 2025, <https://aclanthology.org/2025.acl-long.59/>. Reference implementation: <nlp-unibo/gen-spp>.

generator_parameters()[source]#

The evolvable half, which is narrower than the family’s answer.

A chromosome is what the search may change, so a frozen parameter is not part of one: the embedding table every candidate loads is shared state, and including it would make the chromosome the corpus.

Return type:: List[Parameter]

predictor_parameters()[source]#

What descent moves in a candidate, frozen encoders excluded.

Return type:: List[Parameter]

class pyhighlights.components.models.spp.genspp.GenSPPTrainer(model, n_generations=100, population_size=50, selection_rate=0.5, mutation_probability=1.0, mutation_std=0.05, threshold_mutation_std=None, predictor_epochs=3, task_loss_limit=0.1, stop_threshold=0.01, seed=None, devices=('cpu',))[source]#

Bases: object

External search matching released GenSPP’s selection-rate objective.

Parameters:

model (RegistrationKey[GenSPP])
n_generations (int)
population_size (int)
selection_rate (float)
mutation_probability (float)
mutation_std (float)
threshold_mutation_std (float | None)
predictor_epochs (int)
task_loss_limit (float)
stop_threshold (float)
seed (int | None)
devices (Sequence[str])

static compute_fitness(task_loss, selection_rate, task_loss_limit)[source]#

Selection rate traded against task loss, higher being better.

task_loss_limit is the cross entropy above which a candidate is not competing at all: it gets the floor of 1.0 whatever it selected, so the search cannot buy a sparse selection with a model that has stopped classifying. Set it per corpus: a corpus the predictor nearly solves supports a lower limit than one it does not.

Below the limit the objective is 1 - sqrt((1 - selection_rate) * (1 - task_loss)), and the fitness is its reciprocal: both terms have to be small for it to be large, so a candidate cannot win on sparsity alone.

Return type:

float

Parameters:

task_loss (float)
selection_rate (float)
task_loss_limit (float)

fit(train_loader, val_loader, embeddings=None)[source]#

Search for a generator, and return the best model the search found.

embeddings is the corpus’s token table, when the task read one. It is passed rather than registered for the reason SPPTask passes it: a matrix is data, and no configuration should carry one. Every candidate loads the same table, before _align_initial_state() sees it, so it is shared state rather than part of a chromosome.

Return type:

GenSPP

Parameters:

train_loader (Iterable[InputData])
val_loader (Iterable[InputData])
embeddings (Tensor | None)

selection_rate#: How many couples a generation draws, as a share of the population. Each couple crosses into two children, so 0.5, the reference implementation’s value, adds one child per member: int(0.5 * 50) = 25 couples and 50 children, which then compete with their 50 parents for 50 places.

threshold_mutation_std#: Standard deviation of the perturbation a mutation applies to the selector’s decision threshold, or None to perturb it at mutation_std like every other gene. It names the deviation of the threshold rather than of one gene, and the genes carrying it are the ones the selector declares. None follows the paper; see GenSPP: interlocking-free rationalization through genetic search for what the release does instead.

training_progress: List[float]#

the best objective reached in it, which is 1 / fitness and therefore falls as the search improves. It is what stop_threshold is compared against. The name is the released implementation’s rather than a description, and it is serialized under itself inside search.json, so it is left alone rather than renamed under existing readers.

Type:: One entry per generation

pyhighlights.components.models.spp.genspp.PROBE_SECONDS = 60.0#: How long a worker gets to answer the probe below. Generous, because it is paid once per search and a loaded node can be slow to schedule a fork. It is short against a search measured in hours, which is what it protects. A probe that never answers is a deadlocked fork; the timeout turns that into a search that runs on threads instead. See GenSPPTrainer._open_pool().

GenSPP: an SPP whose generator is searched genetically, plus its trainer.

class pyhighlights.configurations.genspp.ToyGenSPPTrainerConfig(**data)[source]#

A search small enough to finish: two candidates, one generation.

Nothing here is a sensible experiment. It exists so the wiring of corpus, search, scoring and serialization can be exercised in seconds.

Parameters:

model (RegistrationKey[GenSPP])
n_generations (int)
population_size (int)
selection_rate (float)
mutation_probability (float)
mutation_std (float)
predictor_epochs (int)
task_loss_limit (float)
stop_threshold (float)
seed (int | None)
devices (List[str])

model_post_init(_Configuration__context)#

Runs automatically right after Pydantic instantiates an object.

Return type:: None
Parameters:: _Configuration__context (Any)