How kitft/natural_language_autoencoders works
NLA trains paired activation verbalizer (AV) and reconstructor (AR) language models: data generation turns residual-stream vectors into explanation and reconstruction datasets, SFT initializes the models, and GRPO RL improves AV generations using AR reconstruction error. During rollout, the trainer injects activation embeddings and sends them to SGLang; training and scoring are integrated through Miles. nla_inference.py provides a separate checkpoint-based generation and AR-scoring path. The README documents the broader training capabilities; sampled excerpts confirm selected orchestration, rollout, training, storage, and inference behavior, so uncertain internal wiring is omitted.
Main entry points: storage.py Parquet storage · run_pipeline.py Data pipeline · nla_inference.py NLA client · train_actor.py FSDP actor
Components 21
Data generation
- Activation extraction
nla/datagen/stage0_extract.py - Dataset splitting
nla/datagen/stage1_split.py - API explanations
nla/datagen/stage2_api_explain.py - Dataset assembly
nla/datagen/stage3_build.py - Data pipeline
nla/datagen/run_pipeline.py
Training
- Training data source
nla/data_source.py - SFT rollout
nla/rollout/sft_actor.py - FSDP actor
nla/train_actor.py - Megatron actor
nla/megatron/train_actor.py - NLA generation
nla/rollout/nla_generate.py - Reconstruction reward
nla/reward.py - Training losses
nla/loss.py
Inference
- AV / AR checkpoint
- NLA client
nla_inference.py - AR scoring
nla_inference.py
External services
- Text corpus
- Miles trainer
- SGLang serving
Shared contracts
- Parquet storage
nla/storage.py - Metadata schema
nla/schema.py
Other components
- Researcher
Connections 30
- Researcher → Data pipeline: configures
- Text corpus → Activation extraction: provides text
- Data pipeline → Activation extraction: runs stage
- Data pipeline → Dataset splitting: runs stage
- Data pipeline → API explanations: runs stage
- Data pipeline → Dataset assembly: runs stage
- Activation extraction → Parquet storage: writes activations
- Dataset splitting → Parquet storage: writes splits
- API explanations → Parquet storage: writes explanations
- Dataset assembly → Parquet storage: writes datasets
- Parquet storage → Training data source: reads datasets
- Metadata schema → Parquet storage: defines sidecars
- Training data source → Miles trainer: supplies batches
- Miles trainer → FSDP actor: trains actor
- Miles trainer → Megatron actor: trains actor
- Miles trainer → SFT rollout: uses rollout
- Miles trainer → NLA generation: dispatches generation
- NLA generation → SGLang serving: sends embeddings
- Miles trainer → Reconstruction reward: computes reward
- FSDP actor → Training losses: uses loss
- Megatron actor → Training losses: uses loss
- FSDP actor → AV / AR checkpoint: saves models
- Megatron actor → AV / AR checkpoint: saves models
- Researcher → NLA client: submits vector
- NLA client → AV / AR checkpoint: loads AV
- NLA client → SGLang serving: requests generation
- SGLang serving → NLA client: returns text
- Researcher → AR scoring: submits vector and text
- AR scoring → AV / AR checkpoint: loads AR
- AR scoring → Researcher: returns score