Natural Language Autoencoder

Backbone: T5-small  |  Latent dim: 128  |  Dataset: WikiText-2  |  New params: ~2.5M

VAE for text built on a frozen T5-small backbone. Text goes in, gets compressed to a 128-d vector, then reconstructed. The continuous latent space means you can measure semantic similarity, round-trip reconstruct sentences, smoothly blend two pieces of text, or generate new text by sampling.

---

Contents

• Architecture
• How Training Works
• Dataset
• Hyperparameters
• What you can do with it
• Training Metrics
• Quickstart
• Project Layout
• Why it's built this way
• Limitations and Future Work

---

Architecture

flowchart TB
    X(["Input text"]) --> E1

    subgraph ENC ["🔒 Encoder — frozen"]
        direction TB
        E1["T5-small Encoder<br/><code>60M params</code>"]
        E2["hidden states<br/><code>[B, L, 512]</code>"]
        E3["Masked Mean Pool<br/><code>← attention_mask</code>"]
        E4["pooled<br/><code>[B, 512]</code>"]
        E1 --> E2 --> E3 --> E4
    end

    E4 --> MU["fc_mu<br/><code>Linear(512→128)</code>"]
    E4 --> LV["fc_logvar<br/><code>Linear(512→128)</code>"]
    MU --> RP
    LV --> RP

    RP{{"Reparameterize<br/><code>z = μ + ε·σ</code> (train)<br/><code>z = μ</code> (eval)"}}
    RP --> Z(["z &nbsp;<code>[B, 128]</code>"])

    Z --> PROJ["latent_proj<br/><code>Linear(128→512)</code>"]
    PROJ --> MEM["Broadcast<br/><code>[B, L, 512]</code><br/><i>constant cross-attn memory</i>"]

    subgraph DEC ["🔥 Decoder — fine-tuned · lr 5e-5"]
        direction TB
        D1["T5-small Decoder"]
        D2["Shared LM Head<br/><i>weight-tied</i>"]
        D1 --> D2
    end

    MEM -. cross-attention .-> D1
    D2 --> OUT(["logits<br/><code>[B, L, vocab]</code>"])

    classDef frozen   fill:#e0e7ff,stroke:#6366f1,stroke-width:1px,color:#1e293b;
    classDef newp     fill:#d1fae5,stroke:#10b981,stroke-width:1px,color:#064e3b;
    classDef tuned    fill:#fef3c7,stroke:#f59e0b,stroke-width:1px,color:#78350f;
    classDef op       fill:#fae8ff,stroke:#a855f7,stroke-width:1px,color:#581c87;
    classDef io       fill:#f1f5f9,stroke:#64748b,stroke-width:1px,color:#0f172a;

    class E1,E2,E3,E4 frozen;
    class MU,LV,PROJ newp;
    class D1,D2 tuned;
    class RP,MEM op;
    class X,Z,OUT io;

Loading

🔒 frozen (≈60M, no grad)  ·  🟢 new params (fc_mu, fc_logvar, latent_proj — ~2.5M trained)  ·  🔥 fine-tuned T5 decoder + LM head (lr 5e-5)

Only ~2.5M new params are trained (the two FC heads and projection layer). The encoder is fully frozen; the decoder fine-tunes at a low learning rate to preserve what T5 already knows.

Component	Parameters	Trainable
T5-small encoder	~60M	No (frozen)
T5-small decoder	~60M	Yes (lr=5e-5)
Shared LM head (weight-tied)	~65M	Yes (lr=5e-5)
FC heads (μ, log σ²)	~131K	Yes (lr=3e-4)
Projection Linear(128→512)	~66K	Yes (lr=3e-4)
Total	~252M	~2.5M new

---

How Training Works

Loss = reconstruction error + KL penalty (ELBO). The KL term measures how far the encoder's distribution is from a standard Gaussian — keeping the latent space well-organized so you can sample from it later.

The KL penalty ramps up from 0 to 1 over the first 10,000 steps. Starting at 0 lets the model learn decent reconstructions first. If you apply full KL pressure from step 1, the encoder collapses early — it outputs a prior-matched distribution regardless of input, and reconstruction falls back on the decoder's language model priors instead of z.

A per-dimension KL floor of 0.5 nats keeps all 128 dimensions active. Without it, some dimensions drift to exactly match the prior and carry no information — wasted capacity.

Optimizer

AdamW with separate learning rates for new vs. pretrained layers:

Parameter group	Learning rate	Weight decay
New layers (FC heads, projection)	3e-4	0.01
Pretrained T5 (encoder + decoder)	5e-5	0.01

• Warmup: 500 steps (linear 0 → LR)
• Decay: linear from LR → ~0 over remaining steps (capped at max_steps=50,000)
• Gradient clipping: max norm 1.0

Warning signs

After step 10,000, if kl_per_dim < 0.1 nats the trainer emits a warning — the encoder isn't using the latent space effectively.

---

Dataset

Source: wikitext / wikitext-2-raw-v1 via HuggingFace datasets

Splits used: train, validation

Preprocessing

1. Filter — drop empty lines, section headers (lines starting with =), and lines under 20 characters
2. Tokenize — T5Tokenizer, max_length=64, padding="max_length", truncation=True
3. Decoder input — shift-right: prepend pad token, drop last token (standard T5 teacher-forcing)
4. Labels — clone of input_ids; padding positions set to -100 (ignored by CrossEntropy)

---

Hyperparameters

Full config in config.py:

Parameter	Value	Description
model_name	google-t5/t5-small	HuggingFace backbone
latent_dim	128	Dimensionality of z
max_length	64	Max token sequence length
batch_size	32	Training batch size
learning_rate	3e-4	LR for new layers
t5_learning_rate	5e-5	LR for pretrained T5
weight_decay	0.01	AdamW weight decay
num_epochs	20	Maximum training epochs
max_steps	50,000	Hard step cap
warmup_steps	500	LR warmup duration
kl_annealing_end	10,000	Step at which KL weight reaches 1.0
kl_weight_max	1.0	Maximum KL weight
free_bits	0.5	Min KL per latent dim (nats)
grad_clip	1.0	Gradient clipping max norm
save_every_n_steps	1,000	Checkpoint frequency
log_every_n_steps	100	Metrics logging frequency
early_stopping_patience	3	Epochs without val improvement
dataset_name	wikitext	HuggingFace dataset name
dataset_config	wikitext-2-raw-v1	Dataset configuration
num_workers	0	DataLoader workers (set 4 on GPU)

---

What you can do with it

All inference functions are in inference/utils.py. At eval time the model runs deterministically (z = μ, no noise).

encode_text(text) → np.ndarray[128]

Returns the 128-d latent vector for a string. Inputs over 64 tokens get truncated with a warning. Use this for semantic similarity — cosine distance between two vectors tells you how semantically close they are.

reconstruct(text) → str

Full round-trip: encode → z → decode. Shows how much information survives the 128-d bottleneck.

decode_from_z(z) → str

Greedy decoding from a raw latent vector. Applies a repetition penalty to reduce looping. Returns [INCOMPLETE] if EOS isn't reached within 64 tokens.

interpolate(text_a, text_b, steps=10) → list[tuple[float, str]]

Encodes both texts, then walks from z_a to z_b in steps steps, decoding each point. Returns (α, text) pairs for α ∈ [0, 1]. Interpolates along the surface of the latent space (not through the middle), so intermediate points stay in-distribution. Near-parallel and antipodal edge cases are handled.

sample_from_prior(n=1, temp=1.0) → list[str]

Samples random latent vectors and decodes them. Temperature < 1 produces conservative outputs; > 1 produces more varied but less coherent text. Capped at 2.0.

---

Training Metrics

Computed each epoch with full KL weight (no annealing):

Metric	Description
val_loss	ELBO = val_recon + val_kl
val_recon	Cross-entropy reconstruction loss
val_kl	KL divergence (free bits applied)
val_perplexity	exp(val_recon), clipped at 100
val_token_acc	Fraction of non-padding tokens predicted correctly

Early stopping monitors val_loss with patience=3. Best checkpoint saved to checkpoints/best_epoch_N.pt.

---

Quickstart

Requirements

pip install -r requirements.txt

Key dependencies: torch==2.3.1, transformers==4.44.2, datasets==2.21.0, umap-learn==0.5.6

Device auto-detected in order: CUDA → MPS → CPU.

Apple Silicon note: some PyTorch versions crash on the MPS backend with a Metal assertion (MPSNDArray ... buffer is not large enough) when running T5. If this happens, force CPU with --device cpu.

Train

python train.py \
  --model_name google-t5/t5-small \
  --latent_dim 128 \
  --batch_size 32 \
  --num_epochs 20 \
  --max_steps 50000 \
  --kl_annealing_end 10000 \
  --free_bits 0.5

Outputs:

• checkpoints/best_epoch_N.pt — best validation loss checkpoint (full state: weights, optimizer, scheduler, RNG)
• checkpoints/step_XXXX.pt — periodic saves every 1,000 steps (resumable)
• outputs/training_log.json — metrics logged every 100 steps

Demo

python demo.py --checkpoint checkpoints/best_epoch_N.pt

Runs all four capabilities (cosine similarity, reconstruction, interpolation, prior sampling) on default text pairs.

Custom texts:

python demo.py --checkpoint checkpoints/best_epoch_N.pt \
  --text_a "Scientists discovered a new species of deep-sea fish." \
  --text_b "The stock market fell sharply on Tuesday."

Notebook

jupyter notebook notebooks/01_nl_vae_walkthrough.ipynb

Step-by-step walkthrough of architecture, training, and inference.

---

Project Layout

config.py                   VAEConfig dataclass — all hyperparameters
train.py                    Training entry point; parses args, builds trainer
demo.py                     Inference demo; runs four capability checks

model/
  vae.py                    NLVAE: top-level model combining encoder + decoder
  encoder.py                VAEEncoder: T5 encoder + masked mean pool + FC heads
  decoder.py                VAEDecoder: projection + T5 decoder + LM head

training/
  trainer.py                Training loop, validation, checkpointing, early stopping
  loss.py                   ELBO loss: CrossEntropy recon + KL with free bits + annealing

data/
  dataset.py                WikiText-2 loading, filtering, tokenization, DataLoader

inference/
  utils.py                  encode_text, decode_from_z, reconstruct, interpolate, sample_from_prior

visualization/
  plots.py                  UMAP latent space, training curves, interpolation visualization

notebooks/
  01_nl_vae_walkthrough.ipynb   Educational walkthrough

checkpoints/                Saved model states (created at runtime)
outputs/                    Training logs and plots (created at runtime)

---

Why it's built this way

Frozen encoder

T5's encoder already produces rich 512-d contextual representations. Freezing it means the bottleneck learns to compress those representations, not learn to read text from scratch — cutting trainable params from ~252M to ~2.5M and training time substantially.

Constant cross-attention memory

A normal seq2seq decoder attends to varied, position-specific encoder states. Here, z is broadcast as a single constant sequence — every decoder position sees the same memory. This is intentional: if the decoder had rich position-varying context, it'd ignore z entirely and route around the bottleneck. Constant memory forces the decoder to rely on z for all content.

Slerp over linear interpolation

Straight-line interpolation between two high-dimensional vectors passes through the interior of the space, producing points with smaller norms than either endpoint. Since the model trains on vectors drawn from a Gaussian (which concentrate near a hypersphere in high dimensions), linear midpoints land in low-density regions. Slerp stays on the surface, keeping intermediate points in-distribution.

Free bits

Without a per-dimension floor, the KL term drives some latent dimensions to exactly match the prior — they stop carrying information about the input and become dead weight. The 0.5 nat floor removes gradient pressure to collapse those dimensions further, keeping all 128 active.

---

Limitations and Future Work

Current limitations:

• Greedy decoding only. No beam search, top-k, or nucleus sampling. Repetition is partially mitigated by a penalty but can still occur, especially far from the training distribution.
• Short sequences. Max 64 tokens. Longer inputs are truncated; long-range structure isn't representable.
• Single domain. Trained on WikiText-2 (encyclopedic English). Other registers (dialogue, code, social media) are untested.
• No quantitative evaluation. sacrebleu is installed but not wired up. BLEU, BERTScore, and similarity benchmarks aren't reported.
• Fixed prior. p(z) = N(0, I) is assumed, not learned. A flow-based or mixture prior could better match the aggregate posterior.
• No disentanglement. Dimensions aren't encouraged to be independent or semantically interpretable.

Ideas to extend:

• β-VAE: increase β > 1 to push disentanglement (trades reconstruction quality)
• VQ-VAE: discrete codebook instead of continuous bottleneck — lossless round-trips and discrete generation compatibility
• Conditional generation: condition the decoder on a class label (topic, sentiment) alongside z
• Learned prior: normalizing flow to better model the aggregate posterior
• Better decoding: beam search or nucleus sampling to cut repetition
• Longer contexts: learned compression token or hierarchical encoder for documents