안녕하세요, 마흔 번째 X-Review입니다. 이번 논문은 2022년도 ACM MM에 게재된 Reading and Writing: Discriminative and Generative Modeling for Self-Supervised Text Recognition입니다. 바로 시작하도록 하겠습니다. 😵

1. Introduction

Text recognition task는 텍스트 영상에서 어떤 텍스트인지 인식하는 task로 text가 위치하는 영역을 따로 찾을 필요 없이 이미지 전체에 text가 크게 차지하고 있고 이걸 그냥 인식하는 task입니다. 현존하는 대부분의 text recognition 모델들은 대용량의 합성 데이터셋에 크게 의지를 하는데요, 이는 실제 영상을 어노테이션하데 비용이 좀 들기 때문입니다. 하지만 이 합성 데이터와 실제 데이터 사이의 도메인 갭이 성능에 영향을 미치게 되겠죠. 그래서 self-supervised 학습을 통해 라벨이 되어 있지 않은 real 영상을 사용하는 연구들이 나오게 되었습니다.

구체적으로 SeqCLR이나 PerSec이라고 하는 contrastive learning 기반의 text recognition 모델들이 있는데요, 이 두 방법론들은 text 영상에서 구별력을 학습하는데 집중한 것으로 Fig1의 (a)를 보시면 입력으로 PRIVATE라고 적힌 영상이 들어왔을 때 augmentation을 통해 positive 영상을 생성해내고 encoder를 태워 추출한 sequence들간의 거리를 좁히는, 다른 negative 영상 sequcne간의 거리는 멀게하는 식의로 동작합니다.

저자는 사람이 어떤 text인지 인식하는 과정을 학습할 때 읽기와 쓰기를 통해 학습한다는 사실을 발견을 하였습니다. 읽기, 즉 reading이란 다른 시점에서 보이는 text들이나, 다른 모양을 갖는 text들에 대한 외관 차이를 관찰하는 것인데 이를 통해 구별력, 변별력을 학습할 수 있도록 도와줍니다. 또 쓰기 Writing은 텍스트 영상을 더 깊게 인식할 수 있는 generative한 방식이라고 볼 수 있겠습니다. 이 reading과 writing에 대한 조합은 사람들로 하여금 텍스트를 더 잘 읽도록 도와줍니다. 저자들은 이러한 관찰에서 영감을 받아 Discriminative and Generative self-supervised 방식 (이하 DiG)를 제안합니다.

위 Fig 1(b)가 저자들이 제안한 DiG를 보여주고 있는데 앞서 언급했던 두 모델(fig1-(a))와 동일하게 contrastive learning기반이며 여기에 masked image modeling을 더해서 discriminative 모델과 generative 모델의 장점을 모두 활용하고자 하였습니다. 구체적으로 그림을 보며 좀 더 설명해보자면, DiG 모델은 입력 영상의 두 가지 view를 사용하는데 첫 번째 view는 마스킹 된 view이고 두 번째 view는 augmented view입니다. 이 두 view는 encoder에 입력으로 들어가 contrastive learning을 하게 되는 것이구요 동시에 masking된 view는 원본 영상으로 reconstruction하는 과정을 수행하는데 사용됩니다. 즉, 이런 방식을 통해 text 영상의 판별과 생성 (differentiation , generation)을 모두 학습할 수 있게 함으로써 더 강력한 feature representation을 생성하도록 한 것입니다.

본 논문에서 제안된 DiG는 이런 학습 구조로 인해 text recognition을 위한 모델 뿐만 아니라 text와 관련된 다른 task 가령 text segmentation이나 text image super-resolution task에서도 사용이 될 수 있다고 합니다.

주 contribution은 다음과 같습니다.

1. text recognition을 위한 generative self-supervised 학습을 처음으로 제안함으로써 discriminative하고 generative한 self-supervised 모델을 제안함.
2. 제안된 DiG는 constrastive 학습과 masked image modeling을 통합함으로써 discrimination과 generation을 모두 학습할 수 있게 하였으며 성능 향상
3. text recognition task를 포함해 text segmentation, text image super-resolution에 대한 task에서 성능 굳. SOTA.

2. Methodology

본 방법론은 일반적인 self-supervised 학습 절치를 따르는데요, label 없이 사전 학습한 후 task-specific한 fine-tuning을 하는 순입니다. 이제 모델 아키텍처를 구체적으로 살펴도록 하겠습니다.

2.1. Architecture

DiG의 아키텍처는 위 Fig2에서 확인할 수 있습니다. 입력 영상이 처음에 고정된 크기 32 x 128로 resize되게 되구요, 그 다음 입력 영상에 대해 두 view, 마스킹된 뷰인 X_m와 augmentation된 view인 X_a를 생성해냅니다. 이 두 영상은 ViT encoder F(.)의 입력으로 들어가게 되구요 각각 masked feature인 f_m(빨간색 feature)와 augmented feature f_a(초록색 feature)를 추출해냅니다. masked image modeling 브랜치 부분에서는 recognition head가 있어서 이 masked feature f_m를 reconstruction하는데 사용합니다. constrastive learning branch에서 f_m와 f_a는 patch head, projection head, prediction head를 통해 각각의 query q_m, q_a를 생성하구요, 아래 momentum branch부분에서는F_t(.), P_t(.), K_t(.)를 통해 key \hat{k}_m, \hat{k}_a를 생성해냅니다. 마지막으로 이 query와 키를 통해 constrastive loss를 계산해냅니다. 

2.2. Contrastive Learning

contrastive learning에 대해 구체적으로 살펴보도록 하겠습니다. DiG 모델은 MoCo v3를 베이스로 contrastive 학습을 하였으며 몇 부분만 수정하였습니다.

Data Augmentation, Encoder

Intro에서 언급한 바 있는 SeqCLR 모델이 사용한 증강 기법을 사용하며 그 외에 color jitter와 gray 등을 추가하였으며, encoder로는 ViT를 사용하였습니다.

Patch Head, Projection Head, Prediction Head

일반적인 영상에 대한 contrastive learning은 보통 encoding된 feature map을 하나의 단일 벡터로 투영한 다음 이를 통해 loss를 계산하게 되는데요, 근데 text 이미지 같은 경우 일련의 문자를 나타내고 있기 때문에 feature map을 수평으로 4개의 patch로 분할해 각각을 사용하였습니다.

projection head K(.)는 3 layer MLP, prediction head Q(.)는 K(.)와 유사하게 2 layer MLP를 사용하였습니다.

Momentum Branch

Momentum Branch같은 경우 F_t, P_t, K_t로 구성되며 각각은 F(.), P(.), K(.)와 같은 아키텍처를 갖고 있습니다.

2.3. Masked Image Modeling

masked image modling 모듈은 기본적으로 SimMIM을 가져와서 사용하였습니다.

Masking Strategy

masking strategy는 전체 영상에 대해 몇 %의 영상을 마스킹할 것인지에 대한 masking ratio, 마스킹 하는 방식(랜덤하게 할건지 등) 등 마스킹 과정에 대한 전반적인 부분들을 의미합니다. 본 DiG같은 경우 MAE와 동일하게 patch 단위로 랜덤하게 마스킹하도록 하였고 patch 크기는 4×4로, masking ratio는 실험에 따라 0.6으로 설정하였습니다.

Encoder

인코더는 contrastive learning branch와 동일한 ViT encoder를 공유하도록 하여 두 task에서 동일한 feature를 추출하도록 하였습니다.

Prediction Head

prediction head는 앞선 encoder 단계를 통해 잘 만들어진 token을 통해 마스킹된 영역들에 대한 복원을 수행하는 decoding 단계인데요, 이 prediction head는 SimMIM과 같이 단 하나의 linear 만을 사용하였습니다.

Reconstruction Target

마지막으로 reconstruction target은 pretext task로 마스킹된 영역을 복원할 때 어떤 대상을 복원하는 것을 목표로 잡느냐는 의미인데요, 여기서는 마스킹된 RGB 픽셀 값을 예측하도록 하였습니다.

2.4. Optimization

DiG 모델의 전체 logg 함수는 contrastive learning loss L_c와 masked image modeling loss 함수 L_m으로 구성됩니다. 하나씩 살펴보자면 먼저 contrastive learning loss L_c는 아래 식 2번에서 볼 수 있듯이 InfoNCE loss를 사용합니다. 

다음으로 masked image modeling loss L_m는 l2 loss 함수를 사용하였습니다. 식3에서 x_i와 y_i는 각각 예측된 pixel값과 pixel i의 RGB gt 값을 나타냅니다. 

2.5. Text Recognition

다음으로fine-tuning 과정에서 사용되는 text recognition model에 대해 살펴보자면, ViT encoder와 text recognition decoder로 구성이 되어 있는데요 디코더는 인코더를 타고 나온 2D feature를 문자 시퀀스로 변환하는데 이때 사용되는 디코더는 아래 그림 3에서 볼 수 있듯이 CTC, Attention decoder 혹은 transformer decoder 중에 하나입니다.

CTC decoder는 1D feature 시퀀스에서 동작하는 디코더로 보시면 인코더 타고 나온 2D feature map을 수직 방향으로 average pooling 때려서 1D 시퀀스로 변환한 후 입력해 문자 시퀀스를 출력해냅니다.

2.6. Text Segmentation

마지막으로 앞서 짧게 DiG는 text recognition task외에 text와 관련된 text segmentation이라던지 text image super-resolution이 가능하다고 했었는데요. Text segmentation의 모델을 간단히 살펴보자면, ViT 인코더를 사용하고 3개의 multi-head attention layer와 하나의 linear prediction head로 구성된 segmentation decoder를 사용합니다.

2.7. Text Image Super-Resolution

text image super-resolution model도 동일하게 ViT 인코더를 사용하며, decoder도 유사한 segmentation 모델과 유사한 구조를 갖는다고 합니다.

3. Experiments

마지막으로 실험 부분 살펴보겠습니다. 먼저 실험에 사용한 데이터셋은 Unlabeld Real Data(URD)와 Synthetic Text Data(STD), Annotated Real Data(ARD)로 나눌 수 있겠습니다. Text Recognition 벤치마크 데이터셋으로는 아래 table1을 참고하면 좋을 것 같습니다.

크게 정방향으로 되어 있는 regular dataset과 좀 곡선 text위주로 구성된 irregualar dataset, occlusion이 있는 데이터셋으로 나눠볼 수 있겠습니다.

3.1. Self-Supervised Learning

먼저 self-supervised learning에 대한 결과를 보도록 하겠습니다.

Reconstruction of the Masked Image Modeling

위 Fig4가 마스킹된 이미지 reconstruction 결과를 시각화한 것인데요, 보시면 맨 윗 행이 원본 영상 두번째 행이 random mask한 영상 세번째 행이 reconstruction항 형상입니다. 영사잉 흐리거나 좀 가려진 경우에도 잘 reconstruction한 것을 확인할 수 있으며, 이를 통해 좋은 generation feature representation을 학습했다는 결론을 내려볼 수 있겠습니다.

Feature Representation Evaluation

다음으로는 encoder의 가중치를 fix하고 annotation이 되어 있는 실제 데이터로 text recognition decoder를 학습해 feature의 표현력을 평가해본 실험입니다.

Fig5에 정성적 결과가 있는데 (a) ~ (f)는 각각 곡선 text, background가 복잡한 데이터셋, 왜곡이 있는 text, 흐릿한 text, occlusion이 있는 text, word art 등 다양한 경우가 존재합니다. 한 영상에 대해 옆에 3개의 예측된 text가 적혀져 있는데 위에서부터 각각 Gen-ViT-small, Dis-ViT-Small, DiG-ViT-small입니다. 즉, generative 모델과 discriminative 모델 그리고 이 둘을 합친 DiG 모델을 비교하고 있는데요 보시면 DiG 모델이 예측한 결과는 전부 잘 예측됐지만 Gen, Dis은 위와 같이 좀 어려운 영상에 대해 잘못된 예측을 하고 있습니다. 

정량적인 성능은 위 table2에서 확인할 수 있는데요 DiG-ViT-Small이 모든 벤치마크에서 Gen과 Dis에 비해 각각 11.8%, 4.1% 성능을 능가한 것을 확인할 수 있습니다. Gen모델과 Dis 모델을 서로 비교해봤을 때 대부분 벤치마크에서 Dis 모델이 더 높은 정확도를 얻었는데, occlusion이 있는 데이터셋에서는 Gen 모델이 더 나은 성능을 보이고 있습니다.

3.2. Text Recognition

Self-Supervised Text Recognition

위 표4는 기존 self-supervised text recognition 모델과의 정량적인 비교를 보여줍니다. UTI-100M이라고 되어 있는 부분은 추가로 100M 영상으로 사전학습되었다는 의미이며 모든 모델은 Synthetic Text Data(STD)에 대해 학습되고 fine-tuning되었습니다. 디테일한 비교를 위해서 3개의 decoder 모두에 대해 실험을 수행했으며, 결과를 보시면 본 논문에서 제안된 DiG는 regular한 데이터셋 뿐만 아니라 Irregular, occluded, handwritten 데이터셋에 대해 SeqCLR과 PerSec 모델의 성능을 크게 능가하는 것을 볼 수 있습니다. 특히 IC15, SVTP, CUTE와 같은 좀 어려운 벤치마크에서 더 성능 개선 폭이 큰데 예를 들어서 DiG-ViT-Tiny 모델에서 CTC 디코더를 사용하는 경우 PerSec-ViT + UTI-100M보다 8.8%, 4.9%, 13.8% 더 좋은 성능을 보이고 있습니다. 이 persec-ViT가 사전학습할 때 1억개의 라벨링 되지 않은 영상을 사용하는 반면에 DiG는 1577만개의 unlabeled 영상만을 사용하는데서 의의가 있다고 볼 수 있겠죠.

Scene Text Recognition

마지막으로 scene text recognition에 대한 실험결과입니다. 위 그림6에서는 scene text recognition sota 모델들에 대한 모델 파라미터 수와 accuracy가 표시되어 있는데요, 본 논문에서 제안된 DiG를 보면 모델 크기와 정확도 사이의 좋은 trade-off를 갖는 것을 확인할 수 있습니다. 예를 들어 빨간색 네모인 DiG-ViT-Tiny같은 경우 가장 적은 파라미터 수를 갖는데도 기존의 모든 방법론보다 좋은 성능을 보이고 잇습니다. 

마지막으로 table6을 통해 정량적인 성능을 살펴보겠습니다. 제안된 DiG 모델은 11개의 scene text recognition benchmark에서 기존 sota 방법론 성능을 크게 능가하고 있는데요 특히 ABINet이라는 모델이 WiKi 데이터로 학습된 모델인데 DiG-ViT small과 비교했을 때 거의 동일한 모델 크기를 갖음에도 4.4% 더 높은 성능을 보였습니다. 또 ViT encoder와 transformer decoder를 사용하는 SATRN과 비교해본다면 DiG-ViT-Base는 모델 크기가 거의 비슷하면서 (52M vs 55M) 평균적으로 11.5% 더 높은 성능을 보이고 있습니다.