STAIR : Learning Sparse Text and Image Representation in Grounded Tokens (2023)

Abstract: Image and text retrieval is one of the foundational tasks in the vision and language domain with multiple real-world applications. State-of-the-art contrastive approaches, e.g. CLIP, ALIGN, represent images and texts as dense embeddings and calculate the similarity in the dense embedding space as the matching score. On the other hand, sparse semantic features like bag-of-words models are more interpretable, but believed to suffer from inferior accuracy than dense representations. In this work, we show that it is possible to build a sparse semantic representation that is as powerful as, or even better than, dense presentations. We extend the CLIP model and build a sparse text and image representation (STAIR), where the image and text are mapped to a sparse token space. Each token in the space is a (sub-)word in the vocabulary, which is not only interpretable but also easy to integrate with existing information retrieval systems. STAIR model significantly outperforms a CLIP model with +4.9% and +4.3% absolute Recall@1 improvement on COCO-5k text \rightarrow image and image \rightarrow text retrieval respectively. It also achieved better performance on both of ImageNet zero-shot and linear probing compared to CLIP.

## introduction learning high-quality and performant representations from large-scale image-text data has been intensively studied in recent years. improved vision-language representation benefits many applications, including image-text retrieval @xcite , vqa @xcite , and image captioning @xcite . the resurgence of interest in contrastive learning has fueled recent advances. state-of-the-art models like clip @xcite and align @xcite ) employ 1 † work done while at apple figure 1 : learning a sparse and interpretable embedding for a vision-language model. clip @xcite and align @xcite learn a compact dense embedding, while stair aims to learn a sparse and interpretable embedding in a highdimensional space. each dimension in the sparse embedding represents a (sub-)word from a large dictionary, with a non-negative importance weight. despite the impressive benchmark performance, the dense embedding space is usually considered as a black box and challenging to interpret. the meaning of an dense embedding is determined by its vector relationships with other embeddings, rather than corresponding directly to any humanunderstandable concept. this lack of direct correspondence with interpretable concepts hinders the transparency and interpretability. moreover, deploying a retrieval model trained with a contrastive objective to real-world large-scale image-text retrieval is a non-trivial task. despite approximated nearest neighbor search @xcite can be used to retrieve from a dense embedding space, the cost can be high when scaling up to billions of images. additionally, traditional approaches like inverted indexes, commonly employed by search engines, cannot be straightforwardly applied to the dense embeddings. the lack of interoperability further complicates the integration of dense embeddings with other retrieval features without additional training. in comparison, interpretable sparse embeddings possess distinct advantages across numerous applications. more detailed discussion can be found in appendix d. the community has explored sparse and interpretable image representations extensively. early approaches, e.g. bag-of-words and topic models @xcite @xcite were widely coupled with sift descriptor @xcite but found to exhibit inferior performance compared to dense vectors @xcite . other endeavors, like deep visual-semantic embedding @xcite using imagenet topics yielded more interpretability but still failed to outperform dense representations @xcite . we hypothesize that the gap between sparse semantic and dense embedding stems from two primary factors: (1) previous works on semantic embedding inadequately explored large-scale training to capture the rich semantics in the image-text domain. (2) most existing semantic embedding methods are built on a fixed vocabulary (e.g. thousand of concepts), which cannot handle out-of-vocabulary concepts. in this paper, we present a new model, named stair, and a multi-stage training recipe to learn a sparse text and image representation to tackle the aforementioned challenges. our sparse representation not only matches but also outperforms the dense image-text representations. inspired by the recent success of sparse representation in information retrieval field @xcite , stair encodes the image and text into a sparse and grounded token space, as illustrated in figure 1 . specifically, images and text are mapped to sparse and interpretable in a high-dimensional space, where each dimension is associated with an actual (sub-)word from a large dictionary, with non-zero weights. notably, the proposed multi-stage training recipe plays a critical role in the grounding of the sparse embedding. we conduct a comparative analysis between the stair model and a clip model, sharing the same architecture and dataset. experiment results a significant improvement of stair model over clip on image-text retrieval tasks, with +4.9% and +4.3% recall@1 on coco-5k @xcite @xmath0. stair models also demonstrate similar or better performance on zero-shot classification and linear probing tasks including imagenet. to quantify the interpertability of the embeddings, we define an interpretable space using bert vocab. experiments suggest stair is substantially more interpretable, achieving a top-1 accuracy 32.9% on imagenet, outperforming clip's accuracy of 13.7%. ## approach we begin with an overview of the dual-encoder architecture to establish the basis of our research. ## training details the stair model aims to achieve two goals: 1) aligning text and images in the sparse embedding space; 2) grounding the sparse embedding dimension with human-understandable (sub-)word in the vocabulary. in other words, the image and text embeddings should reside in a space defined by basis vectors that correspond to interpretable tokens. however, in practice, we found that simply replacing the dense projection head with a token prediction head alone does not guarantee the 2 nd goal outright. this is because the images and text are distinct modalities with inherent semantic gaps, and contrastive loss alone only encourages text/image alignment. we shown in section 5.2 that the model learns a shortcut by relying on less common tokens to bridge the modality gap. to address this challenge, we propose a multistage training approach that sequentially bridges the discrepancy between the sparse embedding space and the interpretable space defined by the vocabulary as illustrated in figure 3 . ## related work image and text retrieval image-text retrieval approaches can be categorized into dual-encoder and cross-encoder approaches. in the dual-encoder approach, images and text are encoded separately as dense embeddings. devise @xcite was one of the pioneering dual-encoder models. with the advent of transformers, @xcite proposed clip that leverages large-scale pretraining datasets and established new state-ofthe-art across multiple benchmarks. finetuning the visual model @xcite and lan-guage model @xcite further improves performance. on the other hand, the cross-encoder approach like uniter @xcite ) employs a single encoder for both image and text inputs. while cross-encoder models capture fine-grained alignment between image and text but are slower at inference, dual-encoder models offer faster retrieval speed with precomputed embeddings. stair follows the dual-encoder approach but introduces a novel aspect by using sparse embeddings instead of dense embeddings. this choice of sparse embeddings in stair leads to improved interpretability and better retrieval performance compared to dense embeddings. ## conclusion in this paper, we introduced stair (sparse text and image representation), an approach that encodes image and text inputs into sparse embeddings within a sparse token space. we also employed a multi-stage training strategy to ensure that the embeddings are grounded in meaningful tokens. by comparing stair with the clip model, we observed that stair achieved superior performance on image-text retrieval tasks and demonstrated better results on various zero-shot and linear probing classification tasks. furthermore, through quantitative and qualitative analysis, we illustrated the interpretability advantage of our sparse embeddings over dense embeddings, making them more easily understandable for humans.