UAVTB

Uncertainty-Aware Adaptive Visual Token Budgeting experiments built from the local proposal in this folder and the DART pruning baseline.

This repository is set up to answer the first validation question from the proposal:

Can sample-wise adaptive visual token budgeting beat fixed budgets under the same average token cost?

What is in this repo

• research proposal.md and research idea plan.md: the research spec.
• vendor/DART: local checkout of the DART baseline codebase.
• src/uavtb: adaptive budget policies, feature extraction, oracle tooling, and DART integration helpers.
• scripts/run_adaptive_llava_eval.py: per-sample adaptive evaluation wrapper on top of DART's LLaVA codepath.
• scripts/export_budget_features.py: cheap feature export from (image, question) pairs.
• scripts/build_oracle_labels.py: convert fixed-budget sweep scores into oracle budget labels.
• scripts/train_budget_mlp.py: train a lightweight budget predictor from oracle labels.

Validation workflow

1. Install the local package:

pip install -e .

Run the evaluation and training scripts inside the same PyTorch / Transformers environment that you use for DART.

2. Make sure vendor/DART exists.

3. Export cheap features for a benchmark split:

python scripts/export_budget_features.py \
  --question-file /path/to/questions.jsonl \
  --image-folder /path/to/images \
  --output-file outputs/features.jsonl

4. Run a first adaptive baseline with heuristics:

python scripts/run_adaptive_llava_eval.py \
  --model-path /path/to/llava-checkpoint \
  --question-file /path/to/questions.jsonl \
  --image-folder /path/to/images \
  --answers-file outputs/answers.jsonl \
  --decision-log-file outputs/heuristic_budget_log.jsonl \
  --policy heuristic \
  --budget-bins 32,64,128,256

5. After scoring several fixed-budget sweeps, build oracle labels:

python scripts/build_oracle_labels.py \
  --sweep-file outputs/scored_sweeps.jsonl \
  --output-file outputs/oracle_labels.jsonl \
  --sample-id-field sample_id \
  --budget-field budget \
  --score-field score \
  --reference-budget 256 \
  --tolerance 0.0

6. Train a small learned predictor:

python scripts/train_budget_mlp.py \
  --features-file outputs/features.jsonl \
  --label-file outputs/oracle_labels.jsonl \
  --output-file checkpoints/budget_mlp.pt \
  --budget-bins 32,64,128,256

7. Run learned adaptive budgeting:

python scripts/run_adaptive_llava_eval.py \
  --model-path /path/to/llava-checkpoint \
  --question-file /path/to/questions.jsonl \
  --image-folder /path/to/images \
  --answers-file outputs/answers_learned.jsonl \
  --decision-log-file outputs/learned_budget_log.jsonl \
  --policy learned \
  --policy-checkpoint checkpoints/budget_mlp.pt \
  --budget-bins 32,64,128,256

Notes

• In the current DART integration, the adaptive budget is mapped to DART's max_num_trunction field. That is the smallest change that enables per-sample budgets without modifying DART's pruning logic.
• This repo is already initialized as a local git repository. It is ready to push to GitHub once you decide the remote name.