learned
Frozen GPT-2 reads the tokens; the hidden state two blocks before its last is taken and standardized feature by feature.
Why
Contextual. The latent space is built on the hidden states of a frozen pre-trained token encoder. The authors found a diffusion-friendly latent space much easier to learn when the trainable encoder works on contextual hidden states than on shallow token-level features: the hidden states already encode dependencies across the sequence. LDLM paper §5.1
A frozen token encoder also makes the feature loss possible: the latent decoder can be trained to reproduce these hidden states, in place of relying on cross-entropy alone. LDLM paper §5.1
Under a cross-entropy loss, GPT-2’s hidden states give a generative perplexity of 156.3, against 190.4 for GPT-2’s token embeddings and 272.4 for embeddings trained from scratch: the encoder benefits from contextual features, and a continuous input alone does not give that. LDLM paper §6
Which layer. The hidden states are taken from GPT-2’s third-to-last layer. The paper gives no reason for that layer and shows the choice is not critical: every layer it tried works, the third-to-last has the best Mauve and the last the lowest generative perplexity, since the trainable encoder can reshape different hidden states. LDLM paper Appendix B.1
They are normalized to zero mean and unit variance with dataset statistics; the paper states this without a reason. LDLM paper §6
- in
- [B batch, T positions]
- out
- [B batch, T positions, 768 d]
- parameters
- 124,441,344
Parameters
124,441,344 parameters per instance, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Frozen GPT-2 | 124,441,344 | |
| Per-instance total | 124,441,344 | |
Derived from the pinned implementation’s constructors and configuration. The checkpoint’s tensors have not been counted. 306,011,162 of these parameters are trained; frozen GPT-2 and the decoder’s mask vector, 124,442,112, are not.
Explanation
Nothing here is trained: GPT-2 is in evaluation mode and takes no gradient, and the statistics are measured once.
GPT-2, the latent encoder and the two decoders run under bfloat16 autocast. The Perceiver attention switches autocast off around its query and key normalization and its kernel; the queries, keys and values it hands them are already bfloat16.
The standardized features are both what the encoder reads and what the latent decoder is trained to reproduce.
Sources5
Direct evidence
With the feature model frozen, its hidden state is computed without gradient, standardized, and passed to the encoder; the same tensor is the target of the feature loss.
The implementation, pinned at revision e18c42acdfb7, not yet public
Direct evidence
The feature model is wrapped to return one entry of its hidden states, chosen by index; GPT-2’s token table is resized to the configured vocabulary.
The implementation, pinned at revision e18c42acdfb7, not yet public
Direct evidence
With an extraction layer other than −1 the encoder loads the wrapped feature model; when it is frozen it is put in evaluation mode and takes no gradient.
The implementation, pinned at revision e18c42acdfb7, not yet public
Direct evidence
The feature statistics are a mean and a standard deviation per feature over valid positions of the validation set, the variance clamped at 10⁻⁴; features are standardized with them.
The implementation, pinned at revision e18c42acdfb7, not yet public
Direct evidence
The feature model is GPT-2, frozen; the hidden state at index −3 is extracted; the encoder and the decoder are of the Perceiver kind; the file’s vocabulary size and mask token ID are defaults that the trainer overwrites with the tokenizer’s; there are as many latents as positions, of 768 features.
The implementation, pinned at revision e18c42acdfb7, not yet public