UNIT: Unifying Image and Text Recognition in One Vision Encoder

Currently, vision encoder models like Vision Transformers (ViTs) typically excel at image recognition tasks but cannot simultaneously support text recognition like human visual recognition. To address this limitation, we propose UNIT, a novel training framework aimed at UNifying Image and Text recognition within a single model. Starting with a vision encoder pre-trained with image recognition tasks, UNIT introduces a lightweight language decoder for predicting text outputs and a lightweight vision decoder to prevent catastrophic forgetting of the original image encoding capabilities. The training process comprises two stages: intra-scale pretraining and inter-scale finetuning. During intra-scale pretraining, UNIT learns unified representations from multi-scale inputs, where images and documents are at their commonly used resolution, to enable fundamental recognition capability. In the inter-scale finetuning stage, the model introduces scale-exchanged data, featuring images and documents at resolutions different from the most commonly used ones, to enhance its scale robustness. Notably, UNIT retains the original vision encoder architecture, making it cost-free in terms of inference and deployment. Experiments across multiple benchmarks confirm that our method significantly outperforms existing methods on document-related tasks (e.g., OCR and DocQA) while maintaining the performances on natural images, demonstrating its ability to substantially enhance text recognition without compromising its core image recognition capabilities.

Paper

Similar papers

Peer review

Reviewer rNyF5/10 · confidence 4/52024-07-07

Summary

In this paper, the authors propose a novel training framework aiming at unify image and text recognition within a single model. The framework contains both visual and language decoders. In addition, the authors adopt a two-stage training, where the pre-training stage uses intra-scale data while the finetuning stage uses inter-scale samples. The experimental results demonstrate the effectiveness of the framework.

Strengths

(1) The authors aim to unify image and text recognition tasks in one module, which is good for the community. (2) The general pipeline of the paper is clear. (3) The performances improvement in the experiments seems significant.

Weaknesses

(1) Although the authors claim that they unify image and text recognition into one module, I find there are two different decoders. It seems the authors just combine two decoders in the framework, one for image recognition and another for text recognition. Hence, I think this unification is not very elegant. (2) The general pipeline is clear, but some details are not given. For example, the data preparation of both intra-scale and inter-scale training is not clear. In addition, it is not clear why the random feature sampling is effective. Because it will discard tokens, if those discarded tokens contain important information, the recognition will be definitely failed. By the way, I am not sure how many tokes will be discarded commonly. (3) In the paper description, the authors claim that the scale robustness is an important contribution (line 129), but I can not find sufficient experiments about this. In the experiments, I find the input scales are fixed, i.e., documents with 896\*896, and images with 224\*224 (I guess). Why not show more results of flexible inputs about the image and documents? Based on my understanding, the training data is about 896\*896 and 224\*224, so the test data will follow this resolution. If my understanding is correct, I suppose it can not be claimed about the scale robustness. (4) I can not understand the test stage. Does it mean the visual decoder is used for image recognition and the language decoder for text recognition? if it is correct, we need to specify the task for each image during the test stage, which will limit the usage. For example, if an image will be used for both image recognition and text recognition, how to deal with? In addition, this design is not flexible for users.

Questions

(1) In line 162, the authors describe that “low-resolution images, the position embedding are then randomly cropped”. I wonder why low-resolution images will need cropped instead of high-resolution images?

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

Not addressed yet.

Authorsrebuttal2024-08-13

Dear reviewer, Thank you for your thoughtful comments and for engaging in a detailed discussion. **Q2. About random feature sampling** We would like to clarify that the random feature sampling approach is indeed applied **only to natural images**, not document images. As mentioned in lines 242-244 of our paper, this technique is specifically designed to enhance training stability and efficiency when computing the token-wise feature reconstruction loss for high-resolution natural images. The purpose of this feature reconstruction loss is to maintain the original image recognition capability, and it is solely utilized for natural images. Therefore, the concern you have about random sampling causing a loss of text information will not occur. **Q5. About position embedding** Our position embedding is 64x64xd, but for input sizes smaller than 896x896, such as 224x224, a 16x16xd position embedding is required. There are generally three possible solutions: cropping, interpolation, and training another 16x16xd position embedding. We chose cropping for the following reasons: 1. Cropping is a more straightforward approach compared to other methods, as it avoids additional computational costs and memory usage. 2. As mentioned in lines 238-239 of our paper, the low-resolution document dataset consists of images cropped from high-resolution documents. Therefore, we use cropped position embeddings for low-resolution inputs to preserve local structural information, ensuring consistent spatial feature representation across resolutions. For natural images, both interpolation and cropping are viable, and we selected cropping for its simplicity. We hope this explanation addresses your concerns. If our response has clarified the issues and alleviated your concerns, we would sincerely appreciate it if you could kindly consider updating your score.

Reviewer rNyF2024-08-14

thanks for your quick response. Then It make sense for random feature sampling. Since the Table 4 shows the text recognition ability, it very easy to make readers think this is also for document images. I suggest the authors can highlight the sampling for natural images. I plan to increase the rating, but I strongly suggest the authors to make the paper more clear and highlight the contributions.

Authorsrebuttal2024-08-14

Dear reviewer, Thank you for your prompt and thoughtful feedback. We appreciate your understanding regarding the random feature sampling. Following your suggestion, we will make sure to highlight that the sampling is specifically for natural images. We are also committed to revising the paper to ensure that our contributions are clearly presented. Your input has been invaluable, and we're grateful for your increased score. Thank you once again for your time and support!

Reviewer W6eH5/10 · confidence 4/52024-07-09

Summary

This paper presents a new visual backbone training framework named UNIT, which integrate image and text recognition simultaneously. The UNIT framework leverages text recognition, visual feature construction and image captioning as the pre-training tasks. The training pipeline includes two stage: intra-scale pre-training and inter-scale pre-training with different solution for different types of images. Extensive experiments demonstrates UNIT significantly outperforms existing methods on document-related tasks while maintaining performance on natural images without additional inference cost.

Strengths

1. UNIT enables the visual backbone to recognize both text and natural images without the need for additional visual experts. 2. The proposed model retains the original Vision Transformer architecture, making it cost-effective in terms of inference and deployment for various types of images. 3. The two pre-training stages are well-suited for document image pre-training, which typically requires high resolution. 4. The experiments demonstrate that UNIT achieves state-of-the-art performance on various document-related tasks while maintaining strong performance on natural image tasks.

Weaknesses

1. The experiments are insufficient for understanding how the training pipeline works. The paper only leverages resolutions of 1x and 4x and does not show results for training with different resolutions. 2. Document images tend to be non-square. During the pre-training stage, using square images might hinder learning for documents with extreme aspect ratios, such as those found in InfoQA [1]. This issue is not discussed in the paper. 3. The model/data/resolution scaling behaviors of the proposed UNIT is not discussed which is important for pre-training framework. [1] Mathew, Minesh, et al. "Infographicvqa." Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision. 2022.

Questions

1. In Table 3, we observe that even when using image captioning during pre-training, the zero-shot classification performance on natural images remains poor without applying feature reconstruction as the pre-training task. I am curious as to why this approach fails, given that Cappa [1] also utilizes image captioning for pre-training and scales well in natural image classification. Could you provide a detailed discussion on this discrepancy? 2. For the LLaVA experiments, it seems that two different resolutions are applied for different types of images. Why not considering using both 224 or 896 resolution? Is there any negative signals for using large image resolutions for natural image? 3. For the language decoder, using Q-former would be not necessary and not trivial. It seems that using naive MLP layers can also work. [1] Tschannen, Michael, et al. "Image captioners are scalable vision learners too." Advances in Neural Information Processing Systems 36 (2024).

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have addressed the limitation in the manuscripts.

Reviewer W6eH2024-08-12

After reading the authors' feedback and other reviewers' opinions, I would like to thank the authors for their rebuttal. The rebuttal addresses most of my concerns, and thus, I raise my score. I would like to note that even though the authors addressed my concerns, I believe that the paper needs some sort of revision, as it appears that many basic concerns were needed to be clarified in the rebuttal. I strongly encourage the authors to fix the main comments mentioned by the reviewers.

Authorsrebuttal2024-08-12

Dear reviewer, Thank you very much for reconsidering our paper and for raising your score. We appreciate your recognition of our efforts to address your concerns. Your feedback has been instrumental in enhancing the quality of our submission, and we are grateful for the time and thoughtfulness you have dedicated to your review. We have carefully considered your suggestions and are committed to revising the paper accordingly. We believe these revisions will further strengthen our submission.

Reviewer qQPk6/10 · confidence 4/52024-07-11

Summary

The paper proposes to enable vision encoders to support general image recognition and text recognition at the same time, in which a lightweight language decoder is proposed for predicting text outputs and a lightweight vision decoder is proposed to prevent catastrophic forgetting of the original image encoding capabilities. Besides the intra-scale pertaining, an inter-scale fine-tuning stage is utilized to improve the robustness of both image recognition and text recognition. Experiments on multiple datasets demonstrate the effectiveness of the proposed method and its potential to process interleaved text and image information in real-world scenarios.

Strengths

1) The paper proposes to integrate strong text recognition and image recognition into one encoder, which provides good potential for complex document analysis and general vision recognition. 2) The proposed techniques including inter-scale training and feature sampling help the model to deal with diversities for image recognition and text recognition. 3) Experiments on several benchmarks demonstrate the superiority over general encoders (CLIP, DINO) and specialized document models (vary, nougat).

Weaknesses

1) The ability to keep ability from multiple sources is mainly achieved by distillation from a teacher model, which follows existing works like RADIO and is not quite new. 2) As stated in L300-303, the teachers used for other models are different from UNITS, which may lead to unfair comparison. 3) How does the model effectively processing high-resolution document images with the transformer vision encoder? 4) Missing comparison with existing document MLLM, e.g., Monkey. [1] Monkey: Image resolution and text label are important things for large multi-modal models, CVPR 2024. 5) Missing reference for Figure 2. Incomplete information in the references part, e.g., reference [6] A. Brock, S. De, S. L. Smith, and K. Simonyan. High-performance large-scale image recognition without normalization, 2021.

Questions

Please refer to the weakness part.

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

Yes. Positive societal impacts and possible negative societal impacts are well described in the A.3 section.

Authorsrebuttal2024-08-12

Dear Reviewer, Thank you for your thorough review and constructive feedback throughout the process. We appreciate your recognition of our efforts to address your concerns. We have carefully considered your suggestions and are committed to revising the paper as you recommended. These revisions will further strengthen our submission, and we are grateful for your insights that have guided us in this direction. Thank you again for your support!

Authorsrebuttal2024-08-11

Dear reviewer, Thank you for your valuable comments. As the deadline is approaching, we kindly inquire whether our discussions have addressed your concerns. If you have any further questions, we would be happy to continue our conversation. If our response has resolved your concerns, we would greatly appreciate it if you could consider updating your score and providing feedback. Thank you again.

Authorsrebuttal2024-08-11

Dear reviewer, We have provided detailed clarifications on all your questions point to point. If you have any additional questions or suggestions, we would be very happy to discuss them with you. We sincerely hope you can appreciate the significance of our work and consider updating your rating. Thank you!

Reviewer qQPk2024-08-12

Official Comment by Reviewer qQPk

Thank the authors for their response. Regarding W1, in fact, the proposed method integrates both distillation in works like RADIO and supervised learning widely used in OCR methods, e.g., vary, which is not originally designed. Regarding W2-W5, I appreciate the response and believe the details serve as a good supplement to the original manuscript. I'd like to maintain the rating in the current stage.

Reviewer rNyF2024-08-13

I am sorry that it is a little late for the discussion. Thanks authors for the responses. These responses really address some of my concerns, like the usage of encoder and decoder. However, I still can not understand some of your answers, like Q5, why crop 16*16*d for 224*224, but never mind, I can read the original paper CPE for more details. Another question is your answer in Q2 about the random feature sampling. Although you give an example of the MAE paper, it is about the natural image, which is very different with the document image. I do not believe that dropping 75% of patches can be predicted for the document images, because the text patterns are much more complicated than the general objects, like cat, dog.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC