ATOMIS

Introduction

About the Model

We introduce ATOMIS, developed by the Korea Institute of Nuclear Safety (KINS). This model is specifically designed for the nuclear field and is a large language model (LLM) with 32 billion parameters. It achieves state-of-the-art performance among its peers on Logickor, a real-world Korean task benchmark; NuclearQA, a nuclear-domain benchmark; and RAGEval, a RAG benchmark. Please refer to the evaluation results table for details.

Key Features

• Korean Real-World use cases: The model can understand and generate Korean text with high accuracy, making it suitable for practical scenarios.
• Specialized in the Nuclear Domain: The model has been specifically trained on a vast, specialized corpus of nuclear data.
• RAG: The model delivers accurate answers based on real documents through its high RAG performance.

Pre-Training

We created the base model by expanding layers using a passthrough method, building on the gemma-2-27b model. Additionally, we extended the context length to 32K with RoPE and performed continuous pretraining to restore the model’s performance. In particular, to train specialized knowledge in the nuclear domain, we included the following data.

• Atomic Wiki (https://atomic.snu.ac.kr)
• NText (https://paperswithcode.com/dataset/ntext)
• in-house data from KINS (Korea Institute of Nuclear Safety)

Post-Training

The fine-tuning data includes over 1M publicly available instruction datasets as well as high-quality synthetic data. We use this dataset to perform supervised fine-tuning (SFT) and direct preference optimization (DPO).

How to use

# pip install transformers==4.43.4 or later
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("KINS-ai/ATOMIS")
model = AutoModelForCausalLM.from_pretrained(
    "KINS-ai/ATOMIS",
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

messages = [
    {"role": "user", "content": "안녕하세요?"},
]

input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt", return_dict=True).to("cuda")

outputs = model.generate(**input_ids, max_new_tokens=256)
print(tokenizer.decode(outputs[0]))

Evaluation

Overall

LogicKor

We evaluated the performance using the LogicKor code. As the judge model, we employed the officially recommended GPT-4-1106-preview. These scores reflect only the default zero-shot evaluation.

NuclearQA

We employed NuclearQA [1], a human-made benchmark consisting of 100 questions designed by experts to evaluate language models in the nuclear domain.

We then used this question set to assess the LLM’s responses in a manner similar to the Logickor benchmark.

[1] Acharya, A., Munikoti, S., Hellinger, A., Smith, S., Wagle, S. and Horawalavithana, S., 2023. NuclearQA: A Human-Made Benchmark for Language Models for the Nuclear Domain. arXiv:2310.10920.

RAGEval

We used RAGEval [2], a benchmark designed to evaluate RAG performance in terms of factual accuracy, using three novel metrics: Completeness, Hallucination, and Irrelevance.

We evaluated performance using the RAGEval code. As the judge model, we employed the officially recommended gpt-4o. These scores reflect only the completeness metric of the single-document QA evaluation.

[2] Zhu, K., Luo, Y., Xu, D., Wang, R., Yu, S., Wang, S., Yan, Y., Liu, Z., Han, X., Liu, Z. and Sun, M., 2024. RAGEval: Scenario Specific RAG Evaluation Dataset Generation Framework. arXiv:2408.01262.