OpenVLA-OFT Early Exit: Decoder-at-Layer Study (LIBERO-Spatial)
Trained readouts and complete evaluation artifacts for a depth-compression study of
OpenVLA-OFT 7B on LIBERO-Spatial. Core question: is the final LLM layer's advantage
for action decoding better features or feature–decoder compatibility?
Answer: compatibility. Retraining the action head at intermediate layers recovers
near-full closed-loop performance from half the network depth, and tuning a few
attention blocks recovers full performance from layer 12 of 32.
Retrained L1 action head (151M params, fresh init per layer, LLM frozen, trained on
20 episodes/task):
Attach layer
1
2
6
12
16
20
24
28
30
31
32
Success /100
11
24
22
77
97
94
97
99
100
99
100
Reference: full model with original head = 98/100 (matches the published result).
Additional arms:
Configuration
Layer 12
Layer 16
Layer 30
8.4M adapter + frozen original head
77
92
98
8.4M adapter + trained head
81
96
99
Attach@12 + k tuned LLM blocks (k=1/2/4)
93 / 95 / 98
—
—
Key findings:
No cliff at the final layer — the prior diagnostic observation that only layers 31–32 work
with the pretrained head is a head-compatibility artifact, fully repaired by
retraining the head (or even an 8.4M adapter in front of the frozen original head).
The information boundary sits between layers 12 and 16; below it, no static
readout of frozen features passes ~80/100 — but tuning as little as one attention
block below layer 12 jumps success to 93, and four blocks reach baseline parity (98).
Offline action-prediction metrics do not predict closed-loop success (e.g. the
unfreeze arm moves offline L1 by 7% while rollout success moves 77 → 98).
Success vs attach layer: retrained heads, adapter arm, and tuned-block arm; CKA structure below
heads/ head_layer{n}.pt (fresh-init L1RegressionActionHead per layer)
+ per-sample eval predictions head_pred_eval_layer{n}.npy
adapters/ adapter{A|B}_layer{n}.pt (A: pairs with frozen original head)
unfreeze/ unfreeze_k{1,2,4}.pt (tuned blocks 12-k..11 + head, attach@12)
rollout_*.json per-episode rollout records (all 20 configurations)
rollout_*.queries.npz per-query executed vs stock action chunks (drift analysis)
*_results.json offline metrics (heads, adapters, unfreeze, CKA, sweep summary)
offline_check/ fresh offline sweep used for stack validation
hidden/ extraction metadata + normalized GT labels (raw caches excluded, see below)
code_snapshot/ all analysis code (training, rollouts, verification, figure, this study)
env/ setup_env.sh + pip_freeze.txt + conda_env.yml (exact environment)
logs/ complete run logs, from driver install to final rollout
verify: code_snapshot/verify_results.py
Verifying every reported number (no GPU needed)
bash
1pip install numpy # the only requirement for this check2python code_snapshot/verify_results.py --results-dir . --ref-dir baseline_reference --skip-cka
3# expected final line: ALL CHECKS PASSED (exit code 0)
(The min-model provenance check below additionally needs torch, safetensors, and —
for --check-base — huggingface_hub. Both commands were tested on a fresh clone of
this repo exactly as written.)
Fully self-contained: dataset_statistics.json and the baseline-study reference artifacts
(baseline_reference/) are bundled. See REPRODUCING.md for the
three verification tiers (recompute-only → re-rollout with these checkpoints → full
re-derivation from public sources).
Reproducing / regenerating
Environment: env/setup_env.sh (mirrors the exact build; pins in env/pip_freeze.txt).
The raw hidden-state caches (~60 GB) are excluded; they regenerate deterministically
with code_snapshot/extract_hidden2.py and code_snapshot/extract_fullseq.py
(single A100 pass each; built-in gates verify correctness — the study machine
reproduced the base model's predictions bitwise).
min_model/ — the minimum-parameter deployment bundle
min_model/model.safetensors is the assembled exit-at-layer-12 model:
vision encoders + projector + embeddings + LLM blocks 0–11 (blocks 8–11 are the
k=4 tuned copies) + proprio projector + L1 action head. Blocks 12–31, the final
RMSNorm, and the LM head are removed.
Params
Size
LIBERO-Spatial success
Full OpenVLA-OFT stack
7.709 B
~15.5 GB
98/100
min_model (exit@12 + k=4)
3.772 B (48.9%)
7.9 GB
98/100
Forward pass: standard OpenVLA-OFT pipeline, stopping at block 11 and decoding the
action-token hidden states with action_head (bidirectional LLM attention — see
code_snapshot/unfreeze_common.py; the parity rollout used exactly these weights via
code_snapshot/rollout_unfreeze.py, record: rollout_unfreeze_k4.json).
Provenance is fully verifiable: min_model/manifest.json records, per tensor, its
SHA256 and source — either a named tensor of the public base checkpoint or a key of
unfreeze/unfreeze_k4.pt (whose offline predictions and 98/100 rollout are covered by
the main verify script). Check it:
bash
1python code_snapshot/verify_min_model.py --bundle min_model --unfreeze-ckpt unfreeze/unfreeze_k4.pt
2# expected: MIN-MODEL BUNDLE VERIFIED (exit 0); add --check-base to also3# byte-compare every frozen tensor against a fresh download of the base repo
Note: OpenVLA-OFT uses bidirectional attention in the LLM (not causal). If you
re-run decoder blocks standalone (as the unfreeze checkpoints require), follow
code_snapshot/unfreeze_common.py.
License and attribution
This release is licensed under the Llama 2 Community License (see LICENSE):
the OpenVLA backbone derives from Llama-2-7B, and unfreeze/unfreeze_k*.pt contain
modified copies of Llama-2-derived transformer blocks.
