LLM Visual Decoder

This repository trains a lightweight bridge from the hidden states of every transformer block in a frozen causal LLM to the token-level conditioning space of a frozen Stable Diffusion 1.5 model. It is an interpretability experiment testing which visual semantics are decodable from LLM representations.

Stable Diffusion cross-attention outputs cannot be predicted from text alone: they also depend on the noisy image latent and diffusion timestep. The bridge therefore produces the text-side context consumed by every cross-attention block.

flowchart LR
    P["Prompt / context"] --> L["Frozen Qwen3-4B-Base"]
    L --> H["All block × token hidden states"]
    H --> N["Per-layer feature normalization"]
    N --> M["Learned softmax layer mixture"]
    M --> R["Trainable Resampler"]
    R --> C["77 × 768 conditioning tokens"]
    C --> U["Frozen SD 1.5 U-Net"]
    U --> I["Image distribution"]

Loading

What is implemented

• Frozen causal-LLM hidden-state extraction from every transformer block.
• Feature normalization followed by a learned, globally shared softmax weighting over layers. The embedding output can optionally be included.
• A token-preserving Transformer Resampler producing SD 1.5-compatible 77 × 768 conditioning.
• Mean-pooling and last-token MLP ablations.
• Stage 1, align: distillation to the frozen CLIP text-encoder representation.
• Stage 2, diffusion: frozen-U-Net denoising loss on caption-image pairs while retaining the alignment loss.
• Checkpointed inference with bridge-vs-normal-SD generations using identical seeds.
• LLM continuation generation followed by a sliding token-window image trajectory, saved as PNG frames, a grid, an animated GIF, HTML, and JSON metadata.
• Held-out representation metrics and lightweight unit tests.

Linux setup

Recommended hardware is a single NVIDIA GPU with 24 GB VRAM. Python 3.10 or 3.11 is the safest choice.

git clone git@github.com:fumin0ri/llm-visual-decoder.git
cd llm-visual-decoder

python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip

# Install a CUDA build of PyTorch suitable for the NVIDIA driver, then this project.
pip install torch torchvision
pip install -e ".[memory,dev]"

accelerate config
pytest -q

If xformers is unavailable for the selected PyTorch/CUDA combination, install with pip install -e ".[dev]" and set enable_xformers: false in the YAML file.

The first run downloads Qwen3-4B-Base, Stable Diffusion 1.5, and the configured dataset from Hugging Face. Qwen3 requires Transformers 4.51 or newer; the project dependency constrains Transformers to a compatible 4.x release. The default dataset is the 5,000-image MS COCO Captions 2017 validation slice hosted by lmms-lab; each image has several human-written captions. The configured 90/10 split yields 4,500 training images and 500 validation images. Set HF_HOME to a disk with enough space if necessary. The loader requests only the val split, so the unrelated test split is not downloaded.

Quick start

Review configs/mvp.yaml, especially the dataset and output directory. Then run:

# Stage 1: only captions are needed.
accelerate launch -m llm_visual_decoder.train \
  --config configs/mvp.yaml \
  --stage align

Initialize Stage 2 from the final Stage 1 checkpoint:

accelerate launch -m llm_visual_decoder.train \
  --config configs/mvp.yaml \
  --stage diffusion \
  --init-checkpoint runs/qwen3-4b/align/checkpoint-00010000

Generate bridge-conditioned and ordinary SD images with the same random seeds:

llmvd-infer \
  --checkpoint runs/qwen3-4b/diffusion/checkpoint-00010000 \
  --prompt-file prompts/example.txt \
  --seeds 0 1 2 3 \
  --output-dir generated/mvp

Generate an LLM continuation, then slide an 8-token window over the prompt and answer. Every frame uses the same diffusion seed, so image differences are not caused by different initial noise:

llmvd-trajectory \
  --checkpoint runs/qwen3-4b/diffusion/checkpoint-00010000 \
  --prompt "The capital of Japan is Tokyo. The currency used in the country shaped like a boot is" \
  --max-new-tokens 12 \
  --window-size 8 \
  --stride 1 \
  --seed 42 \
  --output-dir generated/currency-trajectory

Open generated/currency-trajectory/index.html to inspect each token window beside its image. trajectory.gif animates the sequence, trajectory-grid.png shows all frames, and manifest.json records the generated continuation, token boundaries, learned layer weights, and generation settings.

Measure held-out context reconstruction:

llmvd-evaluate \
  --checkpoint runs/qwen3-4b/diffusion/checkpoint-00010000 \
  --output runs/qwen3-4b/evaluation.json

Dataset formats

Hugging Face dataset

The default configuration uses:

dataset:
  name: lmms-lab/COCO-Caption2017
  config_name: null
  train_split: val
  caption_column: answer
  image_column: image

The hosted MS COCO slice is a reproducible pipeline baseline. A complete COCO training split or a larger caption-image corpus is preferable for the final experiment. When a row contains multiple captions, one is sampled at each access, exposing the bridge to alternate descriptions over successive epochs.

Local ImageFolder

Set dataset.name: null and point train_data_dir at an ImageFolder directory. Hugging Face ImageFolder accepts a metadata.jsonl file such as:

{"file_name": "000001.jpg", "text": "a red cube above a blue sphere"}
{"file_name": "000002.jpg", "text": "a city after a storm"}

Keep the image files and metadata.jsonl in the same directory, and set caption_column: text and image_column: image.

Layer weighting and ablations

By default, output_hidden_states[0] (the embedding output) is excluded and all transformer block outputs are used. Each state is normalized across its feature dimension, then the bridge learns one softmax weight per block. Qwen3-4B-Base has 36 transformer blocks, so the default bridge learns 36 weights initialized uniformly. Evaluation and inference manifests export the learned weights.

Set bridge.include_embedding_layer: true to add the embedding output. This changes the checkpoint shape and requires retraining. Checkpoints from the earlier single-layer implementation are also incompatible with the new bridge.

Compare the token resampler with pooled ablations:

bridge:
  type: resampler  # mean_mlp or last_mlp

Useful controls include fixed uniform layer weights, one-hot layer weights, an untrained/random LLM, shuffled caption-state pairs, a token-ID-only encoder of comparable capacity, and ordinary SD prompting. A good image is not by itself evidence that the representation is causally used by the LLM.

For a capacity baseline using the previous model, run the same experiment with configs/pythia-1.4b.yaml. Qwen and Pythia use different hidden dimensions, layer counts, and tokenizers, so their bridge checkpoints are not interchangeable.

RTX 4090 memory notes

The base models are frozen. During Stage 2, gradients still pass through the U-Net to its conditioning input, so U-Net activations consume memory even though its weights are not updated. The defaults use:

• batch size 1 and gradient accumulation 8;
• BF16;
• U-Net gradient checkpointing;
• optional xFormers attention;
• Qwen3-4B-Base.

If memory is tight, lower resolution to 384 or 256 (multiples of 8), reduce bridge.num_layers, or use the included Pythia 1.4B baseline config. Do not change num_queries or output_dim for SD 1.5.

Interpretation boundary

The experiment supports the claim that visual-semantic information can be decoded from the learned layer mixture through a particular trained bridge. The sliding window trajectory probes how decodable conditioning changes across token positions; short windows may be outside the caption distribution seen during training. It does not prove that a generated image is a faithful mental image, or that decoded information is causally used by the LLM. Causal activation patching/steering and matched controls are needed for the stronger claim.

License

Code in this repository is MIT licensed. Model weights and datasets retain their own licenses and usage restrictions.