Object Recognition as Next Token Prediction

arXiv | Colab | Documentation | Hugging Face

Top 30 predictions with probabilities from our model on the image of "The Legend of Zelda: Tears of the Kingdom" [^1].

Introduction

This is the official PyTorch implementation for the paper Object Recognition as Next Token Prediction accepted at CVPR 2024 (Highlight).

@inproceedings{nxtp,
  title     = {{Object Recognition as Next Token Prediction}},
  author    = {Kaiyu Yue and Bor-Chun Chen and Jonas Geiping and Hengduo Li and Tom Goldstein and Ser-Nam Lim},
  booktitle = {Computer Vision and Pattern Recognition Conference (CVPR)},
  year      = {2024}
}

Updates

May 26, 2024

• add ImageNet experiments: see src/imagenet
• visualize attention maps in decoder layers during inference: see examples

Mar 17, 2024

• release the best 1.78B model trained on G70M
• export onnx models: docs/onnx-export

Mar 03, 2024

• add examples with top-20 predictions to this readme
• add CLIP ViT- L/14 as the textual embedding model in evaluation metric (Table A.8 of the paper)

Method

This project delves into a fundamental problem in computer vision − object recognition − translating an image into object labels.

Linear models (such as ResNet) and contrastive models (such as CLIP) require predefined labels before inference, limiting their flexibility in real-world applications.

We extend W to cover the entire textual space using language models like LLaMA's 32K token embeddings. Our model predicts labels in a real-open manner through auto-regressive processing.

Additionally, our one-shot sampling technique enables efficient large-scale discriminative predictions, such as the top-100 labels.

The released models have 1.78B parameters. Truncating the model to 0.77B parameters still achieves competitive performance (Table 3 in the paper), which only has one transformer block in the decoder.

Examples

Models

The following table shows the reproduced results of recall (R column in Table 1 of the paper) on the validation splits with top-10 predictions.

Downloading

The checkpoints can be downloaded from the links in the table above. For downloading from Hugging Face, one option is to use git-lfs:

# install git lfs
git lfs install

# download the checkpoint in terminal
git clone https://huggingface.co/kaiyuyue/nxtp

Also, the checkpoint can be downloaded from the model page in the web browser.

Inference

There is an image assets/starbux.jpg for a quick test. First, please follow the instructions in Dependencies to prepare the environment.

To infer an image, please run

python src/infer.py \
  --ckpt-path path/to/model/checkpoint \
  --img-path assets/starbux.jpg \
  --num-labels 20

The output from model trained on G3M will be

top-20 predictions:
| prob: 0.05742 - coffee
| prob: 0.05525 - restaurant
| prob: 0.04402 - shop
| prob: 0.02528 - room
| prob: 0.02468 - store
| prob: 0.02381 - interior
| prob: 0.01732 - area
| prob: 0.01640 - building
| prob: 0.01616 - food
| prob: 0.01408 - bar
| prob: 0.01247 - customer
| prob: 0.01134 - view
| prob: 0.01059 - floor
| prob: 0.01045 - table
| prob: 0.00933 - kitchen
| prob: 0.00926 - home
| prob: 0.00872 - look
| prob: 0.00841 - people
| prob: 0.00693 - cup
| prob: 0.00665 - counter

The output from model trained on G70M is

top-20 predictions:
| prob: 0.15203 - coffee
| prob: 0.09728 - shop
| prob: 0.09182 - counter
| prob: 0.03848 - interior
| prob: 0.03389 - bar
| prob: 0.03215 - restaurant
| prob: 0.02440 - table
| prob: 0.02245 - store
| prob: 0.01950 - area
| prob: 0.01905 - inside
| prob: 0.01590 - starbucks
| prob: 0.01313 - cafe
| prob: 0.01220 - chair
| prob: 0.01172 - floor
| prob: 0.01020 - cup
| prob: 0.00879 - drink
| prob: 0.00794 - room
| prob: 0.00746 - customer
| prob: 0.00635 - wood
| prob: 0.00345 - bakery

License

This project is under the CC-BY-NC 4.0 license. See LICENSE for details.

[^1]: Image credit: ゼルダの伝説 The Legend of Zelda: Tears of the Kingdom. [^2]: Image credit: Space-X. [^3]: Image credit: OpenAI Sora. [^4]: Image credit: Demo in Segment Anything | Meta AI. [^5]: Image credit: Super Mario Bros Wonder. [^6]: Image credit: Photo taken by the author at a Starbucks store.