+2

Paper reading | ZipIt! Merging Models from Different Tasks without Training

Giới thiệu chung

Các model hình ảnh ngày càng phát triển mạnh với khả năng thực hiện chính xác được nhiều task khác nhau, từ classification với hàng nghìn class cho đến những task về object detection, instance segmentation, image generation 😄 Tuy phát triển nhanh và mạnh như vậy nhưng các model này cũng có một số vấn đề tiềm ẩn:

  • Các model chỉ thực hiện được trên các task được train, nếu ta mở rộng việc train trên các task mới sẽ xảy ra vấn đề "lãng quên" (trong bài báo gọi là catastrophic forgetting). Catastrophic forgetting đề cập đến hiện tượng một model được đào tạo về một task cụ thể sẽ quên "kiến thức đã học" trước đó khi nó được đào tạo về một task mới. Đây là một vấn đề trong deep learning, đặc biệt là trong các tình huống continual learning, trong đó model cần học các nhiệm vụ mới theo thời gian mà không cần truy cập vào tất cả dữ liệu trước đó 😄. Sự lãng quên xảy ra nghiêm trọng khi các tham số của model được cập nhật dựa trên dữ liệu mới, khiến model ghi đè lên kiến thức đã học trước đó. Về cơ bản, dữ liệu đào tạo ban đầu của model trở nên ít quan trọng hơn theo thời gian khi model được tinh chỉnh trên dữ liệu mới. Do đó, model có thể thực hiện kém các nhiệm vụ đã học trước đó khi được thử nghiệm lại.
  • Nếu đánh giá model trên những data mới chưa được cập nhật, ta sẽ thường thấy là model có khả năng tổng quát hóa khá kém với những domain mà chưa được train bao giờ.

Các cách giải quyết vấn đề này thường là các kĩ thuật regularization hoặc phải training lại, điều này làm mất thời gian và tài nguyên để training. Nếu như bây giờ có một phương pháp có thể kết hợp các model được train trên trên các task khác nhau thành một model mà không phải train lại thì quá tuyệt 😄 Trong bài báo, nhóm tác giả đề xuất phương pháp kết hợp các model được khởi tạo khác nhau và được train trên các task hoàn toàn riêng biệt, bạn đọc xem hình dưới đươc trích trong bài báo để hiểu hơn về ý tưởng chung 😃.

image.png

Động lực

Mục tiêu của phương pháp là kết hợp các model khác nhau mà không cần phải training thêm lần nào. Cụ thể bài toán được mô tả như sau. Xét một model L\mathcal{L} là một tập hợp của các layer LiLL_i \in \mathcal{L}, mỗi layer này có chứa các tham số (ví dụ như với linear layer thì ta có tham số Wi,biW_i, b_i). Ta định nghĩa kết hợp 2 model LA\mathcal{L}^ALB\mathcal{L}^B là kết hợp các tham số của 2 model này thành một model mới L\mathcal{L}^* sao cho L\mathcal{L}^* vẫn duy trì được độ chính xác của LA\mathcal{L}^ALB\mathcal{L}^B trên các task gốc.

Nếu LA\mathcal{L}^ALB\mathcal{L}^B được finetuned từ cùng checkpoint, một số nghiên cứu nhận thấy rằng việc kết hợp 2 model khá dễ dàng bằng cách thực hiện tính trung bình các trọng số của chúng. Ví dụ, nếu LiL_i là một linear layer, ma trận trọng số mới WiW^*_i được tính như sau:

Wi=12WiA+12WiBW_i^*=\frac{1}{2} W_i^A+\frac{1}{2} W_i^B

Tuy nhiên, nếu LA\mathcal{L}^ALB\mathcal{L}^B không được finetuned từ cùng 1 checkpoint thì cách tính ma trận trọng số trên cho ta độ chính xác ngẫu nhiên 😄 Để giải quyết vấn đề này, một số nghiên cứu trước đây đề xuất phương pháp căn chỉnh (ví dụ như hoán vị) feature space của một model với feature space của một model khác trước khi tính giá trị trung bình. Nếu một trong các feature space của model được căn chỉnh để khớp với feature space của model kia thì hai model có thể trở nên bổ sung cho nhau nhiều hơn, điều này có thể dẫn đến sự kết hợp chính xác hơn các dự đoán của chúng. Cụ thể hơn, gọi PiP_i là một ma trận hoán vị thực hiện hoán vị output của layer LiBL^B_i sang space của LiAL^A_i.

Wi=12WiA+12PiWiBPi1TW_i^*=\frac{1}{2} W_i^A+\frac{1}{2} P_i W_i^B P_{i-1}^T

Nhóm tác giả thực hiện hoán vị output space của WiBW^B_i nhưng cũng cần hoán vị input space của WiBW^B_i. Lý do là input space của WiBW^B_i là output space của layer trước đó và feature dimension của input space đã được hoán vị bởi hoán vị của layer trước đó. Vì vậy, để căn chỉnh input và output space của WiBW^B_i với space tương ứng của model khác, ta cần "undo" hoán vị input space của WiBW^B_i bằng cách sử dụng Pi1TP^T_{i-1}. Tuy nhiên, cách làm này chỉ chính xác với các model được train trên task giống nhau.

image.png

Phương pháp

Phương pháp được sử dụng là kết hợp các trọng số của 2 model thành một bộ trọng số duy nhất. Ý tưởng là ta sẽ thực hiện kết hợp mỗi layer của model này với layer tương ứng của model kia và sẽ điều chỉnh cả hai (khác với cách làm đã được trình bày ở trên, ta chỉ thực hiện hoán vị một trong hai mô hình).

Ví dụ, nếu layer LiLL_i \in \mathcal{L} là một linear layer, ta sẽ có bộ trọng số WiRni×mi,biRniW_i \in \mathbb{R}^{n_i \times m_i}, b_i \in \mathbb{R}^{n_i} và input là $x \in Rmi\mathbb{R}^{m_i}$ với output feature vector fiRnif_i \in \mathbb{R}^{n_i} được tính như sau

fi=Li(x)=Wix+bif_i=L_i(x)=W_i x+b_i

Mục tiêu là lấy LiALAL_i^A \in \mathcal{L}^A từ model A và LiBLBˉL_i^B \in \mathcal{L}^{\bar{B}} từ model B và kết hợp thành một layer LiL_i^*, feature space fif_i^* chứa thông tin của fiAf_i^AfiBf_i^B.

Để xây dựng feature fif_i^* ta giả sử rằng tồn tại các feature "thừa" (redundant feature) trong fiAf_i^AfiBf_i^B. Điều này có nghĩa là tồn tại một số phần tử của 2 feature vector có độ tương quan cao trên một mẫu dữ liệu. Để có thể xét sự tương quan giữa các gặp phần tử trong 2 feature vector, ta thực hiện concat 2 feature vector thành fiAfiBR2nif_i^A \| f_i^B \in \mathbb{R}^{2 n_i} (với nin_i là số cặp các phần tử) và xét sự tương quan giữa các cặp phần tử trong vector được concat mới.

Nếu 2 feature có độ tương quan cao thì ta chỉ cần tính trung bình của chúng. Ta định nghĩa ma trận kết hợp MiRni×2niM_i \in \mathbb{R}^{n_i \times 2 n_i}fif_i^* được tính như sau

fi=Mi(fiAfiB)f_i^*=M_i\left(f_i^A \| f_i^B\right)

Khi đó ma trận MiM_i là một ma trận thưa và tồn tại cặp có chỉ số pp của bộ (j,k)(j,k) với Mi[p,j]=Mi[p,k]=1/2M_{i[p, j]}=M_{i[p, k]}=1 / 2. Ta có thể tìm bộ này sử dụng thuật toán tham lam, vẫn có thuật toán tối ưu độ chính xác hơn nhưng chạy rất chậm và độ chính xác chỉ hơn một chút.

image.png

Một vấn đề gặp phải là với layer tiếp theo Li+1A,Li+1BL_{i+1}^A, L_{i+1}^B không tương thích với biểu diễn tính ở công thức trên. Vì vậy, ta cần "undo" thao tác merge trước khi truyền các feature vào layer tiếp theo. Ta định nghĩa ma trận "unmerge" UiR2ni×niU_i \in \mathbb{R}^{2 n_i \times n_i} thỏa mãn

UififiAfiBU_i f_i^* \approx f_i^A \| f_i^B

Ta cũng có thể tách ma trận unmerge thành UiA,UiBRni×niU^A_i, U^B_i \in \mathbb{R}^{n_i \times n_i} bằng cách chia nửa số hàng của ma trận. Khi đó ta có thể tính toán feature mới như sau:

fi+1A=Li+1A(UiAfi)fi+1B=Li+1B(UiBfi)f_{i+1}^A=L_{i+1}^A\left(U_i^A f_i^*\right) \quad f_{i+1}^B=L_{i+1}^B\left(U_i^B f_i^*\right)

image.png

Từ các thành phần trên ta rút ra một thao tác chung để merge LiAL^A_iLiBL^B_i bất kì như sau:

  • Tính MiM_iUiU_i bằng cách matching feature fiAf^A_ifiBf^B_i. Sau đó truyền UiU_i tới layer tiếp theo và nhận Ui1U_{i-1} từ layer trước đó.
  • Từ MiM_iUiU_i ta tính được parameter mới với linear layer như sau:

fi+1A=Li+1A(UiAfi)fi+1B=Li+1B(UiBfi)f_{i+1}^A=L_{i+1}^A\left(U_i^A f_i^*\right) \quad f_{i+1}^B=L_{i+1}^B\left(U_i^B f_i^*\right)

Trong đó MiAM^A_iMiBM^B_iMiM_i được chia thành 2 nửa theo cột. bib^*_i được tính tương tự như công thức trên nhưng không dùng unmerge.

Tuy nhiên, các model hiện đại không chỉ có một loại layer là linear layer. Một số layer như BatchNorm, ReLU không có ma trận trọng số để tính toán bằng công thức trên. Do đó, ta truyền MiM_iUiU_i qua các layer này như mô tả trong hình dưới.

image.png

Như hình trên ta có một block phổ biến trong mạng CNN, do ReLU và BatchNorm không có ma trận trọng số nên ta bỏ qua, sau đó ta chỉ cần tính ma trận trọng số mới dùng MiM_iUiU_i khi tới Conv layer. Một cách tổng quát hơn, ta cần quan tâm đến cơ chế propagation của từng model để tính toán ma trận trọng số mới sao cho phù hợp.

Trong một số trường hợp, ta có thể không muốn "zip" tất cả layer của 2 mạng neural, đặc biệt nếu như không gian vector của các layer không tương thích hoặc làm giảm độ chính xác đi quá nhiều. Thay vào đó, ta sẽ thực hiện "zip một phần" như trong hình dưới.

image.png

Cách thực hiện rất đơn giản, ta thực hiện zip bình thường đến layer ii nào đó, sau đó các layer mà không được zip sẽ nhận UiU_i thông qua zip propagation. Khi đó, kết quả sẽ có dạng multihead model và thực hiện multple tasks.

Trong trường hợp khác, ta lại muốn merge nhiều hơn 2 model. Khi đó, ta sẽ thực hiện việc lặp matching. Tại đây, khi sử dụng thuật toán tham lam tìm được 2 feature khớp với nhau, ta sẽ xóa chúng và thay bằng merged feature mới. Để ngăn việc một feature được merge quá nhiều lần, nhóm tác giả đặt giá trị tương quan feature mới là giá trị tối thiểu của giá trị tương đồng các feature cũ và được đánh trọng số bởi α(0,1]\alpha \in (0,1].

Nhóm tác giả sử dụng một tham số gọi là "budget" β[0,1]\beta \in [0,1] biểu thị tỷ lệ phần trăm feature được merge trong chính model của chúng trên tổng các feature được merge. Bằng cách chỉ định budget, ta có thể tìm hiểu tác động của việc tự merge feature trong các model riêng lẻ.

image.png

Nhóm tác giả nhận thấy rằng trong tập dữ liệu CIFAR budget tối ưu nhất có giá trị trên 0.8, tức là 80% việc merge là đến từ chính bên trong model riêng lẻ. Đối với ImageNet, điều này lại không xảy ra do độ khó của task làm cho tồn tại ít các feature thừa trong mỗi model.

Thực nghiệm

Nhóm tác giải thực hiện 2 loại thực nghiệm để đánh giá kết quả của việc merge 2 model. Thực nghiệm 1 đó là merge các model được train trên cùng dataset nhưng khác label. Thực nghiệm 2 là merge các model được train khác dataset và tất nhiên là khác cả label.

Bảng dưới là kết quả trong bộ data CIFAR. So sánh ZipIt! n/m với các baseline, trong đó n/m có nghĩa là n trên m layer được zip. Model được zip từ 2 model, một model được train trên một nửa class (Task A) và một model được train trên các class khác. Model được sử dụng là ResNet-20.

image.png

Bảng dưới là kết quả trên tập ImageNet-1k sử dụng ResNet-50 được train từ đầu trên 200 class con (mỗi model) trong 1000 class của bộ data.

image.png

Với model được train trên nhiều bộ dữ liệu khác nhau Stanford Dogs (SD), Oxford Pets (OP), CUB200 (CUB), và NABirds (NAB). Bảng dưới report kết quả độ chính xác trung bình trên mỗi task theo 2 cách, merge theo cặp model và merge cả 4 model.

image.png

Kết luận

Bài báo cho ta một ý tưởng mới để tận dụng tài nguyên bằng cách kết hợp các model được train trên các task khác nhau mà không phải training lại từ đầu. Đây là một ý tưởng khá độc đáo tuy nhiên lại rất "kén" áp dụng cho từng bài toán cụ thể. Hy vọng sau bài báo này các bạn có thêm một hướng đi mới cho việc xây dựng model của mình.

Tham khảo

[1] ZipIt! Merging Models from Different Tasks without Training


All rights reserved

Viblo
Hãy đăng ký một tài khoản Viblo để nhận được nhiều bài viết thú vị hơn.
Đăng kí