Thay thế lan truyền ngược: Sakana AI giới thiệu phương pháp huấn luyện mạng thần kinh mới PC-ALM

Chúng tôi giới thiệu PC-ALM, một giải pháp thay thế cục bộ cho backpropagation. PC-ALM huấn luyện các residual MLP lên tới 1000 lớp, đạt hiệu suất gần tương đương với backprop mặc dù chỉ sử dụng các động lực học cục bộ tại mỗi lớp. PC-ALM trang bị cho mỗi lớp một hệ thống động lực học điều khiển phản hồi giúp phân phối và lan truyền tín hiệu giám sát (supervision credit) xuyên suốt mạng lưới.

Deep learning tiêu chuẩn dựa vào backpropagation. Tuy nhiên, não bộ không thể thực hiện backpropagation, ít nhất là không chính xác như vậy[1, 2]. Làm thế nào não bộ giải quyết được bài toán phân bổ tín hiệu (credit assignment) đa lớp mà không cần sử dụng trực tiếp backprop vẫn là một trong những vấn đề cơ bản chưa có lời giải trong khoa học thần kinh (dù đã có những tiến bộ nhất định[3, 4, 5]).

Có vài lý do khiến não bộ không thể thực hiện backpropagation chính xác. Một trong số đó là "khóa pha" (phase locking)[6, 2]. Backpropagation vận hành theo ba giai đoạn nghiêm ngặt: 1) truyền xuôi (forward pass), sau đó 2) truyền ngược (backward pass), và 3) cập nhật trọng số. Việc cập nhật trọng số bị khóa cho đến khi quá trình truyền xuôi và truyền ngược hoàn tất—một neuron ở lớp đầu phải giữ trạng thái kích hoạt và chờ đợi tín hiệu lỗi truyền đến. Não bộ không có cơ chế nào được biết đến có thể thực thi sự phối hợp thời gian nghiêm ngặt như vậy trên toàn bộ mạng lưới[1].

Trong bài viết này, chúng tôi giới thiệu PC-ALM (Augmented Lagrangian Predictive Coding), một phương pháp huấn luyện mạng lưới thay thế các bước truyền xuôi và truyền ngược của backprop bằng các hệ thống động lực học cục bộ tại mỗi lớp. Mỗi lớp chỉ kết nối với các lớp láng giềng của nó. Thay vì truyền xuôi rồi truyền ngược, chúng tôi vận hành mỗi lớp theo thời gian. Khi đạt đến trạng thái hội tụ, động lực học của toàn bộ hệ thống sẽ phân phối các tín hiệu giám sát một cách nhanh chóng và chính xác trên toàn bộ mạng lưới.

PC-ALM là một phần mở rộng của predictive coding (PC) tiêu chuẩn[7, 8, 9, 10]. PC sử dụng kết nối khuếch tán (tức là dựa trên năng lượng hoặc "dòng nhiệt") giữa các lớp. So với PC, PC-ALM giới thiệu các neuron kép (nhân tử Lagrange) cho mỗi lớp, biến sự tái phát cục bộ của mỗi lớp thành một bộ điều khiển phản hồi PI. Trong trường hợp giới hạn của các mạng lưới tuyến tính, các neuron kép hội tụ về đúng các tín hiệu credit của backprop, mặc dù chỉ sử dụng tính toán cục bộ.

Chúng tôi so sánh PC-ALM với PC và backprop trong một loạt các thí nghiệm. Các phương pháp huấn luyện cục bộ như PC từ trước đến nay rất khó mở rộng quy mô. Theo các tài liệu về PC, chúng tôi sử dụng các tác vụ đơn giản (Fashion-MNIST, CIFAR-10, v.v.) và các mạng lưới như residual MLP.

Chúng tôi chứng minh rằng PC-ALM có thể lan truyền thành công tín hiệu giám sát trong các mạng thần kinh 1000 lớp, khắc phục vấn đề suy giảm tín hiệu của PC tiêu chuẩn

trong khi vẫn duy trì tính cục bộ tại mỗi lớp.

Chúng tôi tập trung vào các mạng lưới sâu, hẹp, một môi trường mà PC thường hoạt động kém hiệu quả.

Cuối cùng, động lực của chúng tôi là tìm hiểu cách các hệ thống phân tán (như não bộ) có thể thực hiện các tính toán gradient mà không cần backpropagation. Ngoài các động lực khoa học, nghiên cứu này có thể đóng góp cho deep learning tiết kiệm năng lượng trên phần cứng neuromorphic, nơi việc mô phỏng hệ thống động lực học có chi phí thấp hơn so với trên GPU[12].

Predictive coding: mỗi lớp là một hệ thống động lực học

Trước khi giải thích về PC-ALM, hãy để chúng tôi giải thích về PC, diễn giải nó từ góc độ hệ thống động lực học để nhấn mạnh vai trò của nó như một giải pháp thay thế backprop.

Predictive coding

Predictive coding bắt nguồn từ các lý thuyết của Helmholtz về nhận thức vô thức[13]. Rao & Ballard (1999) đã phát triển một khung toán học cho PC như một mô hình của vỏ não thị giác[14]. Ý tưởng của PC là mỗi lớp cố gắng mô hình hóa các tín hiệu đầu vào của nó, chỉ gửi lên trên sai số dự đoán (phần mà lớp đó không mô hình hóa được) cho lớp tiếp theo.

Về mặt toán học, predictive coding sử dụng một mô típ chung: lấy một trạng thái và cập nhật nó để giảm sai số dự đoán ở bước tiếp theo,

statet+1=statet−η(statet−targett)⏟sai số dự đoán

Bằng cách áp dụng quy tắc cập nhật này cho vector kích hoạt của mỗi lớp ("trạng thái" là kích hoạt hi của lớp; "mục tiêu" là dự đoán σ(Wihi−1) đến từ lớp bên dưới)1, khung PC thực sự tránh được nhu cầu về một bước truyền xuôi và truyền ngược đồng bộ của backprop.

Để giải thích chi tiết hơn, hãy viết một mạng feedforward dưới dạng một bài toán tối ưu hóa có ràng buộc:

minimizeθ,h12‖y−WLhL−1‖2subject tohi=σ(Wihi−1),i=1,…,L−1.

trong đó L là độ sâu mạng lưới, h0:=x là đầu vào, y là mục tiêu, θ={Wi} là các trọng số, hi là các kích hoạt lớp, và σ là hàm kích hoạt như ReLU. Lưu ý rằng mỗi hi là một biến tối ưu hóa2. Sau đó, chúng tôi xây dựng một hàm mất mát mới bao gồm hàm mất mát giám sát ban đầu, cùng với một hình phạt bậc hai cho các vi phạm ràng buộc của mỗi lớp:

FPC(h,θ)=12‖y−WLhL−1‖2+12∑i=1L−1‖hi−σ(Wihi−1)‖2.

Đây là một sự nới lỏng bậc hai của bài toán có ràng buộc. FPC được gọi là "năng lượng tự do" (free energy) của mạng lưới[15, 9].

Để huấn luyện một mạng thần kinh, PC luân phiên giữa các bước suy luận và học tập:

Predictive coding

suy luận

với t=1,…,T

hi←hi−ηh∇hiFPCvới i=1,…,L−1

học tập

Wi←Wi−ηθ∇WiFPCvới i=1,…,L

Mỗi mini-batch, một bước truyền xuôi khởi tạo các kích hoạt, theo sau là T bước suy luận và một lần cập nhật trọng số duy nhất. Chúng tôi đặt T tỷ lệ thuận với độ sâu mạng lưới; các thí nghiệm 1000 lớp dưới đây sử dụng T=2L.

Mỗi lần cập nhật hi làm giảm sai số dự đoán giữa các lớp liền kề với i. Điều này là do ∇hiFPC chỉ phụ thuộc vào hi−1, hi, và hi+1. Suy luận chỉ yêu cầu giao tiếp giữa các lớp láng giềng ("truyền tin"). Cụ thể, viết ri=hi−σ(Wihi−1) cho sai số dự đoán giữa các lớp i−1 và i, cập nhật suy luận có dạng3:

hi←hi−ηh(ri↑lỗi bên dưới−Wi+1⊤(σ′⊙ri+1↑lỗi bên trên))i=1,…,L−1

Lớp dưới cùng h0 được "kẹp" (cố định) vào một giá trị đầu vào và đỉnh của mạng lưới được kẹp vào mục tiêu y. Chạy T bước cập nhật, mạng lưới ổn định ở các trạng thái hi cho mỗi lớp, sau đó một bước hạ gradient được thực hiện trên cùng FPC đó nhưng bây giờ là đối với các trọng số W (dựa trên các sai số dự đoán còn lại hiện tại và các kích hoạt trạng thái hiện tại).

Wi←Wi+ηθ(σ′⊙ri)hi−1⊤i=1,…,L−1

Cả bước suy luận và cập nhật trọng số đều mang tính cục bộ tại mỗi lớp. Việc cập nhật trọng số giống như Hebbian, ở chỗ nó nhân một sai số hậu synap với hoạt động tiền synap (quy tắc delta), và các động lực học ánh xạ lên một mạch thần kinh với các neuron lỗi rõ ràng[14, 7].

PC huấn luyện các mạng sâu, nhưng thể hiện sự suy giảm tín hiệu

Vì việc giảm thiểu năng lượng tự do đối với mỗi hi không thực thi các ràng buộc theo lớp một cách chính xác, PC dẫn đến một quỹ đạo học tập khác so với backpropagation tiêu chuẩn.

Tuy nhiên, PC đã được chứng minh là huấn luyện thành công các mạng lưới trên các tác vụ đơn giản. Ví dụ, MNIST và Fashion-MNIST trong các residual MLP 128 lớp với độ rộng lớp lớn (512 neuron mỗi lớp)[16].

Tuy nhiên, PC gặp khó khăn với các tác vụ và mạng lưới phức tạp hơn[17]. Hơn nữa, PC thậm chí gặp khó khăn trên các mạng/tác vụ đơn giản nếu độ rộng mạng lưới nhỏ hơn độ sâu của nó[18].

Mỗi lớp điều chỉnh hoạt động của nó để giảm sai số dự đoán với các láng giềng. Sự giám sát đi vào từ đầu ra, nhưng phải đi qua chuỗi các thỏa hiệp cục bộ này để ảnh hưởng đến các lớp trước đó. Trong các mạng sâu và hẹp, tín hiệu credit thu được trở nên yếu đi rất lâu trước khi nó chạm đến đầu vào.