---
library_name: pytorch
license: apache-2.0
pipeline_tag: feature-extraction
tags:
- 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, 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.

## What This Checkpoint Is

`best_val_loss.pt` is an activation autoencoder for Dream hidden states. It compresses the last three Dream hidden-state snapshots and two hidden-state deltas into compact features consumed by the TraceLock policy model.

This checkpoint is not a text generation model and does not contain Dream model weights. Users still need to download Dream from its original repository:

```text
Dream-org/Dream-v0-Instruct-7B
```

## How It Is Used

After downloading this repository into a TraceLock workspace, the expected local path is:

```text
$TRACELOCK_HOME/checkpoints/dream-ae-v1/best_val_loss.pt
```

TraceLock uses this checkpoint in two places:

1. `generate_training_traces.sh`: projects Dream activations while building training traces.
2. `train.sh` / evaluation: reconstructs the same projection stack expected by the TraceLock policy.

## Architecture

The released checkpoint was trained with:

```json
{
  "d_model": 3584,
  "d_hidden_bottleneck": 256,
  "d_delta_bottleneck": 32,
  "dropout": 0.1
}
```

The exported projection state contains:

- hidden-state normalization
- delta-state normalization
- hidden-state projection encoder
- delta-state projection encoder

## Files

- `best_val_loss.pt`: projection autoencoder checkpoint.
- `config.json`: training/configuration metadata for this autoencoder run.
- `data_stats.json`: basic sample count and batch metadata from the run.

## Citation

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}
}
```