ernie-research/AVIOT

Aggregating Visual Information with Optimal Transport for VideoLM Token Compression

Abstract

Video language models process videos as dense visual-token sequences with substantial representational redundancy. Compressing these sequences is therefore essential for reducing the visual-token burden on language-model decoding. The central challenge is to preserve visual information dispersed across frames under such compression. To this end, we introduce Aggregating Visual Information with Optimal Transport (AVIOT), which casts video token compression as transporting a dense empirical measure of frame observations onto a compact target measure. The resulting source-to-target coupling induces a distribution over source observations for each target support, directly specifying how the compressed video representation is constructed. We further adapt this construction along task and spatial axes. Question conditioning modulates the transport cost between source frames and target supports, while influencing how many supports are allocated to each temporal segment, thereby directing representation capacity toward question-relevant content. At multiple spatial granularities, AVIOT computes region-specific temporal transport plans and adaptively fuses the representations they yield, allowing different regions within the same compact representation to draw from different moments. Evaluations across varying compression ratios show that AVIOT matches or outperforms the uncompressed baseline on multiple video-understanding benchmarks while retaining strong performance at higher compression ratios.

Model Package

AVIOT is a video-language model for video question answering with inference-time visual-token compression. This release is a self-contained inference checkpoint: its four safetensors shards include the language model, visual encoder, multimodal projector, and AVIOT modules.

The public AVIOT code defines the custom architecture and the video input pipeline. Install it before loading this checkpoint; generic AutoModel.from_pretrained() loading is not supported by this release.

Installation

AVIOT requires Python 3.10 or newer, CUDA-enabled PyTorch, and Transformers 4.56 or newer. Clone and install the public code repository:

git clone https://github.com/ernie-research/AVIOT.git
cd AVIOT
python -m pip install -e '.[flash]'

The flash extra installs FlashAttention 2, which is used by the default inference command. Install the package without that extra and pass --attn-implementation sdpa if FlashAttention 2 is unavailable.

Loading

Load the checkpoint directly from its Hugging Face repository. The loader downloads and resolves the complete Hub snapshot automatically:

from aviot.inference import AVIOTGenerator

generator = AVIOTGenerator(
    "ernie-research/AVIOT",
    device="cuda",
    torch_dtype="bfloat16",
    attn_implementation="flash_attention_2",
)

prediction = generator.answer(
    "/path/to/video.mp4",
    "What happens in the video?",
    ratio=4,
)
print(prediction.text)

The same loader also accepts a downloaded local directory:

from aviot.inference import AVIOTGenerator

generator = AVIOTGenerator(
    "/path/to/AVIOT-model",
    device="cuda",
    torch_dtype="bfloat16",
    attn_implementation="flash_attention_2",
    local_files_only=True,
)

The command-line interface accepts either the Hub repository ID or a local checkpoint directory:

aviot-infer \
  --checkpoint ernie-research/AVIOT \
  --video /path/to/video.mp4 \
  --question 'What happens in the video?' \
  --ratio 4

For a local checkpoint, pass its directory to --checkpoint and add --local-files-only.

Inference Configuration

The compression ratio controls the target cardinality without requiring a different checkpoint. For F sampled frames and compression ratio r, AVIOT constructs ceil(F / r) compact supports (capped at F). The released training configuration samples ratios from 2 to 10, while --ratio can be selected directly at inference time.

The default generation configuration is deterministic (do_sample=false). The CLI likewise defaults to --temperature 0; sampling can be enabled explicitly by passing a positive temperature.

Video Input

The default video pipeline samples at 2 FPS and retains at most 224 frames. Frames are resized to 384 by 384 pixels and normalized with the preprocessing configuration bundled in vision_tower/. The physical frame timestamps are included in the model prompt.

Citation

If you find AVIOT useful in your research, please cite our paper:

@misc{yin2026aggregatingvisualinformationoptimal,
      title={Aggregating Visual Information with Optimal Transport for VideoLM Token Compression},
      author={Wenti Yin and Xiaotian Han and Junyuan Shang and Yuchen Ding and Shuohuan Wang and Dianhai Yu and Changxin Gao and Nong Sang},
      year={2026},
      eprint={2608.20473},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2608.20473},
}

License

The AVIOT code and this separately distributed checkpoint are released under the Apache License 2.0. See NOTICE for attribution. Datasets remain subject to their respective licenses and terms.