ViTPain: Vision Transformer for Pain Assessment

• ViTPain is a cross-modal framework that employs a ViT-Large backbone with dual branches for PSPI classification and AU regression, using teacher-student distillation.
• It integrates AU-specific query cross-attention and synthetic 3DPain data to address label imbalances and ensure clinical reliability in automated facial pain assessment.
• The framework achieves robust performance improvements, with AUROC up to 0.91 on real datasets, thereby effectively transferring spatial details from heatmap supervision.

Searching arXiv for the cited ViTPain and related pain-transformer papers to ground the article in current literature. ViTPain is a Vision Transformer based cross-modal distillation framework for automated facial pain assessment, introduced alongside the 3DPain synthetic dataset in “Pain in 3D: Generating Controllable Synthetic Faces for Automated Pain Assessment” (Lin et al., 20 Sep 2025). It is designed to address two constraints that have limited progress in automated pain assessment: severe demographic and label imbalance in existing datasets, and the lack of generative pipelines that can precisely control facial action units, facial structure, or clinically validated pain levels. In this formulation, a heatmap-trained teacher guides a student trained on RGB images, with the stated objective of enhancing accuracy, interpretability, and clinical reliability (Lin et al., 20 Sep 2025).

1. Definition and research context

ViTPain is situated within automated pain assessment from facial expressions, a problem motivated by the need to assess pain in non-communicative patients, such as those with dementia (Lin et al., 20 Sep 2025). The framework combines a ViT-Large backbone, dual prediction heads for PSPI classification and AU regression, and a teacher-student training scheme that transfers knowledge from AU heatmaps to RGB images (Lin et al., 20 Sep 2025). Its development is tied directly to 3DPain, a synthetic dataset comprising 82,500 samples, 25,000 pain expression heatmaps, and 2,500 synthetic identities balanced by age, gender, and ethnicity (Lin et al., 20 Sep 2025).

The immediate antecedent in transformer-based pain detection is the fully-attentive pipeline introduced for binary pain detection from facial expressions in “Fully-attentive and interpretable: vision and video vision transformers for pain detection” (Fiorentini et al., 2022). That work proposed ViT and ViViT models trained on the UNBC-McMaster dataset after 3D face registration and frontalization, and reported state-of-the-art performance on binary pain detection from facial expressions (Fiorentini et al., 2022). ViTPain extends this line of work from binary pain detection toward clinically structured estimation using PSPI scores, AU supervision, synthetic data, and cross-modal distillation (Lin et al., 20 Sep 2025).

A recurring source of ambiguity is nomenclature. ViTPain is unrelated to “Visual insTruction Pretraining” (ViTP), which is a domain-specific foundation model pretraining paradigm for remote sensing and medical imaging (Li et al., 22 Sep 2025). The shared “ViT” prefix reflects the use of Vision Transformers, but the methods, objectives, and application domains are distinct.

2. Architectural formulation

ViTPain uses a Vision Transformer backbone specified as ViT-Large (Lin et al., 20 Sep 2025). The input is a $224 \times 224$ RGB image, patchified into $N=196$ tokens of dimension $D=1024$, with standard position embeddings and a classification token for global representation (Lin et al., 20 Sep 2025).

The architecture includes a dual-branch head. The first branch performs 17-class classification over PSPI pain intensity scores from 0 to 16, producing $\hat{\mathbf{y}_{\text{PSPI}}}$ (Lin et al., 20 Sep 2025). The second branch performs regression for the 6 most pain-relevant Action Units, identified as AU4, AU6, AU7, AU9, AU10, and AU43, following the PSPI formulation (Lin et al., 20 Sep 2025). This dual structure reflects the paper’s attempt to align prediction with clinically validated pain coding rather than relying only on an end-task scalar score.

A distinctive component is AU-Specific Query Cross-Attention. Instead of depending only on the CLS token, the model introduces learnable AU query tokens $\mathbf{Q}_{AU} \in \mathbb{R}^{6 \times D}$, each intended to focus on spatial regions relevant for a particular AU (Lin et al., 20 Sep 2025). The attention mechanism is defined as:

$$\mathbf{A}_{i,j} = \frac{\mathbf{Q}_{AU_i} \mathbf{P}_j^T}{\sqrt{D}}$$

$$\alpha_{i,j} = \frac{\exp(\mathbf{A}_{i,j})}{\sum_{k=1}^{N} \exp(\mathbf{A}_{i,k})}$$

$$\mathbf{F}_{AU_i} = \sum_{j=1}^{N} \alpha_{i,j} \mathbf{P}_j$$

with AU prediction given by

$\hat{\mathbf{y}_{AU} = \text{ReLU}(\text{Linear}(\mathbf{F}_{AU}))$

as reported in the source description (Lin et al., 20 Sep 2025). The stated interpretability function of this design is that each query token’s attention map highlights the facial area associated with its AU, allowing inspection of what facial regions contributed to pain estimation (Lin et al., 20 Sep 2025).

This design differs from the earlier ViT and ViViT pain detection pipeline, where a new classification head was added for binary classification, pain versus no pain, and where spatiotemporal modeling was handled through ViViT on 2×2 grids of sequential frames (Fiorentini et al., 2022). ViTPain therefore shifts the task from binary discrimination toward structured pain representation.

3. Teacher-student distillation and learning objectives

ViTPain is explicitly formulated as a multi-modal teacher-student system (Lin et al., 20 Sep 2025). The teacher model is trained on AU heatmaps and is described as achieving the best performance, but real clinical images lack such heatmaps. The student model is trained on RGB images and optimized to mimic the teacher’s outputs and internal representations (Lin et al., 20 Sep 2025). This is the sense in which the framework is “cross-modal”: the source modality for privileged supervision is spatially structured AU heatmaps, while the deployment modality is RGB imagery.

The distillation procedure includes three stated components. For PSPI output distillation, the loss is

$\mathcal{L}_{\text{PSPI}^{\text{distill} = T^2 \cdot \mathrm{KL}\left(\mathrm{softmax}\left(\frac{\mathbf{z}_t}{T}\right) \,\Vert\, \mathrm{softmax}\left(\frac{\mathbf{z}_s}{T}\right)\right)$

where $N=196$0 and $N=196$1 are logits from student and teacher, and $N=196$2 is the softmax temperature (Lin et al., 20 Sep 2025). AU output distillation is defined as

$N=196$3

and feature-level distillation on the global representation is

$N=196$4

again following the provided formulation (Lin et al., 20 Sep 2025).

The total objective is reported as

\begin{align} \mathcal{L}{total} &= \lambda{PSPI}\, \mathcal{L}{CE}(\hat{\mathbf{y}{PSPI}, \mathbf{y}{PSPI}) \nonumber \ &\quad + \lambda{AU}\, \mathcal{L}{MSE}(\hat{\mathbf{y}{AU}, \mathbf{y}{AU}) \nonumber \ &\quad + \lambda{PSPI}^{distill}\, \mathcal{L}{PSPI}^{distill} \nonumber \ &\quad + \lambda{AU}^{distill}\, \mathcal{L}{AU}^{distill} \nonumber \ &\quad + \lambda{feature}^{distill}\, \mathcal{L}_{feature}^{distill} \end{align}

with loss weights selected via validation as $N=196$5, $N=196$6, $N=196$7, $N=196$8, and $N=196$9, and with distillation temperature $D=1024$0 (Lin et al., 20 Sep 2025).

A plausible implication is that ViTPain operationalizes privileged supervision: the teacher exploits a modality unavailable at inference time, while the student is constrained to match both semantic outputs and internal representations. The paper’s stated claim is that this improves generalization and interpretability by transferring detailed structure from heatmap teacher to RGB student (Lin et al., 20 Sep 2025).

4. Clinical targets, labels, and dataset substrate

The clinical target underlying ViTPain is PSPI, a clinically validated pain score (Lin et al., 20 Sep 2025). The reported formula is

$D=1024$1

which is also the formulation used in the earlier transformer pain-detection study for thresholding frames into “pain” and “no pain” in binary classification (Fiorentini et al., 2022). In ViTPain, however, PSPI is not merely a thresholding device; it becomes a 17-class target from 0 to 16, accompanied by explicit AU regression (Lin et al., 20 Sep 2025).

3DPain is the dataset foundation for ViTPain. The dataset is described as large-scale, synthetic, and specifically designed for automated pain assessment (Lin et al., 20 Sep 2025). It is generated through a three-stage framework that produces diverse 3D meshes, textures them with diffusion models, and applies AU-driven face rigging to synthesize multi-view faces with paired neutral and pain images, AU configurations, PSPI scores, and dataset-level annotations of pain-region heatmaps (Lin et al., 20 Sep 2025). Its scale and annotations are central to the method’s claims.

Component	Reported specification	Role
Samples	82,500 samples	Synthetic training data
Heatmaps	25,000 pain expression heatmaps	Teacher supervision and spatial annotations
Identities	2,500 synthetic identities	Demographic diversity
Balance	Balanced by age, gender, and ethnicity	Mitigates demographic skew
Labels	AU configurations, PSPI scores, paired neutral and pain images	Structured supervision

The stated motivation is that existing datasets exhibit severe demographic and label imbalance due to ethical constraints, while existing generative models cannot precisely control AUs, facial structure, or clinically validated pain levels (Lin et al., 20 Sep 2025). This suggests that ViTPain should be understood as inseparable from 3DPain: the model’s architecture, supervision, and intended generalization properties are all tied to the availability of synthetic multi-view, AU-controlled, demographically balanced data.

For comparison, the earlier UNBC-McMaster-based transformer study worked with a real dataset containing more than 48,000 video frames from 25 patients with framewise PSPI scores, using binary labels, oversampling of the minority pain class, and patient-disjoint 5-fold cross-validation (Fiorentini et al., 2022). ViTPain’s use of synthetic data addresses limitations explicitly identified in that broader literature.

5. Training regimen and empirical results

The reported training regimen for ViTPain starts with a ViT pretrained on large-scale data and then trains or fine-tunes on synthetic 3DPain plus real data (Lin et al., 20 Sep 2025). Optimization uses AdamW; the backbone is frozen for the first 5 epochs, followed by end-to-end fine-tuning; and cosine learning rate annealing is used (Lin et al., 20 Sep 2025). Teacher and student are batched separately (Lin et al., 20 Sep 2025).

On the real UNBC-McMaster dataset, the paper reports three main configurations: a ViTPain baseline with AUROC 0.83 when the student is trained on real data only; ViTPain + 3DPain Synthetic Data with AUROC 0.90; and ViTPain + 3DPain + Heatmap Supervision (Full) with AUROC 0.91 and F1 = 0.54 (Lin et al., 20 Sep 2025). The comparison table in the source lists SVM [Lucey11] at AUROC 0.84, CNN-LSTM [Rodriguez17] at AUROC 0.93*, Deep Face Recog [Parkhi15] at F1 0.59, SVR with DCT features [Kaltwang12] at F1 0.48, AFAR [Ertugrul19] at F1 0.59, and Pairwise Contrastive [Rezaei21] at AUROC 0.86 and F1 0.56 (Lin et al., 20 Sep 2025). The note attached to CNN-LSTM states that it used different splits, whereas ViTPain uses 5-fold CV, subject-independent (Lin et al., 20 Sep 2025).

On synthetic 3DPain ablations, the paper reports the following progression: Baseline ViTPain yields Macro AUROC 0.81, PSPI acc. 0.24, $D=1024$2 acc. 0.60, and $D=1024$3 acc. 0.79; adding AU-Query yields Macro AUROC 0.82 and $D=1024$4 acc. 0.81; adding Heatmap Distillation yields Macro AUROC 0.85, PSPI acc. 0.25, $D=1024$5 acc. 0.63, and $D=1024$6 acc. 0.83; and the Teacher (Heatmap only) attains Macro AUROC 0.96 and PSPI acc. 0.67 (Lin et al., 20 Sep 2025). The interpretation given is that the heatmap teacher sets the upper bound and that the best RGB student closely approaches it (Lin et al., 20 Sep 2025).

These results can be situated against the earlier transformer baseline literature. The 2022 fully-attentive pain detection paper reported ViT-1 with F1 score 0.55 ± 0.15 and AUC 0.88, ViViT-1 with F1 score 0.55 ± 0.13 and AUC 0.86, and ViViT-2 with F1 score 0.49 ± 0.04 and AUC 0.76, all on binary pain detection (Fiorentini et al., 2022). Because the task formulations differ, direct numerical comparison should be treated cautiously. What can be stated directly is that ViTPain reports AUROC 0.91 and F1 0.54 on UNBC-McMaster under subject-independent 5-fold CV in its full configuration (Lin et al., 20 Sep 2025), while the earlier study reported state-of-the-art binary detection performance for ViT and ViViT under its own pipeline (Fiorentini et al., 2022).

6. Interpretability, clinical reliability, and limitations

Interpretability is an explicit design objective in ViTPain. The AU query structure is described as allowing visualization of which facial area or action unit drove the decision, and the paper presents this as matching clinical heuristics and criteria (Lin et al., 20 Sep 2025). Clinical reliability is framed through the use of PSPI, identified as a clinically validated score, and through regression of the same metric from interpretable features, namely AUs (Lin et al., 20 Sep 2025).

This emphasis is continuous with the earlier fully-attentive transformer paper, which analyzed attention maps and attention rollout in ViViT and found activations around known pain-relevant facial muscles, including brow, eyes, and cheeks, corresponding to PSPI AUs such as AU4, AU6/AU7, and AU43 (Fiorentini et al., 2022). That earlier work argued that transformer attention provided clinically plausible interpretations for predictions (Fiorentini et al., 2022). ViTPain retains that interpretability agenda but embeds it more deeply in the architecture through AU-specific query tokens and heatmap distillation (Lin et al., 20 Sep 2025).

The method also addresses dataset bias and scarcity through synthetic augmentation and demographic balancing (Lin et al., 20 Sep 2025). The paper states that this makes the model less likely to overfit or perform poorly on underrepresented groups (Lin et al., 20 Sep 2025). Because this is framed as an intended effect rather than a fully enumerated fairness audit, a cautious reading is appropriate: this suggests a pathway toward broader generalization, but the paper excerpt does not provide subgroup performance tables.

Several boundaries are also evident from the reported setup. Real clinical images lack heatmaps, so the teacher’s strongest modality is unavailable at inference time, which is why student distillation is necessary (Lin et al., 20 Sep 2025). The teacher’s performance on synthetic data, with Macro AUROC 0.96 and PSPI acc. 0.67, exceeds the student’s, indicating a remaining gap between privileged and deployable modalities (Lin et al., 20 Sep 2025). A plausible implication is that ViTPain’s practical performance depends on how effectively distillation can transfer spatial priors without direct access to heatmap supervision at test time.

A further misconception worth avoiding is to equate ViTPain with generic ViT usage in pain detection. The 2022 study established that pure transformer architectures could outperform earlier CNN-based approaches in binary pain detection from facial expressions (Fiorentini et al., 2022). ViTPain is more specific: it couples a ViT backbone to AU-query cross-attention, PSPI classification, AU regression, synthetic 3D data generation, and cross-modal teacher-student distillation (Lin et al., 20 Sep 2025). It is therefore a distinct framework rather than merely another ViT baseline.