Sakana AI giới thiệu PC-ALM: Giải pháp thay thế lan truyền ngược, huấn luyện mạng 1000 lớp

Backpropagation là một thuật toán toàn cục: một lượt truyền xuôi, sau đó là một lượt truyền ngược, rồi đến bước cập nhật trọng số, mỗi bước đều bị khóa bởi bước trước đó. Bộ não con người không có cơ chế nào được biết đến cho kiểu khóa pha trên toàn mạng lưới như vậy, đó là lý do tại sao các phương pháp thay thế học tập cục bộ như predictive coding (PC) vẫn thu hút sự quan tâm của giới nghiên cứu. Các nhà nghiên cứu tại Sakana AI đề xuất Augmented Lagrangian Predictive Coding (PC-ALM), một biến thể của PC giúp giữ mọi cập nhật ở phạm vi cục bộ theo lớp nhưng vẫn khôi phục được các tín hiệu phân bổ trọng số (credit signals) tương đương với backprop. Nhóm nghiên cứu báo cáo rằng việc huấn luyện các mạng residual MLP lên tới 1000 lớp đạt kết quả chỉ chênh lệch khoảng 2 điểm phần trăm so với backprop trên tập dữ liệu MNIST.

Liệu nó có thể triển khai được không? Có, dưới dạng mã nguồn nghiên cứu: một bản triển khai tham chiếu JAX theo giấy phép MIT có thể chạy trên CPU và tái lập lưới chiều rộng-chiều sâu trong bài báo. Đây là một phương pháp huấn luyện, không phải là một mô hình, và mới chỉ được thử nghiệm trên các tập dữ liệu benchmark hình ảnh nhỏ.

Tại sao PC tiêu chuẩn bị đình trệ trong các mạng sâu và hẹp

PC coi mỗi kích hoạt ẩn là một biến tối ưu hóa và phạt sự sai lệch bình phương giữa kích hoạt của mỗi lớp và dự đoán đến từ lớp bên dưới. Suy luận là quá trình hạ gradient trên năng lượng đó; học tập là một bước cập nhật trọng số kiểu Hebbian. Vấn đề là sự giám sát (supervision) đi vào từ đầu ra và phải khuếch tán qua một chuỗi các thỏa hiệp cục bộ. Trong các mạng sâu và hẹp, tín hiệu phân bổ trọng số mờ dần từ rất lâu trước khi nó chạm tới đầu vào. Innocenti và cộng sự đã mô tả khoảng cách PC-BP này như một hàm số của chiều rộng và chiều sâu, và nó tồi tệ nhất khi chiều rộng nhỏ hơn chiều sâu.

Những thay đổi của PC-ALM

PC-ALM bắt đầu từ quan điểm huấn luyện có ràng buộc: tối thiểu hóa hàm mất mát có giám sát với điều kiện $h_i = \sigma(W_i h_{i-1})$ tại mỗi lớp. PC là phương pháp nới lỏng ràng buộc đó bằng hình phạt bậc hai. Thay vào đó, PC-ALM sử dụng Augmented Lagrangian, gắn một nhân tử Lagrange $\lambda_i \in \mathbb{R}^{d_i}$ sao cho $\text{dim}(\lambda_i) = \text{dim}(h_i)$ vào mỗi ràng buộc lớp trong khi vẫn giữ hình phạt của PC. Việc đặt $\lambda = 0$ sẽ khôi phục chính xác phương pháp PC.

Suy luận luân phiên giữa 2 bước cục bộ: một bước gradient sơ cấp trên các kích hoạt, và một bước đối ngẫu $\lambda_i \leftarrow \lambda_i + \alpha r_i$ tích lũy sai số dự đoán của lớp. Việc hoàn thiện bình phương cho thấy mỗi bước sơ cấp là một bước PC tiêu chuẩn với mục tiêu dự đoán bị dịch chuyển bởi $-\lambda_i/\rho$. Sau T bước, cập nhật trọng số tác động lên tín hiệu tổng hợp $\lambda_i + \rho r_i$. Nhóm nghiên cứu coi đây là một bộ điều khiển PI cho mỗi lớp: sai số dự đoán là thành phần tỉ lệ và nhân tử là thành phần tích phân. $\alpha = 0$ cho ra PC; $\alpha = \rho$ với bài toán bên trong được giải chính xác cho ra phương pháp nhân tử cổ điển.

Gradient backprop chính xác trong trường hợp tuyến tính

LeCun đã quan sát vào năm 1988 rằng các nhân tử Lagrange của một mạng có ràng buộc bằng với các adjoint của backprop tại một điểm KKT. Nhóm nghiên cứu chứng minh rằng trong các mạng PC tuyến tính, dưới điều kiện ổn định bán kính phổ, PC-ALM hội tụ về điểm KKT đó: các kích hoạt quay trở lại giá trị truyền xuôi của chúng trong khi mỗi $\lambda_i$ tích phân thành adjoint BP chính xác. Giới hạn ổn định theo chế độ là $\eta_h \sigma_i^2 (2\rho + \alpha) < 4$, điều này rút gọn về điều kiện của PC khi $\alpha = 0$. Không giống như dòng gradient đơn điệu của PC, ma trận lặp của PC-ALM có các giá trị riêng phức tạo ra các dao động tắt dần; $\alpha$ thiết lập tần số của chúng nhưng không phải tốc độ suy giảm.

Kết quả

Nhóm nghiên cứu quét các mạng residual MLP với chiều rộng và chiều sâu từ 8 đến 128 trên Fashion-MNIST và MNIST theo tham số hóa trường trung bình của Innocenti và cộng sự, huấn luyện trong 1 epoch. Với ngân sách suy luận T = 2L, PC-ALM tương đương với backprop trên mọi chiều rộng, chiều sâu và hàm kích hoạt (identity, tanh, ReLU), trong khi PC giảm mạnh ở các ô sâu và hẹp. Ô tham chiếu của kho lưu trữ (chiều rộng 32, chiều sâu 32, ReLU, Fashion-MNIST) báo cáo độ chính xác kiểm thử 78,66% cho BP, 68,13% cho PC và 77,75% cho PC-ALM, với cosine gradient so với BP tăng từ 0,604 lên 0,909.

Nghiên cứu mở rộng bức tranh: các mạng residual MLP 1000 lớp trên MNIST (chiều rộng 32, ReLU, 5 epochs) duy trì trong khoảng 2 điểm so với BP, và PC-ALM cải thiện hơn PC trên mọi benchmark đã thử nghiệm, bao gồm ResNet-18 trên CIFAR-10 và Tiny ImageNet.

Những điểm chính cần lưu ý

Hãy xem Bài báo, Blog và GitHub Repo. Mọi công lao thuộc về các nhà nghiên cứu của dự án này. Ngoài ra, hãy thoải mái theo dõi chúng tôi trên Twitter và đừng quên tham gia SubReddit ML 150k+ của chúng tôi và Đăng ký Bản tin của chúng tôi. Khoan đã! Bạn có dùng Telegram không? Bây giờ bạn cũng có thể tham gia cùng chúng tôi trên Telegram.

Cần hợp tác với chúng tôi để quảng bá GitHub Repo HOẶC trang Hugging Face HOẶC Ra mắt sản phẩm HOẶC Hội thảo trực tuyến, v.v.? Hãy kết nối với chúng tôi.