Add pipeline tag and improve model card

README.md

@@ -1,10 +1,15 @@
 ---
-license: apache-2.0
 language:
 - en
+license: apache-2.0
+pipeline_tag: feature-extraction
+base_model:
+- google/gemma-2-2b
+tags:
+- mechanistic-interpretability
+- sparse-autoencoder
 ---
 
-
 # 🧩 Model Weights for *Towards Atoms of Large Language Models*
 
 [![arXiv](https://img.shields.io/static/v1?label=Paper&message=Arxiv&color=red&logo=arxiv)](https://arxiv.org/abs/2509.20784)  [![GitHub](https://img.shields.io/static/v1?label=Code&message=Github&color=blue&logo=github)](https://github.com/ChenhuiHu/towards_atoms)  [![Hugging Face](https://img.shields.io/static/v1?label=Models&message=HuggingFace&color=yellow&logo=huggingface)](https://huggingface.co/Ericccccc/towards_atoms)
@@ -13,10 +18,22 @@ This repository contains the model weights associated with the paper:
 
  👉 [***Towards Atoms of Large Language Models***](https://arxiv.org/abs/2509.20784)
 
-Specifically, it provides the weights of threshold-activated sparse autoencoders (SAEs) trained on activations across layers of Gemma2-2B, using the CounterFact dataset. 
+Specifically, it provides the weights of threshold-activated sparse autoencoders (TSAEs) trained on activations across layers of [Gemma 2 2B](https://huggingface.co/google/gemma-2-2b), using the CounterFact dataset. 
+
+## Description
+The paper introduces **Atom Theory** to define and identify the fundamental representational units (FRUs) of large language models, termed **atoms**. Using threshold-activated sparse autoencoders (TSAEs) and a non-Euclidean metric called the atomic inner product (AIP), the authors identify units with near-perfect faithfulness and stability across layers of Gemma2.
 
-Note that only the model weights are included in this repository. 
+## Usage
+Note that only the model weights are included in this repository. For the complete implementation, including training scripts, data preprocessing, and evaluation pipelines, please refer to the main codebase:
 
-For complete implementation, including training scripts, data preprocessing, and evaluation pipelines, please refer to the main codebase:
+👉 [https://github.com/ChenhuiHu/towards_atoms](https://github.com/ChenhuiHu/towards_atoms)
 
- 👉 https://github.com/ChenhuiHu/towards_atoms
+## Citation
+```bibtex
+@article{hu2025towards,
+  title={Towards Atoms of Large Language Models},
+  author={Hu, Chenhui and Cao, Pengfei and Chen, Yubo and Liu, Kang and Zhao, Jun},
+  journal={arXiv preprint arXiv:2509.20784},
+  year={2025}
+}
+```