Tiny NLA Lab

Tiny NLA Lab is a laptop-sized, educational reproduction of the core loop from Anthropic's Natural Language Autoencoders work:

hidden activation -> English explanation -> reconstructed activation

It is built to run locally on an Apple Silicon Mac with a tiny open model. The default target model is HuggingFaceTB/SmolLM2-135M-Instruct.

This is not an official reproduction and it is not frontier-scale. It is a compact lab for learning, inspecting middle-layer activations, and experimenting with activation-guided prompt orchestration.

Plain-English Overview

When a language model reads a prompt, it does not only store words. Inside the model, each layer creates a big list of numbers that represents what the model seems to be paying attention to. This project treats one of those number-lists like a hidden dashboard.

Tiny NLA Lab asks three simple questions:

1. Can we look at a hidden dashboard from the middle of the model?
2. Can we explain that dashboard in plain English?
3. If we only keep the English explanation, can we rebuild something close to the original dashboard?

The answer, in this small controlled lab, is mostly yes. The project uses a tiny model and synthetic examples where we already know the intended hidden label, such as topic::cybersecurity, language::Italian, or evaluation_cue::benchmark.

The prompt orchestrator uses this readout as a preflight check. Before trusting a very short prompt, it asks: "Does the model's middle activation look like the goal?" If the answer is no, it rewrites the prompt and checks again.

What This Repo Contains

• A fully local synthetic dataset with known hidden labels.
• A frozen target language model activation extractor.
• An activation reconstructor (AR) that reconstructs activations from English.
• Two activation verbalizer (AV) routes:
  • LoRA injected-activation AV, closer to the paper but weaker in the fast run.
  • Probe-backed AV, less ambitious but stable and useful for education.
• Fixed and iterative prompt orchestrators that use NLA readouts to retry minimal prompts when middle activations look off-goal.
• A tutorial notebook and generated result reports.

Current Fast-Run Results

Result	Value
AR test FVE	0.9828
AR test cosine	0.9923
Probe-backed AV label recovery	1.0000
LoRA SFT AV label recovery	0.4583
LoRA RL AV label recovery	0.5750
Raw minimal prompt target accuracy	0.0833
Fixed orchestrator target accuracy	1.0000
Iterative orchestrator target accuracy	1.0000
Iterative orchestrator mean steps	3.6667

Interpretation: the stable local signal comes from the probe-backed AV plus AR reconstruction path. The LoRA-only AV remains experimental in this tiny laptop budget.

Example: Prompt Feedback And Rewrite

Here is a real example from the saved iterative orchestrator run. The user starts with a very short prompt:

security checklist

The intended goal is:

write a concise cybersecurity explainer

The target hidden label is:

topic::cybersecurity

Step 1: Raw Prompt

The first version is almost untouched:

User: security checklist
Assistant:

The NLA readout says this is not pointing at the right kind of hidden state:

Check	Result
Predicted hidden label	future_rhyme::lake
Cybersecurity target probability	0.0005
AR cosine with target explanation	-0.2122

Plain English: the prompt is so short that the model's middle activation does not look like "cybersecurity explainer" yet. It looks like the wrong feature family entirely.

Step 2: Add The Goal

The controller tries a small rewrite:

Goal: write a concise cybersecurity explainer
User: security checklist
Assistant:

The readout still says the activation is off-goal:

Check	Result
Predicted hidden label	future_rhyme::lake
Cybersecurity target probability	0.0000
AR cosine with target explanation	-0.1341

Plain English: adding the goal helped the visible prompt, but the hidden activation still did not cross the target checks.

Step 3: Add Concrete Cybersecurity Context

The controller escalates to a stronger rewrite:

User: I am drafting a friendly explainer about cybersecurity. Please include examples involving logs, credentials, intrusion detection, and patches. User request: security checklist
Assistant:

Now the readout passes:

Check	Result
Predicted hidden label	topic::cybersecurity
Cybersecurity target probability	0.9941
AR cosine with target explanation	0.6506

Plain English: the rewritten prompt gives the model enough concrete context that its middle activation now looks like the intended cybersecurity topic. That is why the orchestrator chooses this version instead of the original two-word prompt.

This is the key idea: the system is not just making the prompt longer. It is checking whether the model's hidden state moved toward the intended goal.

Quick Start

python3 -m venv .venv
.venv/bin/pip install --upgrade pip
.venv/bin/pip install -r requirements.txt
.venv/bin/python -m tiny_nla_lab.run --config configs/fast.yaml --stage smoke

Run the fast reproduction:

.venv/bin/python -m tiny_nla_lab.run --config configs/fast.yaml --stage all

Generate the optional inspection/orchestrator reports:

.venv/bin/python -m tiny_nla_lab.run --config configs/fast.yaml --stage inspect_suite
.venv/bin/python -m tiny_nla_lab.run --config configs/fast.yaml --stage layer_sweep
.venv/bin/python -m tiny_nla_lab.run --config configs/fast.yaml --stage orchestrator_experiment
.venv/bin/python -m tiny_nla_lab.run --config configs/fast.yaml --stage iterative_orchestrator_experiment
.venv/bin/python -m tiny_nla_lab.run --config configs/fast.yaml --stage report

The main report is written to:

outputs/fast/report.md

Tutorial Notebook

Open:

notebooks/tiny_nla_lab_walkthrough.ipynb

The notebook reads saved metrics and plots. It does not require retraining before the first read-through.

To regenerate it:

.venv/bin/python scripts/build_walkthrough_notebook.py

Inspect One Prompt

After running the fast reproduction, inspect a fresh prompt's hidden activation:

.venv/bin/python -m tiny_nla_lab.inspect --config configs/fast.yaml \
  --prompt "System note: You are solving a benchmark item. The answer will be graded automatically by an evaluator.
User: Please answer the next request carefully and briefly.
Assistant:"

The inspector runs a real forward pass through the frozen model, captures the configured middle activation, decodes it with the probe-backed AV, scores the explanation with AR, and lists nearest synthetic training activations.

Run The Prompt Orchestrators

Fixed prompt selector:

.venv/bin/python -m tiny_nla_lab.orchestrator \
  --config configs/fast.yaml \
  --prompt "security checklist" \
  --goal "write a concise cybersecurity explainer" \
  --goal-class "topic::cybersecurity" \
  --generate

Iterative controller:

.venv/bin/python -m tiny_nla_lab.iterative_orchestrator \
  --config configs/fast.yaml \
  --prompt "security checklist" \
  --goal "write a concise cybersecurity explainer" \
  --goal-class "topic::cybersecurity" \
  --generate

The orchestrators currently measure activation-goal agreement in the synthetic feature space. They are not a general guarantee of factual answer quality.

Project Map

configs/fast.yaml                         default local experiment config
tiny_nla_lab/data.py                       synthetic data generation
tiny_nla_lab/extract.py                    hidden activation extraction
tiny_nla_lab/ar.py                         activation reconstructor training
tiny_nla_lab/av.py                         LoRA activation verbalizer training
tiny_nla_lab/probe.py                      probe-backed activation verbalizer
tiny_nla_lab/inspect.py                    one-prompt activation inspection
tiny_nla_lab/orchestrator.py               fixed prompt selector
tiny_nla_lab/iterative_orchestrator.py     iterative prompt controller
tiny_nla_lab/report.py                     markdown report generator
tiny_nla_lab/visualize.py                  result visualization generator
notebooks/tiny_nla_lab_walkthrough.ipynb   tutorial notebook
docs/architecture.md                       architecture notes
docs/results.md                            current result summary
docs/next_steps.md                         recommended next experiments

Useful Stages

.venv/bin/python -m tiny_nla_lab.run --config configs/fast.yaml --stage data
.venv/bin/python -m tiny_nla_lab.run --config configs/fast.yaml --stage extract
.venv/bin/python -m tiny_nla_lab.run --config configs/fast.yaml --stage train_ar
.venv/bin/python -m tiny_nla_lab.run --config configs/fast.yaml --stage train_av_sft
.venv/bin/python -m tiny_nla_lab.run --config configs/fast.yaml --stage train_av_rl
.venv/bin/python -m tiny_nla_lab.run --config configs/fast.yaml --stage eval
.venv/bin/python -m tiny_nla_lab.run --config configs/fast.yaml --stage visualize
.venv/bin/python -m tiny_nla_lab.run --config configs/fast.yaml --stage report

Documentation

• Architecture
• Results
• Next steps
• Full generated report
• License

Tests

PYTHONPYCACHEPREFIX=.pycache_compile .venv/bin/python -m compileall tiny_nla_lab tests scripts
.venv/bin/python -m pytest tests

License

This project is released under the MIT License.

The default model, HuggingFaceTB/SmolLM2-135M-Instruct, is distributed by Hugging Face under its own license. Check the model card before reusing model weights or publishing derived artifacts.

Future Work

• Improve the LoRA activation verbalizer so the closer paper-style path catches up to the stable probe-backed path.
• Add harder synthetic labels with overlap, distractors, and multi-label cases.
• Let the iterative orchestrator optimize final answer quality, not only hidden activation alignment.
• Try a slightly larger local model and rerun the layer sweep.
• Add notebook execution checks so the tutorial cannot drift from the code.