Improve model card and add metadata

README.md

@@ -1,17 +1,20 @@
 ---
-license: apache-2.0
 library_name: pytorch
+license: apache-2.0
+pipeline_tag: feature-extraction
 tags:
-  - tracelock
-  - dream
-  - diffusion-language-model
-  - activation-autoencoder
-  - pytorch
+- tracelock
+- dream
+- diffusion-language-model
+- activation-autoencoder
+- pytorch
 ---
 
 # TraceLock Dream Activation Autoencoder
 
-This repository contains the projection autoencoder checkpoint used to reproduce TraceLock on Dream.
+This repository contains the projection autoencoder checkpoint used to reproduce TraceLock on Dream, as presented in the paper [The Path Matters: Learning a Token-Commitment Policy for Diffusion Language Models](https://huggingface.co/papers/2605.24697).
+
+**Code**: [https://github.com/BobSun98/TraceLock](https://github.com/BobSun98/TraceLock)
 
 TraceLock is a token-level acceptance policy for Dream-style masked diffusion generation. Dream proposes candidate tokens during the denoising loop, and TraceLock decides which positions should be locked now versus kept masked for later refinement.
 
@@ -66,5 +69,16 @@ The exported projection state contains:
 
 ## Citation
 
-If you use this checkpoint, please cite the TraceLock paper/repository once available.
-
+If you use this checkpoint, please cite the following paper:
+
+```bibtex
+@misc{sun2026pathmatters,
+  title={The Path Matters: Learning a Token-Commitment Policy for Diffusion Language Models},
+  author={Bohang Sun and Max Zhu and Francesco Caso and Jindong Gu and Junchi Yu and Philip Torr and Pietro Liò and Jialin Yu},
+  year={2026},
+  eprint={2605.24697},
+  archivePrefix={arXiv},
+  primaryClass={cs.CL},
+  url={https://arxiv.org/abs/2605.24697}
+}
+```