ECF: Entropy-Based Candidate Filter for Multimodal RAG

ECF is a training-free evidence-admission framework for retrieval-augmented multimodal diffusion language models (DLMs). A retriever supplies a ranked candidate pool, while the DLM decides whether an additional page makes its initial answer-denoising distribution more confident than a geometry-matched blank control. The final generator receives top-1 and at most one admitted companion page.

This repository provides one evaluation pipeline for three masked multimodal DLMs and five visual question-answering datasets. Model-specific code is confined to backend adapters; retrieval, slicing, gating, logging, and answer scoring are shared.

Method

For top-1 evidence c1 and rank-j candidate cj, ECF computes:

H1j   = H(c1 + cj)
H1b_j = H(c1 + blank_like(cj))
Gj    = H1b_j - H1j

H is weighted mean entropy over the first answer block, with weight exp(-0.5 i) for answer-mask position i. Lower entropy indicates a more committed first-step denoising distribution.

k = 1: use c1
k >= 2:
  if G2 <= 0: use c1
  otherwise: use c1 + c2
k >= 3:
  replace c2 with the largest-gain cj only if Gj > G2 and H1j < H1

The rank-2 decision is the admission anchor. A failed rank-2 probe blocks lower-ranked pages. This policy is shared by all model backends.

Supported Configurations

Component	Values
Method	ecf, fixed_topk
Model	llada2uni, llada_v, dream_vl
Dataset	chartqa, infochartqa, docvqa, infovqa, tatdqa
Slicing	none, pre_retrieval, post_retrieval

fixed_topk --k K evaluates every fixed input size from top-1 through top-K. ecf-eval --method ecf --k K evaluates one dynamic candidate pool of size K.

Repository Layout

ecf/
  backends/                 model-specific official inference adapters
    llada2_uni.py
    llada_v.py
    dream_vl.py
    llada2/                 namespaced LLaDA2 tokenization/denoising helpers
  data/
    loader.py               unified parquet schema and prompt construction
    retrieval.py            VisRAG encoding and retrieval replay
    evidence.py             page/fragment evidence units
    slicing.py              standalone fragment-cache command
  evaluation/
    ecf.py                  model-independent dynamic gate runner
    fixed_topk.py           model-independent fixed top-k runner
    common.py               logging, resume, retrieval and scoring fields
  metrics/answers.py        scorer used by every experiment
  policy/ecf.py             pure dynamic top-k selection policy
  cli.py                    unified public CLI
evaluate/
  ecf/                      five dataset launchers
  fixed_topk/               five dataset launchers
scripts/                    setup, downloads, and dataset preparation
tests/                      policy, scoring, runner, and hygiene tests

Installation

The setup script creates a Conda environment named ecf, installs PyTorch for the detected CUDA version, checks out the required model repositories, and downloads:

• inclusionAI/LLaDA2.0-Uni
• GSAI-ML/LLaDA-V
• Dream-org/Dream-VL-7B
• openbmb/VisRAG-Ret

Set ECF_MODELS_ROOT, ECF_DATA_ROOT, or ECF_THIRD_PARTY_ROOT to put large artifacts outside the repository.

Configure a dataset source, then run the setup script:

export ECF_DATASET_REPO=organization/ECF-Ret-Test-2500
bash scripts/setup.sh
conda activate ecf

To prepare datasets from an existing local export instead, set ECF_AUTHORIZED_DATA_ROOT to the directory containing the five datasets.

Data Format

Each dataset uses the following Parquet layout:

Dataset-Ret-Test-2500/
  corpus/train-*.parquet
    corpus-id, image={bytes,path}
  queries/train-00000-of-00001.parquet
    query-id, query, answer
  qrels/train-00000-of-00001.parquet
    query-id, corpus-id, score

The dataset preparation script also accepts records.jsonl with query_id, question, answer, corpus_id, and image_path. It deterministically selects 2,500 examples for each evaluation set.

Evaluation

Dataset launchers expose MODEL, K, DEVICE, and OUTPUT_DIR:

MODEL=llada2uni K=3 DEVICE=cuda:0 bash evaluate/ecf/chartqa.sh
MODEL=llada_v K=3 DEVICE=cuda:0 bash evaluate/fixed_topk/infochartqa.sh
MODEL=dream_vl K=5 DEVICE=cuda:1 bash evaluate/ecf/infovqa.sh

The equivalent unified CLI is:

ecf-eval \
  --method ecf \
  --model llada2uni \
  --dataset chartqa \
  --k 3 \
  --device cuda:0 \
  --output-dir outputs/ecf/llada2uni/chartqa \
  --resume

Important overrides include:

--dataset-dir PATH
--model-path PATH_OR_HF_ID
--model-source PATH
--retriever-path PATH_OR_HF_ID
--slicing auto|none|pre_retrieval|post_retrieval
--fragment-cache PATH
--retrieval-replay CSV
--answer-length N
--steps N
--temperature FLOAT
--max-samples N

When ECF slicing is auto, DocVQA, InfoVQA, and TATDQA use pre-retrieval slicing; chart datasets do not. Fixed top-k defaults to unsliced full pages. If slicing is enabled without --fragment-cache, a cache is created under the dataset directory using the default settings: at most three fragments, minimum area ratio 0.03, minimum side 96 px, padding 8 px, with the full page retained.

ecf-build-cache --dataset-dir data/InfoVQA-Ret-Test-2500

Retrieval can be frozen and replayed across model backends:

ecf-eval --method fixed_topk --model llada2uni --dataset chartqa --k 3 \
  --retrieval-only --retrieval-output outputs/chartqa_top3.csv

ecf-eval --method ecf --model dream_vl --dataset chartqa --k 3 \
  --retrieval-replay outputs/chartqa_top3.csv

Outputs

Each evaluation writes per-example CSV results, a JSON summary, run_config.json, a retrieval replay CSV, and run.log. With --resume, ECF continues from an existing CSV after validating its schema and query order.

Development

python -m pytest -q
python -m compileall -q ecf tests