Wrenformer

[janosh]

This PR adds a new variant of the Wren model called Wrenformer, a rewrite that drops the custom self-attention in favor of builtin PyTorch TransformerEncoder. It also prepares Matbench submissions for Roost and Wren (structure tasks only for Wren).

The Wren rewrite as a transformer encoder was necessary to run robustly across Matbench tasks as the original Wren would run into out-of-memory errors during training and inference on materials with large numbers of Wyckoff positions (< 16 was a safe cutoff). This happened even on A100 GPUs with 80 GB RAM.

Initial performance testing (training for 100 epochs and with only 3 transformer layers) suggests Wrenformer slightly beats Wren across all Matbench tasks I recorded for both.

More sets of hyperparameters

Speed difference between Wren and Wrenformer

According to @CompRhys, Wren could run 500 epochs in 5.5 h on a P100 training on 120k samples of MP data (similar to the matbench_mp_e_form dataset with 132k samples). Wrenformer only managed 207 epochs in 4h on the more powerful A100 training on matbench_mp_e_form. However, to avoid out-of-memory issues, Rhys constrained Wren to only run on systems with <= 16 Wyckoff positions. The code below shows that this lightens the workload by a factor of about 7.5, likely explaining the apparent slowdown in Wrenformer.

import pandas as pd
from aviary.wren.utils import count_wyks
from examples.mat_bench import DATA_PATHS

df = pd.read_json(DATA_PATHS["matbench_mp_e_form"])

df["n_wyckoff"] = df.wyckoff.map(count_wyks)


sum_wyckoffs_sqr = (df.n_wyckoff**2).sum()
sum_wyckoffs_lte_16_sqr = (df.query("n_wyckoff <= 16").n_wyckoff ** 2).sum()
print(f"{sum_wyckoffs_sqr=}")
print(f"{sum_wyckoffs_lte_16_sqr=}")
print(f"{sum_wyckoffs_sqr/sum_wyckoffs_lte_16_sqr=:.3}")
# prints 7.45, so Wrenformer has to do 7.45x more work, explaining the about 2x slow down
# on a more powerful GPU (Nvidia A100 vs Wren on a P100)