1. Mở đầu

Như chúng ta đã biết thì các phương pháp semi-supervised learning đã góp công không nhỏ trong việc cải thiện hơn nữa các model state-of-the-art trong rất nhiều computer vision tasks như image classification, object detection, và semantic segmentation. Các phương pháp như Pseudo Labels hay self-training chắc cũng khá quen thuộc với những người từng làm về semi-supervised learning. Hôm nay mình sẽ giới thiệu đến các bạn một phiên bản nâng cấp của Pseudo Labels, giúp đưa semi-supervised learning lên đỉnh của Imagenet. Nói sơ qua thì cách hoạt động của Pseudo Label khá đơn giản: chúng ta cần 2 model, một gọi là teacher và một là student. Đầu tiên, ta cần huấn luyện teacher model với dữ liệu có nhãn, sau đó sử dụng model teacher để predict ra nhãn giả - pseudo label của dữ liệu chưa có nhãn, từ bây giờ ta sẽ gọi dữ liệu có nhãn chuẩn là labeled data và dữ liệu được sinh nhãn giả là pseudo data cho ngắn gọn nhé. Pseudo data sẽ được kết hợp với labeled data để huấn luyện cho model student, nhờ có sự bổ sung này mà student có thể mang lại kết quả tốt hơn so với teacher.

Mặc dù phương pháp kể trên khá hiệu quả nhưng vẫn tồn tại một nhược điểm lớn là sẽ có những pseudo label mà teacher sinh ra không chính xác, kéo theo student cũng sẽ học từ dữ liệu sai lệch đấy và kết quả là performance của student bị giảm sút. Điểm yếu này được gọi là confirmation bias trong pseudo-labeling.

Paper mà hôm nay mình muốn thảo luận với mọi người là một phiên bản nâng cấp xịn xò của Pseudo Label - Meta Pseudo Labels. Những gì mà Meta Pseudo Label muốn làm là cải thiện nhược điểm kể trên của teacher thông qua việc quan sát pseudo label mà nó sinh ra sẽ có ảnh hưởng gì đến student, nghĩa là nó sẽ nhận lại feedback của student sau khi học từ pseudo label và tự chỉnh sửa lại bản thân để cho ra những phiên bản pseudo label tốt hơn. Vì phần chứng minh của paper khá nhiều toán nên mình sẽ cố gắng đi chậm, nếu các bạn phát hiện ra mình sai thì đừng ngần ngại góp ý nhé 😁

2. Meta Pseudo Labels

Trái: Pseudo Labels, teacher được cố định sau khi train với labeled data sau đó sinh pseudo label cho student học. Phải: Meta Pseudo Labels, teacher được train song song với student.

Ký hiệu:

• $T, S$ : mô hình teacher và student
• $\theta_t,\theta_s$ : tham số của teacher và student
• $\theta_s^{PL}$: tham số của mô hình student được train với pseudo label tạo bởi teacher
• $|T|, |S|$ : dimension của student và teacher
• $(x_L,y_L)$ : labeled data gồm image và label
• $x_U$ : unlabeled data chỉ gồm image
• $T(x_U,\theta_T)$: soft prediction của teacher với unlabeled data
• $S(x_U,\theta_S), S(x_L,\theta_S)$: soft prediction của student với $x_U$ và $x_L$
• $\operatorname{CE}(q,p)$: cross-entropy loss giữa 2 phân phối q, p với q là label
• $\mathbb{E}_x[f]$ : giá trị kỳ vọng của phương trình f với biến ngẫu nhiên x.
• $\nabla$ : gradient

2.1 Revisit Pseudo Label

Trước khi nói về Meta Pseudo Label, ta sẽ mở đầu với việc ôn lại 1 chút về Pseudo Label nhé. Như đã giới thiệu ở phần mở đầu, Pseudo Label huấn luyện student model với unlabeled data để tối thiểu hóa hàm cross-entropy:

$$\theta_{S}^{\mathrm{PL}}=\underset{\theta_{S}}{\operatorname{argmin}} \underbrace{\mathbb{E}_{x_{u}}\left[\operatorname{CE}\left(T\left(x_{u} ; \theta_{T}\right), S\left(x_{u} ; \theta_{S}\right)\right)\right]}_{:=\mathcal{L}_{u}\left(\theta_{T}, \theta_{S}\right)} \tag{1}$$

• $\mathcal{L}_{u}\left(\theta_{T}, \theta_{S}\right)$ : loss của student khi train với pseudo label tạo bởi teacher trên unlabeled data.

Giả sử ta đã có model teacher được train tốt với tập labeled, mục tiêu của pseudo label là tạo ra $\theta_s^{PL}$ tối ưu trên tập labeled data:

$$\mathbb{E}_{x_{l}, y_{l}}\left[\operatorname{CE}\left(y_{l}, S\left(x_{l} ; \theta_{S}^{\mathrm{PL}}\right)\right)\right]:=\mathcal{L}_{l}\left(\theta_{S}^{\mathrm{PL}}\right) \tag{2}$$

2.2 Solution for confirmation-bias

Với Pseudo Labels, muốn cho student $\theta_s^{PL}$ tối ưu thì bắt buộc phải phụ thuộc vào teacher $\theta_T$ thông qua pseudo label $T(x_U,\theta_T)$. Để miêu tả sự phụ thuộc này ta sẽ dùng ký hiệu $\theta_{S}^{\mathrm{PL}}(\theta_{T})$. Như vậy hàm loss của student trên labeled data có thể được viết gọn lại như sau: $\mathcal{L}_{l}(\theta_{S}^{\mathrm{PL}}(\theta_{T}))$ và tất nhiên nhiệm vụ của hàm này sẽ là tối ưu 2 tham số $\theta_{S}^{\mathrm{PL}}$ và $\theta_{T}$. Từ đó ta có thể tối ưu hóa $\mathcal{L}_{l}$ theo $\theta_T$ như sau:

$$\begin{aligned}\min _{\theta_{T}} & \mathcal{L}_{l}\left(\theta_{S}^{\mathrm{PL}}\left(\theta_{T}\right)\right), \\\text { trong đó } & \theta_{S}^{\mathrm{PL}}\left(\theta_{T}\right)=\underset{\theta_{S}}{\operatorname{argmin}}\; \mathcal{L}_{u}\left(\theta_{T}, \theta_{S}\right)\end{aligned} \tag{3}$$

Theo như công thức trên thì ta có thể tối ưu hóa teacher thông qua biểu hiện của student, từ đó pseudo label dùng để train student cũng sẽ dần được cải thiện. Tuy nhiên do mối phụ thuộc $\theta_{S}^{\mathrm{PL}}(\theta_{T})$ và $\theta_T$ là vô cùng phức tạp nên việc tính gradient $\nabla_{\theta_{T}}(\theta_{S}^{\mathrm{PL}}(\theta_{T}))$ nếu muốn diễn ra thì bắt buộc phải thay đổi toàn bộ quá trình training của student.

Để đơn giản hóa việc này, ta sẽ áp dụng ý tưởng của meta-learning : xấp xỉ $\underset{\theta_{S}}{\operatorname{argmin}}$ bằng cách update từng bước gradient của $\theta_T$:

$$\theta_{S}^{\mathrm{PL}}\left(\theta_{T}\right) \approx \theta_{S}-\eta_{S} \cdot \nabla_{\theta_{S}} \mathcal{L}_{u}\left(\theta_{T}, \theta_{S}\right)\tag{4}$$

$$\text{ với } \eta_S \text{ là learning rate của student }$$

Thay biểu thức trên vào phương trình (3) ta sẽ có hàm tối ưu của teacher trong Meta Pseudo Labels:

$$\min _{\theta_{T}} \quad \mathcal{L}_{l}\left(\theta_{S}-\eta_{S} \cdot \nabla_{\theta_{S}} \mathcal{L}_{u}\left(\theta_{T}, \theta_{S}\right)\right) \tag{5}$$

Về cơ bản thì quá trình training của student vẫn phụ thuộc vào phương trình (1) của Pseudo Labels, ngoại trừ việc tham số của teacher sẽ không còn cố định mà thay đổi dần dựa vào student. Từ đó chúng ta sẽ rút ra được quá trình tối ưu hóa song song teacher - student:

• Student: sử dụng pseudo label từ teacher - $T(x_U,\theta_T)$ để tối ưu hóa hàm mục tiêu với SGD:

$$\theta_S^{\prime}=\theta_S-\eta_{S} \cdot\nabla_{\theta_{S}} \mathcal{L}_{u}\left(\theta_{T}, \theta_{S}\right)\tag{6}$$

• Teacher: sử dụng labeled data kết hợp với feedback của student để cải thiện pseudo label và tối ưu hóa hàm mục tiêu với SGD:

$$\theta_T^{\prime}=\theta_T-\eta_{T} \cdot\nabla_{\theta_{T}} \mathcal{L}_{l}\left(\theta_S-\nabla_{\theta_{S}} \mathcal{L}_{u}\left(\theta_{T}, \theta_{S}\right)\right)\tag{7}$$

2.3 Teacher's auxiliary losses

Các tác giả thấy rằng Meta Pseudo Labels tự thân nó đã khá tốt rồi, tuy nhiên nếu thêm một task phụ vào quá trình training của teacher thì performance sẽ còn tốt hơn. Do đó khi train teacher với labeled data, ta có thể thêm một auxiliary task dạng self-supervised để tận dụng unlabeled data giúp tăng độ generalization của model teacher. Auxiliary task này được thực hiện theo paper UDA (Unsupervised Data Augmentation for Consistency Training) với tổng quan như sau:

Ta có thể mô tả 1 cách đơn giản về UDA như sau:

• B1 : Với labeled data $(x,y)$, ta để model predict label $\hat{y} = P_{\theta}(y|x)$ và tính supervised loss $L_{sup}= \operatorname{CE}(y, \hat{y})$
• B2 : Với unlabeled data $(x)$, ta tiến hành augment $x$ để có $\hat{x}$, sau đó để model predict label cho $x$ và $\hat{x}$ : $P_{\theta}(y|x)$ và $P_{\theta}(y|\hat{x})$ rồi tính unsupervised loss với 2 label trên : $L_{unsup}(P_{\theta}(y|x), P_{\theta}(y|\hat{x}))$
• B3: tính loss tổng : $L_{final}=L_{sup} + \alpha\cdot L_{unsup}$ và optimize model dựa trên loss tổng

2.4 Derivation of the Teacher’s Update Rule

Nhắc lại một số ký hiệu toán học:

• cho hàm khả vi $f : R ^ { m } \rightarrow R^ {n}, x \mapsto f(x) , x \in R ^ { m }$, ta sẽ tìm được ma trận jacobi của $f$ dựa trên đạo hàm từng phần hàm $f$ với $x$:

$$\nabla _ { x } f = \operatorname { g r a d } f = \frac { d f } { d x } = \left [ \frac { \partial f \left ( x \right ) } { \partial x _ { 1 } } \quad \frac { \partial f \left ( x \right ) } { \partial x _ { 2 } } \quad \ldots \quad \frac { \partial f \left ( x \right ) } { \partial x _ { n } } \right ]$$

Giờ ta sẽ vào món chính: tính gradient cho quá trình cập nhật teacher. Giả sử với một batch unlabeled data $x_u$, teacher sẽ sinh pseudo label $\hat { y } _ { u } \sim T ( x _ { u } ; \theta _ { T } )$, sau đó student sử dụng $(x_u,\hat{y}_u)$ để cập nhật tham số $\theta_S$ của nó. Chúng ta kỳ vọng tham số mới của student sẽ có dạng $\mathbb{E}_{\widehat{y}_{u} \sim T\left(x_{u} ; \theta_{T}\right)}\left[\theta_{S}-\eta_{S} \nabla_{\theta_{S}} \mathbf{C E}\left(\widehat{y}_{u}, S\left(x_{u} ; \theta_{S}\right)\right)\right]$. Ta sẽ cập nhật tham số của teacher trên tập labeled data thông qua cross-entropy của sự thay đổi giữa tham số của student cũ và student mới:

$$\underbrace{\frac{\partial R}{\partial \theta_{T}}}_{1 \times|T|}=\frac{\partial}{\partial \theta_{T}} \operatorname{CE}\left(y_{l}, S\left(x_{l} ; \mathbb{E}_{\widehat{y}_{u} \sim T\left(x_{u} ; \theta_{T}\right)}\left[\theta_{S}-\eta_{S} \nabla_{\theta_{S}} \operatorname{CE}\left(\widehat{y}_{u}, S\left(x_{u} ; \theta_{S}\right)\right)\right]\right)\right)\tag{8}$$

Đặt:

$$\underbrace{\bar{\theta}_{S}^{\prime}}_{|S| \times 1}=\mathbb{E}_{\widehat{y}_{u} \sim T\left(x_{u} ; \theta_{T}\right)}\left[\theta_{S}-\eta_{S} \nabla_{\theta_{S}} \mathbf{C E}\left(\widehat{y}_{u}, S\left(x_{u} ; \theta_{S}\right)\right)\right]\tag{9}$$

phương trình (6) trở thành:

$$\underbrace{\frac{\partial R}{\partial \theta_{T}}}_{1 \times|T|}=\frac{\partial}{\partial \theta_{T}} \operatorname{CE}(y_l, S(x_l;\theta_{S}^{\prime}))\tag{8'}$$

Áp dụng quy tắc đạo hàm của hàm hợp với $\frac{\partial R}{\partial \theta_{T}}$:

$$\begin{aligned}\frac{\partial R}{\partial \theta_{T}}&=\frac{\partial R}{\partial \theta_{S}}\times\frac{\partial \theta_S}{\partial \theta_{T}}\\&=\underbrace{\left.\frac{\partial \operatorname{CE}\left(y_{l}, S\left(x_{l} ; \bar{\theta}_{S}^{\prime}\right)\right)}{\partial \theta_{S}}\right|_{\left.\theta_{S} =\bar{\theta}_{S}^{\prime}\right)}}_{1\times| S \mid} \cdot \underbrace{\frac{\partial \bar{\theta}_{S}^{\prime}}{\partial \theta_{T}}}_{|S| \times|T|} \\&=\qquad\qquad\qquad A \qquad\qquad\times \quad B\end{aligned}\tag{10}$$

Xét phương trình (8), phần A chính là quá trình train student $\theta_S^{\prime}$ với labeled data sau khi đã train $\theta_S$ với pseudo data để có $\theta_S^{\prime}$, phần này hoàn toàn có thể tính thông qua backprop thông thường.

Ta xét tiếp phần B:

$$\begin{aligned}B=\underbrace{\frac{\partial \bar{\theta}_{S}^{\prime}}{\partial \theta_{T}}}_{|S| \times|T|} &=\frac{\partial}{\partial \theta_{T}} \mathbb{E}_{\widehat{y}_{u} \sim T\left(x_{u} ; \theta_{T}\right)}\left[\theta_{S}-\eta_{S} \nabla_{\theta_{S}} \operatorname{CE}\left(\widehat{y}_{u}, S\left(x_{u} ; \theta_{S}\right)\right)\right] \\&=\frac{\partial}{\partial \theta_{T}} \mathbb{E}_{\widehat{y}_{u} \sim T\left(x_{u} ; \theta_{T}\right)}\left[\theta_{S}-\eta_{S} \cdot\left(\left.\frac{\partial \operatorname{CE}\left(\widehat{y}_{u}, S\left(x_{u} ; \theta_{S}\right)\right)}{\partial \theta_{S}}\right|_{\theta_{S}=\theta_{S}}\right)^{\top}\right]\end{aligned}\tag{11}$$

Chú ý : với pt (11), jacobian của $\operatorname{CE}(\widehat{y}_{u}, S(x_{u} ; \theta_{S}))$ có $\text{dim}=1\times |S|$ cần được chuyển vị để khớp với $\text{dim}_{\theta_S}=|S|\times1$.

Vậy thì tại sao $\theta_S$ có $dim=|S|\times1$ và $\nabla_{\theta_S}\operatorname{CE}$ có $dim=1\times|S|$ ? Ở đây $|S|$ chính là số lượng tham số có trong student. Với $\theta_S$ là tham số của student nên dĩ nhiên dim của nó sẽ là $|S|$ và mỗi tham số trong student là duy nhất nên $\text{dim}_{\theta_S}=|S|\times1$. Còn $\nabla_{\theta_S}\operatorname{CE}$ là gradient của hàm loss với biến là $\theta_S$ và chỉ có 1 $\theta_S$ được xét đến, trong $\theta_S$ có $|S|$ lượng tham số nên $\text{dim}_{\nabla_{\theta_S}\operatorname{CE}}=1\times|S|$.

Xét phương trình (11), để đơn giản thì ta đặt $g_{S}$ là ký hiệu gradient của student:

$$\underbrace{g_{S}\left(\widehat{y}_{u}\right)}_{|S| \times|1|}=\left(\left.\frac{\partial \operatorname{CE}\left(\widehat{y}_{u}, S\left(x_{u} ; \theta_{S}\right)\right)}{\partial \theta_{S}}\right|_{\theta_{S}=\theta_{S}}\right)^{\top}\tag{12}$$

Do $\theta_S$ không phụ thuộc vào $\theta_T$, nên $\theta_S$ ở phương trình (11) sẽ bị triệt tiêu khi đạo hàm theo $\theta_T$, do đó phương trình (11) sẽ trở thành:

$$\underbrace{\frac{\partial \bar{\theta}_{S}^{\prime}}{\partial \theta_{T}}}_{|S| \times|T|}=-\eta_{S} \cdot \frac{\partial}{\partial \theta_{T}} \mathbb{E}_{\widehat{y}_{u} \sim T\left(x_{u} ; \theta_{T}\right)}[\underbrace{g_{S}\left(\widehat{y}_{u}\right)}_{|S| \times 1}]\tag{13}$$

Bây giờ chúng ta sẽ đi giải quyết "củ khoai" này nhé, theo như paper thì đúng ra sẽ dùng REINFORCE algorithm, nhưng mà mình có đọc qua paper gốc được viết năm 1992 thì thấy khó nuốt quá nên có thử tự giải theo cách "dễ nhai" hơn. Mọi người xem qua và cho ý kiến về cách giải của mình nhé.

Với phương trình (13) thì ta sẽ đi giải quyết đạo hàm của hàm kỳ vọng $\mathbb{E}_{\widehat{y}_{u}}[g_{S}\left(\widehat{y}_{u}\right)]$. Một cách tổng quát thì kỳ vọng của hàm $f(x)$ với biến ngẫu nhiên rời rạc $x$ sẽ có dạng:

$$\mathbb{E}[f(x)]=\sum_{x}{P(x)f(x)}$$

Áp dụng công thức trên vào phương trình (13) với $\mathbb{E}_{\widehat{y}_{u}}[g_{S}\left(\widehat{y}_{u}\right)]$:

$$\frac{\partial }{\partial \theta_T} \mathbb{E}_{\hat{y}_u \sim T(x_u;\theta_T)}[g_s(\hat{y}_u)] \\ = \frac{\partial}{\partial \theta_T} \sum_{\hat{y}_u} p(\hat{y}_u|x_u;\theta_T)g_s(\hat{y}_u) \\ = \sum_{\hat{y}_u} \frac{\partial}{\partial \theta_T}p(\hat{y}_u|x_u;\theta_T)g_s(\hat{y}_u) \tag{14}$$

Để tính đạo hàm của $p$ , ta sẽ cần công thức ($*$) ở dưới. Tuy nhiên để có công thức ($*$) ta cần thực hiện một số bước. Tính đạo hàm của $y = log(f(x))$:

• Đặt $y = log(f(x))$ và $u=f(x)$

$$y=log(u) \text{ => }\frac{dy}{du}=\frac{\text{d }log(u)}{d\;u}=\frac{1}{u}\text{ và }\frac{du}{dx}=f^{\prime}(x)$$

$$\text{=> } \frac { d y } { d x } = \frac { d y } { d u }\cdot \frac { d u } { d x } = \frac{1}{f(x)}\cdot f^{\prime}(x)$$

Viết lại công thức tổng quát:

$$\frac{\partial{logf(x)}}{\partial{x}}=\frac{1}{f(x)}\cdot \frac{\partial{f(x)}}{\partial{x}}$$

$$\frac{\partial{f(x)}}{\partial{x}}=f(x)\cdot \frac{\partial{logf(x)}}{\partial{x}}\tag{*}$$

Áp dụng (*) vào (14):

$$\frac{\partial }{\partial \theta_T} \mathbb{E}_{\hat{y}_u \sim T(x_u;\theta_T)}[g_s(\hat{y}_u)]\\ = \sum_{\hat{y}_u} \frac{\partial}{\partial \theta_T}p(\hat{y}_u|x_u;\theta_T)g_s(\hat{y}_u)\\=\sum_{\hat{y}_u} p(\hat{y}_u|x_u;\theta_T) \frac{\partial}{\partial \theta_T} log(p(\hat{y}_u|x_u;\theta_T)g_s(\hat{y}_u) \\ =\mathbb{E}_{\hat{y}_u \sim T(x_u;\theta_T)} [g_s(\hat{y}_u) \frac{\partial}{\partial \theta_T} log(p(\hat{y}_u|x_u;\theta_T)] \tag{15}$$

Cuối cùng ta có thể diễn giải $\frac{\partial \bar{\theta}_{S}^{\prime}}{\partial \theta_{T}}$ như sau:

$$\begin{aligned}\underbrace{\frac{\partial \bar{\theta}_{S}^{(t+1)}}{\partial \theta_{T}}}_{|S| \times|T|} &=-\eta_{S} \cdot \frac{\partial}{\partial \theta_{T}} \mathbb{E}_{\widehat{y}_{u} \sim T\left(x_{u} ; \theta_{T}\right)}\left[g_{S}\left(\widehat{y}_{u}\right)\right] \\&=-\eta_{S} \cdot \mathbb{E}_{\widehat{y}_{u} \sim T\left(x_{u} ; \theta_{T}\right)}[\underbrace{g_{S}\left(\widehat{y}_{u}\right)}_{|S| \times 1} \underbrace{\cdot \underbrace{\frac{\partial \log P\left(\widehat{y}_{u} \mid x_{u} ; \theta_{T}\right)}{\partial \theta_{T}}}_{1 \times|T|}]}\\&=\eta_{S} \cdot \mathbb{E}_{\widehat{y}_{u} \sim T\left(x_{u} ; \theta_{T}\right)}[\underbrace{g_{S}\left(\widehat{y}_{u}\right)}_{|S| \times 1} \cdot \underbrace{\frac{\partial \operatorname{CE}\left(\widehat{y}_{u}, T\left(x_{u} ; \theta_{T}\right)\right)}{\partial \theta_{T}}}_{1 \times|T|}]\end{aligned}\tag{17}$$

Đến đây thì ta đã có thể sử dụng đạo hàm của cross-entropy để tính brackrop như thông thường. Thay phương trình (17) vào (10):

$$\begin{aligned}\underbrace{\frac{\partial R}{\partial \theta_{T}}}_{1 \times|T|} &=\underbrace{\left.\frac{\partial \mathbf{C E}\left(y_{l}, S\left(x_{l} ; \bar{\theta}_{S}^{\prime}\right)\right)}{\partial \theta_{S}}\right|_{\theta_{S}=\bar{\theta}_{S}^{\prime}}}_{1 \times|S|} \underbrace{\frac{\partial \bar{\theta}_{S}^{\prime}}{\partial \theta_{T}}}_{|S| \times|T|} \\&=\eta_{S} \cdot \underbrace{\left.\frac{\partial \operatorname{CE}\left(y_{l}, S\left(x_{l} ; \bar{\theta}_{S}^{\prime}\right)\right)}{\partial \theta_{S}}\right|_{\theta_{S}=\bar{\theta}_{S}^{\prime}}}_{1 \times|S|} \cdot \mathbb{E}_{\widehat{y}_{u} \sim T\left(x_{u} ; \theta_{T}\right)}[\underbrace{g_{S}\left(\widehat{y}_{u}\right)}_{|S| \times 1} \cdot \underbrace{\frac{\partial \operatorname{CE}\left(\widehat{y}_{u}, T\left(x_{u} ; \theta_{T}\right)\right)}{\partial \theta_{T}}}_{1 \times|T|}]\end{aligned} \qquad\tag{18}$$

Cuối cùng, ta sẽ sử dụng phép xấp xỉ Monte-Carlo cho mọi biểu thức trong pt(18) với $\hat{y}_u$ đã tính từ trước. Cụ thể hơn thì ta sẽ tính xấp xỉ $\bar{\theta}_{S}^{\prime}$ với $\theta_S$ bằng cách cập nhật tham số student với $(x_u,y_u)$: $\theta_S^{\prime}=\theta_S-\eta_{S} \cdot\nabla_{\theta_S}\operatorname{CE}(\hat{y}_u, S(x_u;\theta_s))$. Đồng thời ước lượng $\mathbb{E}$ cũng với $\hat{y}_u$. Với kết quả ước lượng vừa rồi, ta sẽ tính được gradient của $\nabla_{\theta_T}\mathcal{L}_u(\theta_T, \theta_S)$.

Pt(18) là dạng tổng quát cho 1 batch dữ liệu. Để tường minh hơn ta sẽ lấy 1 mẫu ngẫu nhiên trong batch để tính gradient:

$$\begin{aligned}\nabla_{\theta_{T}} \mathcal{L}_{l} &=\eta_{S} \cdot \underbrace{\frac{\partial \operatorname{CE}\left(y_{l}, S\left(x_{l} ; \theta_{S}^{\prime}\right)\right)}{\partial \theta_{S}}}_{1 \times|S|} \cdot \underbrace{\left(\left.\frac{\partial \mathbf{C E}\left(\widehat{y}_{u}, S\left(x_{u} ; \theta_{S}\right)\right)}{\partial \theta_{S}}\right|_{\theta_{S}=\theta_{S}}\right)^{\top}}_{|S| \times 1} \cdot \underbrace{\frac{\partial \operatorname{CE}\left(\widehat{y}_{u}, T\left(x_{u} ; \theta_{T}\right)\right)}{\partial \theta_{T}}}_{1 \times|T|} \\&=\underbrace{\eta_{S} \cdot\left(\left(\nabla_{\theta_{S}^{\prime}} \operatorname{CE}\left(y_{l}, S\left(x_{l} ; \theta_{S}^{\prime}\right)\right)^{\top} \cdot \nabla_{\theta_{S}} \operatorname{CE}\left(\widehat{y}_{u}, S\left(x_{u} ; \theta_{S}\right)\right)\right)\right.}_{\text {A scalar }:=h} \cdot \nabla_{\theta_{T}} \mathbf{C E}\left(\widehat{y}_{u}, T\left(x_{u} ; \theta_{T}\right)\right)\end{aligned}$$

Đến đây là hết phần diễn giải cách cập nhật của teacher dựa trên gradient của student rồi nhỉ, các bạn thấy scalar $h$ bên trên chứ ? Đấy chính là thứ mà chúng ta mong muốn từ đầu đến giờ : feedback của student để teacher cải thiện performance. Khi các bạn xem phần pseudo code với UDA bên dưới thì sẽ thấy một $h$ tương tự:

Tuy nhiên, khi xem code của Meta Pseudo Label thì các bạn sẽ thấy $h$ được tính như thế này: . Biến dot_product chính là công thức tính $h$ lằng nhằng phía trên đấy :v Nếu viết lại theo công thức toán học dựa trên đoạn code thì $h$ sẽ được tính như sau: $h = L(\theta_S) - L(\theta_S')$. Vậy tại sao từ h dài dòng lại có thể biến đổi thành phép trừ 2 hàm loss đơn giản như vậy? Thử chứng minh 1 chút nhé:

$$\theta'_S = \theta_S - \eta_S \nabla_{\theta_S}CE(\hat{y}_u, S(x_u;\theta_S))$$

$$\text{Đặt }\eta = \eta_S \nabla_{\theta_S}CE(\hat{y}_u, S(x_u;\theta_S)) \text{ ta có: }$$

$$\theta'_S = \theta_S - \eta$$

Áp dụng công thức xấp xỉ taylor: $f(x+h)=f(x) + hf'(x)$

$$L(\theta_S') = L(\theta_S - \eta) \approx L(\theta_S) - \eta \nabla_{\theta_S}L(\theta_S) \\ =L(\theta_S) - \eta_S \nabla_{\theta_S}CE(\hat{y}_u, S(x_u;\theta_S)) \nabla_{\theta_S}L(\theta_S) \\ =L(\theta_S) - \eta_S \nabla_{\theta_S}CE(\hat{y}_u, S(x_u;\theta_S)) \nabla_{\theta_S}CE(y_l, S(x;\theta_S)) \\ = L(\theta_S) - h$$

$$h = L(\theta_S) - L(\theta_S')$$

Done!!

Dưới đây là toàn bộ quá trình train teacher với UDA và feedback từ student:

Và kết quả SoTA của MPL với EfficientNet-L2:

Lời kết

Bài viết của mình đến đây là đã hoàn thành mục đích ban đầu: cố gắng thử thách bản thân với một paper kinh điển do các idol người Việt viết và mang paper này đến với mọi người một cách dễ hiểu nhất. Nếu có thắc mắc thì các bạn có thể comment bên dưới, mình sẽ cố gắng trả lời trong tầm kiến thức của bản thân. Hoặc nếu các bạn phát hiện lỗi sai thì cứ thẳng thắn góp ý nhé. Cảm ơn các bạn đã đọc bài.

References

1. Meta Pseudo Labels
2. Unsupervised Data Augmentation for Consistency Training