def fit(*, signal: bool, seed: int, steps: int | None, accelerator: str) -> rf.Model:
lit.seed_everything(seed, workers=True)
model = rf.Model(
name="state",
d_model=16,
n_layers=1,
n_heads=4,
batch_size=128,
measurement=rf.Number,
target=rf.Boolean(mask=True),
)
model.optimizer = lambda module: torch.optim.AdamW(module.parameters(), lr=5e-3)
data = rf.SyntheticDataModule(
model=model,
train=partial(records, rows=1024, seed=seed + 1, signal=signal),
validate=partial(records, rows=512, seed=seed + 2, signal=signal),
seed=seed,
)
trainer = lit.Trainer(
accelerator=accelerator,
max_epochs=8,
max_steps=steps if steps is not None else -1,
logger=False,
enable_progress_bar=False,
enable_model_summary=False,
enable_checkpointing=False,
deterministic=True,
)
trainer.fit(model=model, datamodule=data)
return model
def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
signal_model = fit(signal=True, seed=seed, steps=steps, accelerator=accelerator)
control_model = fit(signal=False, seed=seed, steps=steps, accelerator=accelerator)
state_auc, state_accuracy = score(
signal_model, partial(records, rows=2048, seed=seed + 3, signal=True), accelerator
)
control_auc, control_accuracy = score(
control_model, partial(records, rows=2048, seed=seed + 3, signal=False), accelerator
)
prefilled_auc, prefilled_accuracy = score(
signal_model, partial(records, rows=2048, seed=seed + 3, signal=True, fill_nulls=True), accelerator
)
return {
"state_auc": state_auc,
"state_accuracy": state_accuracy,
"independent_auc": control_auc,
"independent_accuracy": control_accuracy,
"prefilled_auc": prefilled_auc,
"prefilled_accuracy": prefilled_accuracy,
"auc_gap": state_auc - control_auc,
}, {
"Null state reaches 0.99 AUC": state_auc >= 0.99,
"Null state reaches 0.98 accuracy": state_accuracy >= 0.98,
"Independent validity AUC remains between 0.42 and 0.58": 0.42 <= control_auc <= 0.58,
"Independent validity accuracy remains between 0.45 and 0.55": 0.45 <= control_accuracy <= 0.55,
"Null-to-zero AUC remains between 0.42 and 0.58": 0.42 <= prefilled_auc <= 0.58,
"Null-to-zero accuracy remains between 0.45 and 0.55": 0.45 <= prefilled_accuracy <= 0.55,
"State exceeds its control by at least 0.40 AUC": state_auc - control_auc >= 0.40,
}