diff --git a/CITATION.cff b/CITATION.cff
new file mode 100644
index 0000000..ed8e281
--- /dev/null
+++ b/CITATION.cff
@@ -0,0 +1,74 @@
+cff-version: 1.2.0
+message: "If you use Trainax in your research, please cite the NeurIPS 2024 APEBench paper below."
+title: "Trainax: Learning Methodologies for Autoregressive Neural Emulators"
+type: software
+authors:
+  - given-names: Felix
+    family-names: Koehler
+    affiliation: "Technical University of Munich"
+repository-code: "https://github.com/Ceyron/trainax"
+url: "https://fkoehler.site/trainax/"
+license: MIT
+keywords:
+  - "autoregressive models"
+  - "neural emulators"
+  - "neural operators"
+  - "unrolled training"
+  - "differentiable physics"
+  - "time series forecasting"
+  - "JAX"
+  - "Equinox"
+  - "Optax"
+  - "scientific machine learning"
+identifiers:
+  - type: url
+    value: "https://pypi.org/project/trainax/"
+    description: "PyPI package"
+preferred-citation:
+  type: conference-paper
+  title: "APEBench: A Benchmark for Autoregressive Neural Emulators of PDEs"
+  authors:
+    - given-names: Felix
+      family-names: Koehler
+      affiliation: "Technical University of Munich"
+    - given-names: Simon
+      family-names: Niedermayr
+      affiliation: "Technical University of Munich"
+    - given-names: Rüdiger
+      family-names: Westermann
+      affiliation: "Technical University of Munich"
+    - given-names: Nils
+      family-names: Thuerey
+      affiliation: "Technical University of Munich"
+  collection-title: "Advances in Neural Information Processing Systems 37 (NeurIPS 2024)"
+  conference:
+    name: "Thirty-eighth Conference on Neural Information Processing Systems"
+  year: 2024
+  volume: 37
+  start: 120252
+  end: 120310
+  doi: "10.52202/079017-3822"
+  url: "https://proceedings.neurips.cc/paper_files/paper/2024/hash/d9875ebcf74bccdc5076acab0dbee62c-Abstract-Datasets_and_Benchmarks_Track.html"
+  identifiers:
+    - type: other
+      value: "arXiv:2411.00180"
+      description: "arXiv preprint"
+  abstract: >-
+    We introduce the Autoregressive PDE Emulator Benchmark (APEBench), a
+    comprehensive benchmark suite to evaluate autoregressive neural emulators
+    for solving partial differential equations. APEBench is based on JAX and
+    provides a seamlessly integrated differentiable simulation framework
+    employing efficient pseudo-spectral methods, enabling 46 distinct PDEs
+    across 1D, 2D, and 3D. Facilitating systematic analysis and comparison of
+    learned emulators, we propose a novel taxonomy for unrolled training and
+    introduce a unique identifier for PDE dynamics that directly relates to
+    the stability criteria of classical numerical methods. APEBench enables
+    the evaluation of diverse neural architectures, and unlike existing
+    benchmarks, its tight integration of the solver enables support for
+    differentiable physics training and neural-hybrid emulators. Moreover,
+    APEBench emphasizes rollout metrics to understand temporal generalization,
+    providing insights into the long-term behavior of emulating PDE dynamics.
+    In several experiments, we highlight the similarities between neural
+    emulators and numerical simulators. The code is available at
+    github.com/tum-pbs/apebench and APEBench can be installed via
+    `pip install apebench`.
\ No newline at end of file
diff --git a/README.md b/README.md
index d76257e..f393d10 100644
--- a/README.md
+++ b/README.md
@@ -222,10 +222,11 @@ it:
 
 ```bibtex
 @article{koehler2024apebench,
-  title={{APEBench}: A Benchmark for Autoregressive Neural Emulators of {PDE}s},
-  author={Felix Koehler and Simon Niedermayr and R{\"}udiger Westermann and Nils Thuerey},
-  journal={Advances in Neural Information Processing Systems (NeurIPS)},
-  volume={38},
+  title={Apebench: A benchmark for autoregressive neural emulators of pdes},
+  author={Koehler, Felix and Niedermayr, Simon and Westermann, R{\"u}diger and Thuerey, Nils},
+  journal={Advances in Neural Information Processing Systems},
+  volume={37},
+  pages={120252--120310},
   year={2024}
 }
 ```
diff --git a/docs/index.md b/docs/index.md
index aa5d383..79ff1aa 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -90,12 +90,14 @@ it:
 
 ```bibtex
 @article{koehler2024apebench,
-  title={{APEBench}: A Benchmark for Autoregressive Neural Emulators of {PDE}s},
-  author={Felix Koehler and Simon Niedermayr and R{\"}udiger Westermann and Nils Thuerey},
-  journal={Advances in Neural Information Processing Systems (NeurIPS)},
-  volume={38},
+  title={Apebench: A benchmark for autoregressive neural emulators of pdes},
+  author={Koehler, Felix and Niedermayr, Simon and Westermann, R{\"u}diger and Thuerey, Nils},
+  journal={Advances in Neural Information Processing Systems},
+  volume={37},
+  pages={120252--120310},
   year={2024}
 }
+
 ```
 
 (Feel free to also give the project a star on GitHub if you like it.)