Vision Transformers

Deep Learning and Computer Vision in R: A Practical Introduction

SSA ACT workshop - Vision Transformers

Weihao (Patrick) Li

RSFAS, ANU

Content summary

• Overview of computer vision (CV)
• reticulate and torch basics
• Image classification, neural networks and Keras
• Convolutional neural networks
• > Vision transformers
• Object detection
• Image segmentation

What a convolutional net cannot see

A 3 \times 3 convolution sees only a small local region. Each block expands the receptive field while halving the spatial resolution. By the time a unit sees the entire 32 \times 32 image, the feature map has shrunk to 4 \times 4, compressing spatial detail into a highly summarised representation. CNNs trade spatial detail for broader context.

Vision transformer

ViT splits an image into patches, treating each patch like a word in a sentence. It then uses attention to learn how different patches relate to each other, before using the final representation to predict the image class.

In essence, ViT applies the transformer architecture from language models to images, replacing words with image patches.

Patch embedding

Patchification

• Split the image into fixed-size, non-overlapping patches (e.g., 16 × 16).
• A 224 × 224 image produces 196 patches.

Patch embedding

• Linearly project each patch into a d-dimensional vector (token).
• Output shape: (N_tokens, d).
• Equivalently, use a convolutional layer with d filters and a stride equal to the patch size.

Patchification (ViT) — Dosovitskiy et al. (2020), An Image is Worth 16×16 Words.

Self-attention

Self-attention allows each token to update its representation using information from other tokens, capturing relationships across the image.

For a token embedding matrix X (N_{tokens} \times d),

Q = XW_Q,\qquad K = XW_K,\qquad V = XW_V

Attention matrix:

A = \mathrm{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)

Output:

\mathrm{Attention}(Q,K,V) = AV

• Q: queries (“what am I looking for?”)
• K: keys (“what information is available?”)
• V: values (“the information to retrieve”)
• Each token attends to other relevant tokens.

An example of attention matrix visualization — Vaswani et al. (2017), Attention Is All You Need.

The encoder

Transformer block

A transformer is built by stacking many transformer blocks, each containing

1. Layer normalization
2. multi-head self-attention (each head has d/N_{heads})
3. Residual addition
4. Layer normalization
5. Multiple dense layers (d \to 4d \to d)
6. Residual addition

The output of one block is the input of the next. ViT-B/16 stacks 12 of them, with d = 768 and 12 heads, 86M parameters.

An example of a transformer block — Dosovitskiy et al. (2020), An Image is Worth 16×16 Words.

The class token

So far, every token represents an image patch. But the classifier needs one vector to represent the whole image.

• Prepend a learnable [CLS] token immediately after patch embedding.
• Through self-attention, the [CLS] token gathers information from all image tokens.
• After the final transformer block, pass only the [CLS] token to the classification head.

Later models often omit the class token and use mean-pooling to average the patch tokens instead. This is computationally simpler and naturally supports tasks that produce spatial outputs, such as segmentation masks.

Position embedding

You just saw that attention matrix compares tokens with one another but has no information about their original positions. Shuffling the tokens simply shuffles the outputs.

Attention is permutation-invariant, so we need to tell each token where it came from.

• Add a learnable positional embedding (a N_{tokens} \times d matrix) to each token, after prepending [CLS] token.

Without positional embeddings, the model becomes a bag-of-patches classifier: it may recognise a cat or a dog, but it loses information about where things are in the image.

Vision transformer variants

Vision transformer variants differ along four design decisions:

decision	options	examples
Tokenisation	fixed patches · hierarchical merging · convolutional stem	ViT · Swin · CoAtNet
Positional information	learned absolute · sinusoidal · relative bias · rotary	ViT · Transformer · Swin · RoPE
Attention scope	global · local windows · local then global	ViT · Swin · Hiera
Classification	[CLS] token · mean pooling · distillation token	ViT · Hiera · DeiT

CNN vs. ViT

	CNN	ViT
Inductive bias	locality, translation equivariance	relatively little
Global context	receptive field grows with depth	every attention layer
Strengths	limited data, smaller datasets	large datasets, large-scale pretraining
Multimodal integration (e.g. text, audio)	Typically requires modality-specific architectures	Naturally supports a shared architecture across modalities

A pretrained ViT in R

• Code
• torchinfo

# Wrap a torch model as a Keras layer
vit_layer <- function() {
  vit <- torchvision$models$vit_b_16(weights = "DEFAULT") # Download the ViT from torchvision
  vit$heads$head <- torch$nn$Identity() # Remove the original classification head
  keras$layers$TorchModuleWrapper(vit, output_shape = tuple(NULL, 768L)) # The CLS token has 768 dimensions
}

model <- keras$Sequential(list(
  keras$layers$Input(shape = tuple(224L, 224L, 3L)),
  keras$layers$Permute(tuple(3L, 1L, 2L)),        # NHWC -> NCHW for torch
  vit_layer(),
  keras$layers$Dense(1L, activation = "sigmoid")
))

invisible(model$summary())

Model: "sequential"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type)                    ┃ Output Shape           ┃       Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ permute (Permute)               │ (None, 3, 224, 224)    │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ torch_module_wrapper            │ (None, 768)            │    85,798,656 │
│ (TorchModuleWrapper)            │                        │               │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense (Dense)                   │ (None, 1)              │           769 │
└─────────────────────────────────┴────────────────────────┴───────────────┘
 Total params: 85,799,425 (327.30 MB)
 Trainable params: 85,799,425 (327.30 MB)
 Non-trainable params: 0 (0.00 B)

torchinfo <- import("torchinfo", convert = FALSE)
torchinfo$summary(model$layers[1]$module, verbose = 0L)

==========================================================================================
Layer (type:depth-idx)                                            Param #
==========================================================================================
VisionTransformer                                                 768
├─Conv2d: 1-1                                                     590,592
├─Encoder: 1-2                                                    151,296
│    └─Dropout: 2-1                                               --
│    └─Sequential: 2-2                                            --
│    │    └─EncoderBlock: 3-1                                     7,087,872
│    │    └─EncoderBlock: 3-2                                     7,087,872
│    │    └─EncoderBlock: 3-3                                     7,087,872
│    │    └─EncoderBlock: 3-4                                     7,087,872
│    │    └─EncoderBlock: 3-5                                     7,087,872
│    │    └─EncoderBlock: 3-6                                     7,087,872
│    │    └─EncoderBlock: 3-7                                     7,087,872
│    │    └─EncoderBlock: 3-8                                     7,087,872
│    │    └─EncoderBlock: 3-9                                     7,087,872
│    │    └─EncoderBlock: 3-10                                    7,087,872
│    │    └─EncoderBlock: 3-11                                    7,087,872
│    │    └─EncoderBlock: 3-12                                    7,087,872
│    └─LayerNorm: 2-3                                             1,536
├─Sequential: 1-3                                                 --
│    └─Identity: 2-4                                              --
==========================================================================================
Total params: 85,798,656
Trainable params: 85,798,656
Non-trainable params: 0
==========================================================================================

Exercises

06-vision-transformers-ex.R — 15 minutes (5 to walk through, 10 to work)

1. The script loads a pretrained ViT and drops its 1000-class head, so the encoder returns one 768-dimensional vector per image.
2. It wraps that encoder as a Keras layer and stacks the resize, crop, normalise and permute steps in front of it — read those four lines before you run anything.
3. One thing is left as ???: freeze the backbone so that only the two dense layers learn. Check model$summary() before and after — the trainable count should fall from 86M to about 197k.
4. Train for 3 epochs, then compare accuracy and the ROC curve with the convolutional network from exercise 04.