• 현재의 Visual Document Understanding (VDU) 기법은 텍스트 추출을 OCR 엔진에 맡기고, OCR이 반환한 텍스트를 이해하는 방식에 집중해왔다.
• 하지만 OCR을 활용하는 것에는 다음의 3가지 문제가 존재한다.
  1. OCR을 사용하기 위해 높은 컴퓨팅 리소스가 필요하다.
  2. OCR 모델은 문서의 종류 혹은 언어에 대한 유연성을 갖고있지 않다.
  3. OCR에서의 에러가 다음 절차에 영향을 미친다.
• 따라서 OCR을 쓰지 않는 VDU 모델인 Donut(Document nderstanding transformer)을 제안한다.

github: clovaai/donut paper: <OCR-free Document Understanding Transformer>, Kim et al., 2022

1. Introduction

• Visual Documnet Understanding (VDU): 문서 이미지 (송장, 영수증, 명함 등)에서 정보를 추출하는 작업으로, 산업에서 꼭 필요한 작업이면서도 아래의 세부 작업이 수반된다는 점에서 연구자들에게 어려운 주제이다.
  • 문서 분류 (document classification)
  • 정보 추출 (information extraction)
  • 시각질의응답 (Visual Question Answering (VQA)

• 기존 VDU 방법은 two-stage 방식으로 풀이된다. Fig.1. The schema of the conventional document information extraction (IE) pipeline
1. 문서 이미지에서 텍스트 읽기
2. 문서의 총체적인 이해

• OCR-dependent한 방법의 문제점.

  • OCR 전처리는 비용이 높다.
    • 사전학습된 기성 OCR 엔진을 사용할 수 있지만, 고품질의 OCR 결과를 위한 추론 비용이 높다.
    • 기성 OCR은 다양한 언어 및 도메인의 변화에 맞춘 유연한 대응이 어렵다.
    • 직접 OCR 모델을 훈련시키는 것 또한 상당한 규모의 데이터셋과 비용이 수반된다.
  • OCR에서의 에러가 VDU 시스템에 전파됨으로써 후속 절차에 부정적인 영향을 끼친다.
    • 한국어, 중국어와 같이 OCR의 품질이 상대적으로 낮은 경우에 이 문제는 더 심각해진다.
    • 이를 위해 OCR 후처리 (교정) 모듈을 도입하는 것이 일반적이나, 전체 시스템의 규모와 유지보수 비용의 측면에서 실용적인 해결책은 아니다.

• 제안

  • 이에 OCR 없이 raw input image를 의도된 output으로 직접 매핑시키는 모델을 제안한다.
  • 아키텍쳐: Transformer-only
  • 간단한 아키텍쳐와 사전훈련 방법으로 기존 방법보다 전반적으로 더 나은 성능을 얻을 수 있다 (아래 Fig.2. 참조).
    Fig.2. The pipeline overview and benchmarks

• 훈련 스키마

  • pre-training: 어떻게 텍스트를 읽는지 학습 (이미지와 이전 텍스트 문맥을 함께 조건화하여 다음 단어를 예측하는 방식)
    • 합성 데이터를 사용하여 도메인과 언어에 대한 유연성을 확보
  • fine-tuning: downstream task에 맞게 전체 문서를 이해하는 법을 학습

• 본 논문의 기여점

  1. End-to-End 방식으로 훈련된 최초의 Transformer 아키텍쳐의 OCR-free VDU 방법
  2. 합성 데이터의 활용을 가능케 한 단순한 사전학습 스키마
     • 본 논문에서 제안한 SynthDoG (합성 데이터 생성기)를 통해 편리하게 다국어 설정을 확장할 수 있다.
  3. 공공 벤치마크와 특정 산업의 데이터셋 모두에 대한 다양한 실험을 진행함으로써 본 방법이 단순히 SOTA 성능을 낼 뿐만 아니라 실용적인 이점(e.g. 비용효율)을 제공한다는 점도 확인했다.
  4. GitHub에 소스코드, 사전학습 모델, 합성 데이터를 공개하였다.

2. Method

Fig.3. The pipeline of Donut

2.1 Preliminary: background

(위 Introduction에서의 내용과 겹치므로 생략)

2.2 Document Understanding Transformer

• Donut은 문서 이미지에 대한 전반적인 이해를 하는 End-to-End의 VDU 모델이다.
• 아키텍쳐는 Transformer 기반의 visual encoder와 textual decoder로 구성되어 있다(전체 흐름은 위 Fig.3 참조).
  • visual encoder: 주어진 문서 이미지에서 특징 추출을 수행하며, 그 어떤 OCR 기능과 연관된 모듈을 사용하지 않는다.
  • textual decoder: 파생된 특징(visual encoder에서 나온 feature)를 일련의 subword 토큰과 매핑하여 원하는 구조화된 형식을 구성한다(e.g., JSON).

Encoder

• Visual encoder는 문서 이미지(input)를 embeddings(output)으로 변환한다.

  • input: $x \in \mathbb{R}^{H \times W \times C}$
    • $H$: 이미지 높이
    • $W$: 이미지 너비
    • $C$: 이미지 채널
  • output: ${z_i \mid z_i \in \mathbb{R}^d, 1 \leq i \leq n }$
    • $n$: feature map의 크기 or 이미지 패치의 수
    • $d$: 인코더의 latent vectores의 차원 수

• 아키텍쳐는 CNN-based 모델, Transformer-based 모델 모두 사용 가능하며, 본 논문에서는 Swin Transformer를 사용하였음.

  • 사유: 문서 파싱과 관련된 선행 연구에서 가장 좋은 연구를 보였기 때문
    Swin Transformer의 활용

    1. input image x를 non-overlapping patches로 분할
    2. Swin Transformer blocks을 patchtes가 통과
    • Swin Transformer block의 구성: shifted window-based multi-head self-attention 모듈과 2겹의 MLP
    3. patch를 병합하는 레이어가 각 단계에서 patch tokens에 대해 적용
    4. 최종 block의 output인 {z}는 textual decoder에 전달

Decoder

• ${z}$가 주어졌을 때, textual decoder는 token sequence를 생성한다.
  • token sequence: $(y_i)_{i=1}^{m}$
    • $y_i \in \mathbb{R}^v$ : $i$번째 토큰의 one-hot vector
    • $v$: 토큰 사전의 크기
    • $m$: 하이퍼파라미터 (각각)
• 아키텍쳐: BART
  • 디코더 모델의 initial weight를 사전학습된 multi-lingual BART 모델로 초기화

Model Input

• 전통적인 Transformer의 방법을 따라 teacher-forcing 스키마를 사용
  • teacher-forcing: 학습 간 input으로 이전 time step의 모델 output이 아닌 ground-truth를 사용하는 전략
• Test 때는 GPT-3에서 영감을 받아, 주어진 프롬프트에 맞게 모델이 토큰을 생성
  • 실험 간, 각 downstream task의 프롬프트에 스페셜 토큰 추가
  • 위 Fig.3의 output sequence에 실험 때 사용한 프롬프트 확인 가능
    • teacher-forcing과 decoder의 출력 형식에 관한 추가적인 설명은 논문의 Appendix A.4에서 확인 가능

Output Conversion

• output으로 나온 토큰 시퀀스를 원하는 구조의 형식으로 변환한다.
  • 본 논문에서는 JSON 형식을 채택함 (높은 표현 능력 때문)
• 위 Fig.3에서 확인이 가능하듯, 하나의 토큰 시퀀스는 JSON 데이터로 1:1 변환이 가능하다.
• 스페셜 토큰은 [START_*], [END_*] 2개만을 사용하였다.
  • *은 추출할 각 영역을 나타낸다.
  • output 토큰 시퀀스가 잘못 구성되면, 해당 영역이 유실된 것으로 처리한다.
    • e.g.) [START_name]은 존재하지만 [END_name]이 없다면 "name" 영역은 추출에 실패한 것으로 간주하였다. - 이는 간단한 정규표현식으로 구현 가능

2.3 Pre-training

Task

• 모델은 이미지 내 모든 텍스트를 순서에 맞게 읽어내는 것을 학습한다(기본적으로 좌상단 -> 우하단 순 진행).
  • 학습 목표: 이미지와 이전 문맥을 함께 조건화하여 다음 토큰 예측의 cross-entropy loss를 최소화하는 것
  • pseudo-OCR task로도 해석할 수 있다.
• 모델은 시각적 말뭉치(visual corpora), 즉 문서 이미지에 대한 시각 언어 모델 (visual language model)로 학습된다.

Visual Corpora

• 데이터셋 구축
  • IIT-CDIP: 11M 규모의 스캔된 영어 문서 이미지 데이터셋
  • CLOVA OCR API (상업용): pseudo text label 획득을 위해 활용
• Synthetic Document Generator (SynthDoG): 하지만 위와 같은 데이터셋은 (특히 다른 언어에 대해) 매번 구축 및 사용이 어렵기 때문에 본 논문의 연구진은 합성 데이터 생성기를 만들었다.
  • SynthDoG와 위키피디아(중국어, 일본어, 한국어, 영어)를 사용하여 각 언어별 0.5M 샘플 생성

Synthetic Document Generator

Fig.4. Generated English, Chinese, Japanese, and Korean samples with SynthDoG.

• 이미지 렌더링 파이프라인은 <SynthTIGER: Synthetic Text Image GEneratoR Towards Better Text Recognition Models>, 2021, Yim et al.에서의 방법을 따랐다.
  • 위 Fig.4에서 볼 수 있듯, 생성된 샘플은 여러 개의 컴포넌트로 구성되어 있다; 배경, 문서, 텍스트, 레이아웃.
  • 배경 이미지: ImageNet에서 샘플링
  • 문서의 질감: 수집한 종이 사진에서 샘플링
  • 단어와 어구: 위키피디아에서 샘플링
  • 레이아웃: 그리드를 무작위로 쌓는 간단한 rule-based 알고리즘에 의해 생성
  • 실제 이미지를 모방하기 위해 여러 이미지 렌더링 기술이 적용되었다.
    활용된 이미지 렌더링 기술

    • <Synthetic Data for Text Localisation in Natural Images>, A. Gupta et al., 2016
    • <UnrealText: Synthesizing Realistic Scene Text Images from the Unreal World>, Long, S. et al., 2020
    • <Synthtiger: Synthetic text image generator towards better text recognition models>, Yim et al., 2021
  • SynthDoG에 대한 자세한 내용과 코드는 논문의 Appendix A.2에서 확인할 수 있다.

Fine-tuning

• 모델은 문서 이미지를 어떻게 이해할지 학습한다.
• 위 Fig.3에서 볼 수 있듯, 본 논문에서 모든 downstream tasks는 JSON 예측 문제로 간주한다.
  • decoder는 JSON으로 변환이 가능하면서도 원하는 output 정보를 나타내는 토큰 시퀀스를 생성하도록 학습된다.

3. Experiments and Analysis

3.1 Downstream Tasks and Datasets

Fig.5. Samples of the downstream datasets.
(a) Document Classification.
(b) Document Information Extraction.
(c) Document Visual Question Answering

Document Classification

• 본 논문에서는 모델이 서로 다른 종류의 문서를 구별할 수 있는지 확인하기 위해 classification task를 테스트하였다.
  • 일반적으로 인코딩된 임베딩에 대해 softmax를 통해 클래스 라벨을 예측하는 다른 모델들과 달리, Donut은 문제해결 방법의 일관성을 유지하기 위해 클래스 정보가 담긴 JSON을 생성하도록 하였다.

Document Information Extraction (IE)

• 모델이 문서의 복잡한 레이아웃과 문맥을 완전히 이해했는지 확인하기 위해, 본 논문에서는 다양한 실제 문서 이미지에 대해 IE task를 테스트하였다.

• 모델은 단순히 문자를 잘 읽는 것을 넘어서, 레이아웃과 의미를 이해하여 텍스트 간 복잡한 구조를 도출해낼 수 있어야 한다.

• 평가지표는 2가지를 채택하였다: field-level F1 score, Tree Edit Distance (TED) based accuracy

  1. F1 Score: 추출한 영역 정보가 ground truth에 속하는지를 확인한다.
     • 본 점수는 하나의 문자라도 놓치면 영역 추출에 실패한 것으로 간주한다.
     • 본 점수의 한계
       1. 부분적인 중복을 고려하지 않는다.
       2. 예측한 구조(그룹, nested hierarchy 등)를 측정할 수 없다.
  2. TED based metric
     • 트리로 표현되는 모든 문서에 사용할 수 있다.
     • $max(0, 1-TED(pr, gt)/TED(\emptyset, gt)$
       • $gt$: ground-truth
       • $pr$: predicted
       • $\emptyset$: 빈 트리

• 데이터셋은 2개의 공개된 벤치마크 데이터셋과 2개의 비공개된 산업용 데이터셋을 사용하였다.

  The Consolidated Receipt Dataset (CORD)
  Ticket
  Business Card (In-Service Data)
  Receipt (In-Service Data)

Document Visual Question Answering (DocVQA)

• 본 논문에서는 모델의 성능을 더 검증하기 위해 DocVQA를 수행하였다.
• 문서의 이미지와 질문 한 쌍이 주어지면 모델은 이미지 내 시각 및 텍스트 정보를 모두 참고하여 질문의 답변을 반환한다.
• 일관된 방법을 유지하기 위해, 질문을 시작 프롬프트로 세팅하고 decoder가 답변을 생성하도록 하였다.
  데이터셋: DocVQA

3.2 Setups

3.3 Experimental Results

Document Classification

Document Information Extraction (IE)

Document Visual Question Answering (DocVQA)

3.4 Further Studies

On Pre-training Strategy

On Encoder Backbone

On Input Resolution

On Text Localization

On OCR System

On Low Resourced Situation

4. Related Work

4.1 Optical Character Recognition (OCR)

4.2 Visual Document Understanding (VDU)

5. Conclusions

• Donut은 기존의 VDU 기법들과 달리 OCR에 의존하지 않고 end-to-end 방식으로 학습되었다.
• 문서 이미지 합성 데이터를 생성하는 SynthDoG를 제안하였다.
  • 이를 통해 Donut은 실제 문서 이미지 데이터셋의 규모와 여러 언어 확장을 용이하게 하였다.
• 학습 단계를 "어떻게 읽을 것인지"와 "어떻게 이해할 것인지"로 나누어 점진적으로 학습시켰다.
• 실험과 분석을 통해 공공 벤치마크와 내부 서비스 데이터셋에 대해 본 모델이 더 높은 성능과 더 나은 비용 효율성을 가지는 것을 확인하였다.

[그림-1] 논문 Abstract

아래 내용을 확인해보시면 본 논문 이해에 도움이 될 수 있습니다.

• ResNet
• Swin Transformer

Abstract

• Hierarchical Transformer(e.g., Swin Transformers)는 바닐라 ViT(Vision Transformer)의 문제를 깨고 일반적인 CV(Computer Vision) Task에 실질적으로 적용될 수 있도록 함.

• 하지만 이는 convolution에 내재된 inductive bias가 아닌 Transformer 자체의 우수성에 기인한 결과임.

• 이 논문의 목표는 ConvNet의 한계를 시험하는 것임에 따라, 표준 ResNet을 ViT의 방향으로 점진적으로 현대화(modernize)시키면서 성능 차이를 야기하는 핵심 요소를 파악하고자 함.

• 그 결과가 순수 ConvNet 구조인 ConvNeXt.

1. Introduction

1) ImageNet moment

• 2010년대에 _ImageNet moment_라고도 불리우는 AlexNet의 등장 이후로 시각 인식(visual recognition) 분야는 특징 추출(engineering features) 중심에서 아키텍쳐 설계(ConvNet 기반) 중심으로 이동함.

• 특히 고해상도 이미지에서의 시각처리에 필수적인 sliding window 전략과 ConvNet에 내장된 inductive bias가 잘 맞음.

2) Transformer

• NLP(자연어 처리) 분야에서 Transformer가 RNN 구조를 대체하고 주류 backbone 구조로 자리잡고, 이후 2020년 ViT의 등장으로 NLP와 CV의 흐름이 수렴함.

• 이미지를 일련의 patch들로 분할하는 ViT의 첫 patchify 레이어를 제외하면 ViT에는 이미지에 특화된 inductive bias도 없고 기존 Transformer에서의 큰 변경점도 없음.

• ViT의 주목할만한 점은 확장성인데, 모델과 데이터셋 규모가 커질수록 Transformer는 ResNet보다 훨씬 좋은 성능을 낼 수 있음.

• 하지만 ConvNet의 inductive bias가 없는 상황에서 가장 큰 문제점은, ViT의 global attention 구조가 input size에 대해 이차적 복잡성(quadratic complexity)를 지닌다는 것이고, ImageNet 수준에서는 괜찮았어도 고해상도 이미지에서는 극복할 수 없는 문제가 됨.

3) Hierarchical Transformer

• 이 문제를 해소하기 위해 하이브리드 접근법을 도입함. 예를 들어, sliding window(e.g. local window 내에서의 attention) 전략이 Transformer에 도입되었음.

• 그 흐름의 초석인 Swin Transformer가 Image Classification 외 일반적인 CV Task에서도 SOTA 성능을 기록하는 것을 보고 convolution의 존재가 사라지긴 커녕, 오히려 필요함을 확인함.

• 하지만, sliding window 자체에 대한 단순한 구현은 큰 비용을 요구할 수 있음. 가령 cyclic shifting과 같은 고수준의 접근법을 사용하면 속도가 최적화될 수 있지만, 시스템 설계가 매우 복잡해짐.

4) Modernize ResNet

• ConvNet이 밀리는 것처럼 보이게 하는 핵심 요인으로써 Transformer의 확장성에 크게 기여하는 MHSA(Multihead self-attention)가 있음.

• 시스템 단계에서의 비교를 하자면, ResNet과 Swin Transformer는 공통적으로 유사한 inductive bias를 지니고 있지만, 학습 절차와 거시적인(macro)/미시적인(micro) 단계에서의 아키텍쳐 설계가 크게 다름.

• 이에 표준 ResNet을 개선된 절차로 학습시킨 후, 점진적으로 hierarchical ViT 구조의 방향으로 현대화(modernize)시키면서 _"Transformer 내 어떤 설계 결정이 ConvNet 성능에 영향을 끼칠 것인가?"_라는 질문의 답을 찾고자 함.

• 그 결과 몇가지 핵심 요소를 찾아 적용한 ConvNeXt는 순수 ConvNet 구조임에도 불구하고 accuracy, 확장성 및 견고함 측면에서 모두 Transformer를 뛰어넘음.

2. Modernizing a ConvNet: a Roadmap

[그림-2] 논문 내 Figure2

ResNet이 Transformer와 유사한 ConvNet 구조로 이동하는 기록이 위 표에 담겨있으며, 성능 비교는 각각 ResNet-50과 Swin-T(tiny 모델로 추정), ResNet-200과 Swin-B(base 모델로 추정)에 대해 진행함. 후술되는 내용은 ResNet-50과 Swin-T의 비교에 대한 것임.

2.1 Training Techniques

최신 학습 기법을 적용하면 ResNet-50의 성능을 끌어올릴 수 있다는 최근 연구에 입각하여, DeiT와 Swin Transformer와 유사한 학습 방식을 적용함.

2.2 Macro Design

• Swin Transformer의 거시적인 네트워크 구조(block 및 네트워크 단위)를 분석해보면, ConvNet을 따라 여러 stage를 구성했고, 각 단계(stage)는 각기 다른 feature map 크기(resolution)을 지님.

• 이때 흥미로운 점은 단계에 대한 연산 비율(stage compute ratio)과 줄기세포 구조(stem cell structure)임.

1) Changing stage compute ratio

• Swin-T는 ResNet의 기본적인 원칙은 따르나, 각 단계의 연산 비율은 1:1:3:1로, ResNet-50의 3:4:6:3과 다름.

• 따라서 ResNet-50도 (3, 3, 9, 3)으로 수정하여 Swin-T에 맞춰 FLOPs를 조정함.

Changing stem to "Patchify"

• Swin-T는 patchify 전략으로, 처음에 _큰 커널 크기(e.g. kernel size = 14 or 16)_와 _non-overlapping convolution_을 사용함.

• 이에 맞춰 ResNet 스타일의 stem cell_((7 X 7), stride 2 convolution, max pool)을 patchify stem((4 X 4), stride 4 convolution)_으로 변경함.

2.3 ResNeXt-ify

• ResNeXt가 바닐라 ResNet보다 더 나은 FLOPs/accuracy trade-off를 지닌다는 점에 착안하여 그 구조를 참조함.

• 핵심 요소는 grouped convolution으로, ResNeXt에서는 3 X 3 conv layer에 grouped convolution을 사용함.

• 여기서는 depthwise convolution(그룹 수 = 채널 수)을 사용함. 이는 채널 단위로 수행되는 self-attention의 가중합 연산과 유사함(공간 차원의 정보만 혼합).

• 이어서 1 X 1 conv와 조합하면 공간과 채널의 혼합으로 분리되는데, ViT처럼 공간 및 채널 차원에 대해 정보를 혼합하되 둘 모두에 대해서 혼합하지는 않는 특성을 지닌다(depthwise conv에서 공간 정보를 혼합한 뒤 1X1 conv로 채널 정보 혼합).

• ResNeXt의 전략대로 네트워크의 너비를 증가시키되, Swin-T에 맞춰 채널 수를 64 -> 96으로 변경

2.4 Inverted Bottleneck

[그림-3] 논문 Figure3
(a) ResNeXt block, (b) inverted bottleneck, (c) depthwise conv layer 위로 이동

• 채널 수를 줄였다가 늘리는 기존의 bottleneck과 달리 채널 수를 늘렸다가 줄이는 방식이 inverted bottleneck이며, 확장 비율은 4배([그림-3]의 (b) 참고).
• MobileNetV2에 의해 대중화된 방법이며, 모든 Transformer block은 이 inverted bottleneck을 만들어낸다고 함.
• depthwise conv의 FLOPs가 상승했지만, downsampling하는 residual block의 shortcut 1X1 conv layer의 FLOPs가 대폭 감소하여 전체 네트워크의 FLOPs가 감소함. 하지만 성능은 약간 향상됨(80.5% -> 80.6%).

2.5 Large Kernel Sizes

• ViT의 가장 도드라지는 특징 중 하나는 바로 non-local self-attention이며, 이는 각 layer가 global receptive field를 지니도록 함.
• Swin-Transformer가 비록 self-attention block에 local window 개념을 다시 가져오긴 했지만, 그래도 window size는 최소 7 X 7로, ResNe(X)t의 3 X 3에 비하면 매우 큰 차이임.
• 따라서 VGGNet에 의해 ConvNet의 황금율이 된 kernel size 3 X 3을 깨고 다시 큰 kernel size를 살펴보고자 함.

1) Moving up depthwise conv layer

• 우선 depthwise conv layer를 위로 올림([그림-3] (b)에서 (c)로 변경). 이는 MSA block이 MLP block 이전에 위치하는 Transformer를 근거로 함.
• 복잡하고 효과적이지 않은 모듈(MSA, large-kernel conv)의 채널 수는 적게 할당하면서 효과적인 1 X 1 layers에는 힘을 줄 수 있음.

2) Increasing the kernel size

• kernel size는 7 X 7일 때가 가장 효과적.
• 3, 5, 7, 9, 11에 대해 시험을 해봤으며, 더 큰 커널 크기의 이점은 7 X 7일 때가 최대 수준이었음.

2.6 Micro Design

앞서 Macro Design이 block 및 네트워크 단위에서의 이야기였다면, 여기서는 layer 단위의 분석이 진행되며, activation function과 normalizaion layer에 초점을 맞춤.

1) Replacing ReLU with GELU

• NLP와 CV 아키텍쳐의 한 가지 차이점은 사용하는 activation function이 다르다는 점.
• 수많은 activation function의 발전이 있었음에도 불구하고, ConvNet에서는 단순하면서 효과적인 ReLU를 사용함.
• 하지만 BERT, GPT-2, ViT 등 최근의 진보된 Transformers에서는 모두 GELU(Gaussian Error Linear Unit)를 사용함.
• 변경해도 성능의 차이는 없으며, 적용해도 문제 없음을 확인함.

2) Fewer activation functions

[그림-4] 논문 Figure4 (activation function 마킹)
Swin Transformer, ResNet, ConvNeXt의 block 구조.
Transformer의 MLP block은 편의상 "1X1 convs"로 표기함.

• Transformer는 MLP block 통틀어서 activation function이 하나만 있는 반면에 ConvNet은 1X1를 포함한 모든 conv 레이어마다 activation function이 할당되어 있음.

• 따라서 두 1X1 conv layer 사이에 하나를 제외한 residual block 내 모든 GELU 레이어를 없앰.

3) Fewer normalization layers

[그림-5] 논문 Figure4 (normalization layer 마킹)

• 마찬가지로 noramalization layer도 ResNet에 비해 Swin Transformer가 더 적음.
• 따라서 1X1 conv 레이어 이전에 1개의 BN(Batch Normalization) 레이어만 배치하고 나머지는 제거함.
• 이때 성능은 Swin Transformer를 뛰어넘었으며, normalizaion layer의 개수도 더 적음.
• 경험적으로, block 시작 전에 BN 레이어를 하나 더 두는 것은 성능 개선에 기여하지 않음을 확인함.

4) Substituting BN with LN

• BN(Batch Normalization)은 ConvNet에서 수렴을 개선하고 과적합을 줄여준다는 점에서 애용하는 normalization 방법이나, Transformers에서는 더 단순한 LN(Layer Normalization)을 사용하고 더 나은 성능을 보임.
• 표준 ResNet에서 BN을 LN으로 직접 바꾸는 것에는 무리가 있겠으나, 이 모델(ConvNeXt)의 경우, 직접 바꿨을 때 성능이 소폭 향상함.

5) Separate downsampling layers

• Swin Transformer에서 downsampling layer를 각 단계와 분리시킨 아이디어를 차용함.
• 2X2 conv layers(stride 2)를 spatial downsampling에 사용하며, 앞에 LN이 선행되었을 때 더 안정적인 학습이 가능했음.
• 추가적으로 Swin Transformer를 따라 LN을 추가 배치하였으며, 위치는 다음과 같음.
• 각 downsampling layer 앞, stem 이후, global average pooling 이후.

3.1. Settings

아래의 데이터셋에 대해 훈련을 진행함

• ImageNet-1K (1000 클래스, 1.2M 훈련 이미지)
• ImageNet-22K (21841 클래스(ImageNet-1K 클래스의 슈퍼셋), 14M 사전학습 이미지, ImageNet-1K로 fine-tune)

1) Training on ImageNet-1K

• EMA(Exponential Moving Average)가 더 큰 모델의 과적합을 완화시키는 것을 확인함.

2) Pre-training on ImageNet-22K

• warmup은 5epoch만 진행
• EMA 미사용
• 그 외는 ImageNet-1K 때와 동일.

3) Fine-tuning on ImageNet-1K

• pre-training, fine-tuning, test 시 기본 해상도는 224 X 224.
• 추가적으로, 384 X 384의 해상도로 fine-tune 진행(ImageNet-22K와 ImageNet-1K pretrained models).
• ViTs/Swin Transformers와 비교했을 때 ConvNeXt는 다양한 해상도에 대해 fine-tune하기 훨씬 단순함. Fully-convolutional하고, input patch size를 맞추거나 position biases에 대해 보간(interpolate)할 필요도 없기 때문.

6. Conclusions

• 제시된 방법들은 완전히 새로운 것이 아니며, 다만 개별적으로 확인된 각 선택들을 종합적으로 고려한 것일 뿐임.
• 단순하면서도 효과적인 ConvNeXt를 제안함.

주요 특이사항

• Tensorflow 구현
• 싱글 GPU 기준으로 수정

아키텍쳐

AlexNet 구조도(2 GPU 기준)

Subclass형 모델 구현

import tensorflow as tf
from tensorflow.keras import datasets, layers, models


class AlexNet(tf.keras.Model):
    def __init__(self):
        super().__init__()
        initializer = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.01)

        self.alexnet = models.Sequential([
            layers.experimental.preprocessing.Resizing(227, 227, input_shape=train_images.shape[1:]),

            layers.Conv2D(filters=96, kernel_size=(11, 11), strides=4, padding='valid', activation='relu', kernel_initializer=initializer, name='Conv-1'), # 1층
            layers.Lambda(tf.nn.local_response_normalization, name='LRN-1'),
            layers.MaxPooling2D(pool_size=(3, 3), strides=2, padding='valid', name='Pooling-1'),

            layers.Conv2D(filters=256, kernel_size=(5, 5), strides=1, padding='same', activation='relu', kernel_initializer=initializer, bias_initializer='ones', name='Conv-2'), # 2층
            layers.Lambda(tf.nn.local_response_normalization, name='LRN-2'),
            layers.MaxPooling2D(pool_size=(3, 3), strides=2, padding='valid', name='Pooling-2'),

            layers.Conv2D(filters=384, kernel_size=(3, 3), strides=1, padding='same', activation='relu', kernel_initializer=initializer, name='Conv-3'), # 3층
            layers.Conv2D(filters=384, kernel_size=(3, 3), strides=1, padding='same', activation='relu', kernel_initializer=initializer, bias_initializer='ones', name='Conv-4'), # 4층
            layers.Conv2D(filters=256, kernel_size=(3, 3), strides=1, padding='same', activation='relu', kernel_initializer=initializer, bias_initializer='ones', name='Conv-5'), # 5층
            layers.MaxPooling2D(pool_size=(3, 3), strides=2, padding='valid', name='Pooling-3')
        ], name='Alexnet')

        self.fc = models.Sequential([
            layers.Flatten(name='Flatten'),
            layers.Dense(units=4096, activation='relu', kernel_initializer=initializer, bias_initializer='ones'), # hidden layer
            layers.Dropout(rate=0.5, name='Dropout-1'),
            layers.Dense(units=4096, activation='relu', kernel_initializer=initializer, bias_initializer='ones'), # hidden layer
            layers.Dropout(rate=0.5, name='Dropout-2'),
            layers.Dense(units=1000, activation='softmax', name='Classifier') # output layer
        ], name='Fully_connected')

    def call(self, x):
        return self.fc(self.alexnet(x))

    def build_model(self):
        x = tf.keras.layers.Input(shape=(227, 227, 3))
        model = tf.keras.Model(inputs=[x], outputs=self.call(x))
        return model

• Summary 확인

alexnet = AlexNet()
model = alexnet.build_model()
model.summary()

Model: "model_2"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 input_5 (InputLayer)        [(None, 227, 227, 3)]     0         

 Alexnet (Sequential)        (None, 6, 6, 256)         3747200   

 Fully_connected (Sequentia  (None, 1000)              58631144  
 l)                                                              

=================================================================
Total params: 62378344 (237.95 MB)
Trainable params: 62378344 (237.95 MB)
Non-trainable params: 0 (0.00 Byte)

• plot_model

  tf.keras.utils.plot_model(model)

  plot_model(subclass)

Sequential형 모델 구현

initializer = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.01)

model = models.Sequential()

# 227 x 227 x 3으로 resize
model.add(layers.experimental.preprocessing.Resizing(227, 227, input_shape=train_images.shape[1:]))

model.add(layers.Conv2D(filters=96, kernel_size=(11, 11), strides=4, padding='valid', activation='relu', kernel_initializer=initializer, name='Conv-1')) # 1층
model.add(layers.Lambda(tf.nn.local_response_normalization, name='LRN-1'))
model.add(layers.MaxPooling2D(pool_size=(3, 3), strides=2, padding='valid', name='Pooling-1'))

model.add(layers.Conv2D(filters=256, kernel_size=(5, 5), strides=1, padding='same', activation='relu', kernel_initializer=initializer, bias_initializer='ones', name='Conv-2')) # 2층
model.add(layers.Lambda(tf.nn.local_response_normalization, name='LRN-2'))
model.add(layers.MaxPooling2D(pool_size=(3, 3), strides=2, padding='valid', name='Pooling-2'))

model.add(layers.Conv2D(filters=384, kernel_size=(3, 3), strides=1, padding='same', activation='relu', kernel_initializer=initializer, name='Conv-3')) # 3층
model.add(layers.Conv2D(filters=384, kernel_size=(3, 3), strides=1, padding='same', activation='relu', kernel_initializer=initializer, bias_initializer='ones', name='Conv-4')) # 4층
model.add(layers.Conv2D(filters=256, kernel_size=(3, 3), strides=1, padding='same', activation='relu', kernel_initializer=initializer, bias_initializer='ones', name='Conv-5')) # 5층
model.add(layers.MaxPooling2D(pool_size=(3, 3), strides=2, padding='valid', name='Pooling-3'))

# FC layer
model.add(layers.Flatten(name='Flatten')) 
model.add(layers.Dense(units=4096, activation='relu', kernel_initializer=initializer, bias_initializer='ones')) # hidden layer
model.add(layers.Dropout(rate=0.5, name='Dropout-1'))
model.add(layers.Dense(units=4096, activation='relu', kernel_initializer=initializer, bias_initializer='ones')) # hidden layer
model.add(layers.Dropout(rate=0.5, name='Dropout-2'))
model.add(layers.Dense(units=1000, activation='softmax', name='Classifier')) # output layer

• summary 확인

  Model: "sequential"

Layer (type) Output Shape Param #

resizing (Resizing) (None, 227, 227, 3) 0

Conv-1 (Conv2D) (None, 55, 55, 96) 34944

LRN-1 (Lambda) (None, 55, 55, 96) 0

Pooling-1 (MaxPooling2D) (None, 27, 27, 96) 0

Conv-2 (Conv2D) (None, 27, 27, 256) 614656

LRN-2 (Lambda) (None, 27, 27, 256) 0

Pooling-2 (MaxPooling2D) (None, 13, 13, 256) 0

Conv-3 (Conv2D) (None, 13, 13, 384) 885120

Conv-4 (Conv2D) (None, 13, 13, 384) 1327488

Conv-5 (Conv2D) (None, 13, 13, 256) 884992

Pooling-3 (MaxPooling2D) (None, 6, 6, 256) 0

Flatten (Flatten) (None, 9216) 0

dense (Dense) (None, 4096) 37752832

Dropout-1 (Dropout) (None, 4096) 0

dense_1 (Dense) (None, 4096) 16781312

Dropout-2 (Dropout) (None, 4096) 0

Classifier (Dense) (None, 1000) 4097000

================================================================= Total params: 62378344 (237.95 MB) Trainable params: 62378344 (237.95 MB) Non-trainable params: 0 (0.00 Byte)

- plot_model
<p align="center" style="color:gray;">
  <img src="https://velog.velcdn.com/images/leo-moooon/post/85145257-9adf-4cb4-b829-a042a862e100/image.png" alt="plot_model(sequential)" style="margin:0;">
  plot_model(sequential)
</p>