fermi-bert-1024: Pretrained BERT for Nuclear Power

A BERT model optimized for the nuclear energy domain, fermi-bert-1024 is pretrained on a combination of Wikipedia (2023), Books3, and a subset of the U.S. Nuclear Regulatory Commission’s ADAMS database. It is specifically designed to handle the complex technical jargon and regulatory language unique to the nuclear industry. Trained on the Oak Ridge National Laboratory Frontier supercomputer using 128 MI250X AMD GPUs over a 10-hour period, this model provides a robust foundation for fine-tuning in nuclear-related applications.

Training

fermi-bert-1024 is a BERT model pretrained on wikipedia (2023), Books3, and ADAMS with a max sequence length of 1024.

We make several modifications to the standard BERT training procedure:

• We use a custom nuclear-optimized WordPiece tokenizer to better represent the unique jargon and technical terminology specific to the nuclear industry.
• We train on a subset of U.S. Nuclear Regulatory Commission’s Agency-wide Documents Access and Management System (ADAMS).
• We train on Books3 rather than BookCorpus.
• We use larger batch size and other improved hyper parameters as described in RoBERTa.

Evaluation

We evaluate the quality of fermi-bert-1024 on the standard GLUE benchmark (script). We find it performs comparably to other BERT models but with the advantage of performing better on documents in the nuclear energy space as demonstrated by our downstream fine-tuning.

Model	Bsz	Steps	Seq	Avg	Cola	SST2	MRPC	STSB	QQP	MNLI	QNLI	RTE
bert-base-uncased	256	1M	512	0.81	0.56	0.82	0.86	0.88	0.91	0.84	0.91	0.67
roberta-base	8K	500k	512	0.84	0.56	0.94	0.88	0.90	0.92	0.88	0.92	0.74
fermi-bert-512	4k	100k	512	0.83	0.60	0.93	0.88	0.89	0.91	0.87	0.91	0.68
fermi-bert-1024	4k	100k	1024	0.83	0.6	0.93	0.86	0.89	0.91	0.86	0.92	0.69

Pretraining Data

We train on 40% Wikipedia, 30% Books3, 30% ADAMS. We pack and tokenize the sequences to 1024 tokens. If a document is shorter than 1024 tokens, we append another document until it is 1024 tokens. If a document is longer than 1024 tokens we split it into multiple documents. For 10% of the Wikipedia documents, we do not concatenate short documents. See M2-Bert for rationale behind including short documents.

Usage

from transformers import AutoModelForMaskedLM, AutoConfig, AutoTokenizer, pipeline

tokenizer = AutoTokenizer.from_pretrained('atomic-canyon/fermi-bert-1024') # `fermi-bert` uses a nuclear specific tokenizer
model = AutoModelForMaskedLM.from_pretrained('atomic-canyon/fermi-bert-1024')

# To use this model directly for masked language modeling
classifier = pipeline('fill-mask', model=model, tokenizer=tokenizer, device="cpu")

print(classifier("I [MASK] to the store yesterday."))

Acknowledgement

This research used resources of the Oak Ridge Leadership Computing Facility, which is a DOE Office of Science User Facility supported under Contract DE-AC05-00OR22725.