นักวิจัย Sakana AI เปิดตัว PC-ALM ทางเลือกการเรียนรู้แบบ Layer-Local แทน Backpropagation รองรับเครือข่ายลึกถึง 1,000 ชั้น

Backpropagation เป็นอัลกอริทึมแบบโกลบอล (global algorithm) ที่เริ่มจาก forward pass ตามด้วย backward pass แล้วจึงอัปเดตน้ำหนัก (weight update) ซึ่งแต่ละขั้นตอนจะถูกล็อคไว้ตามลำดับก่อนหลัง

สมองของมนุษย์ไม่มีกลไกที่ทราบแน่ชัดสำหรับการทำ phase locking ทั่วทั้งเครือข่ายในลักษณะนั้น นั่นคือสาเหตุที่ทางเลือกการเรียนรู้แบบโลคัล เช่น predictive coding (PC) ยังคงดึงดูดความสนใจในงานวิจัย

นักวิจัยจาก Sakana AI เสนอ Augmented Lagrangian Predictive Coding (PC-ALM) ซึ่งเป็นตัวแปรหนึ่งของ PC ที่ยังคงรักษาการอัปเดตแบบ layer-local ไว้ได้ทุกขั้นตอน แต่ยังสามารถกู้คืนสัญญาณ credit signals ที่สอดคล้องกับ backprop ได้ โดยทีมวิจัยรายงานว่าสามารถฝึก residual MLPs ได้สูงถึง 1,000 ชั้น โดยมีผลลัพธ์ห่างจาก backprop ไม่เกิน 2 เปอร์เซ็นต์บน MNIST

มันสามารถนำไปใช้งานได้จริงหรือไม่? ใช่ ในรูปแบบของโค้ดวิจัย โดยมี JAX reference implementation ภายใต้สัญญาอนุญาต MIT ที่สามารถรันบน CPU และทำซ้ำผลลัพธ์ width-depth grid ตามในงานวิจัยได้ แต่มันเป็นวิธีการฝึก (training method) ไม่ใช่ตัวโมเดล และเพิ่งผ่านการทดสอบกับชุดข้อมูลรูปภาพขนาดเล็กเท่านั้น

ทำไม standard PC ถึงหยุดชะงักในเครือข่ายที่ลึกและแคบ

PC ปฏิบัติกับ activation ของเลเยอร์ที่ซ่อนอยู่ทุกตัวเสมือนเป็นตัวแปรการเพิ่มประสิทธิภาพ (optimization variable) และลงโทษค่า squared mismatch ระหว่าง activation ของแต่ละเลเยอร์กับค่าพยากรณ์ที่ส่งมาจากเลเยอร์ด้านล่าง

การอนุมาน (Inference) คือการทำ gradient descent บนพลังงานนั้น ส่วนการเรียนรู้คือขั้นตอนการปรับน้ำหนักแบบ Hebbian-like

ปัญหาคือการควบคุม (supervision) จะเข้าสู่ระบบที่เอาต์พุตและต้องแพร่กระจายผ่านสายโซ่ของการประนีประนอมระดับโลคัล ในเครือข่ายที่ลึกและแคบ สัญญาณ credit signal จะจางหายไปนานก่อนที่จะถึงอินพุต

Innocenti et al. ได้ระบุลักษณะของช่องว่าง PC-BP นี้ว่าเป็นฟังก์ชันของความกว้างและความลึก และจะแย่ที่สุดเมื่อความกว้างน้อยกว่าความลึก

สิ่งที่ PC-ALM เปลี่ยนแปลง

PC-ALM เริ่มต้นจากมุมมองของข้อจำกัดในการฝึก คือการลด supervised loss ให้เหลือน้อยที่สุดภายใต้ข้อจำกัดในทุกเลเยอร์

ในขณะที่ PC คือการผ่อนคลาย (relaxation) ของปัญหานั้นแบบ quadratic-penalty แต่ PC-ALM ใช้ augmented Lagrangian แทน โดยการเพิ่ม Lagrange multiplier เข้าไปในข้อจำกัดของแต่ละเลเยอร์ในขณะที่ยังคง penalty ของ PC ไว้ ซึ่งการตั้งค่า λ = 0 จะได้ผลลัพธ์เท่ากับ PC ทุกประการ

การอนุมานจะสลับกันระหว่าง 2 ขั้นตอนโลคัล คือขั้นตอน primal gradient บน activations และขั้นตอน dual ที่สะสมค่าความผิดพลาดในการพยากรณ์ของเลเยอร์

การทำ completing the square แสดงให้เห็นว่าแต่ละขั้นตอน primal คือขั้นตอน standard PC ปกติ โดยที่เป้าหมายการพยากรณ์ถูกเลื่อนไปตามค่า หลังจากผ่านไป T ขั้นตอน การอัปเดตน้ำหนักจะทำงานบนสัญญาณรวม

ทีมวิจัยอ่านค่านี้เป็นเสมือน PI controller ต่อเลเยอร์ โดยมีความผิดพลาดในการพยากรณ์เป็นเทอมสัดส่วน (proportional term) และตัวคูณคือเทอมปริพันธ์ (integral term) การตั้งค่า α = 0 จะได้ PC ส่วน α = ρ พร้อมกับการแก้ปัญหาภายในอย่างแม่นยำจะนำไปสู่วิธีดั้งเดิมของ multipliers

ค่า Exact backprop gradients ในกรณี Linear

LeCun สังเกตในปี 1988 ว่า Lagrange multipliers ของเครือข่ายที่มีข้อจำกัดจะมีค่าเท่ากับ backprop adjoints ที่จุด KKT

ทีมวิจัยพิสูจน์ว่าในเครือข่าย linear PC ภายใต้เงื่อนไขความเสถียรของรัศมีสเปกตรัม (spectral-radius stability condition) PC-ALM จะลู่เข้าสู่จุด KKT นั้น โดยที่ activations จะกลับไปสู่ค่า forward-pass ในขณะที่แต่ละค่าจะรวมเข้ากับ BP adjoint ที่แม่นยำ

ขอบเขตความเสถียรต่อโหมดคือ ซึ่งจะลดรูปเหลือเงื่อนไขของ PC ที่ α = 0 ต่างจาก monotone gradient flow ของ PC เมทริกซ์การทำซ้ำของ PC-ALM มีค่า eigenvalues เชิงซ้อนที่สร้างการแกว่งแบบลดทอน (damped oscillations) โดยค่า α จะกำหนดความถี่ของมันแต่ไม่ใช่กำหนดอัตราการสลายตัว

ผลลัพธ์

ทีมวิจัยทดสอบ residual MLPs ที่มีความกว้างและความลึกตั้งแต่ 8 ถึง 128 บน Fashion-MNIST และ MNIST ภายใต้ mean-field parameterization ของ Innocenti et al. โดยฝึกเป็นเวลา 1 epoch ด้วยงบประมาณการอนุมานที่ T = 2L

พบว่า PC-ALM สามารถเทียบชั้น backprop ได้ในทุกความกว้าง ความลึก และฟังก์ชันกระตุ้น (identity, tanh, ReLU) ในขณะที่ PC ประสิทธิภาพตกลงอย่างมากในเซลล์ที่ลึกและแคบ

เซลล์อ้างอิงของ repo (กว้าง 32, ลึก 32, ReLU, Fashion-MNIST) รายงานความแม่นยำในการทดสอบที่ 78.66% สำหรับ BP, 68.13% สำหรับ PC และ 77.75% สำหรับ PC-ALM โดยค่า gradient cosine เมื่อเทียบกับ BP เพิ่มขึ้นจาก 0.604 เป็น 0.909

งานวิจัยยังขยายขอบเขตภาพให้ชัดเจนขึ้น โดย residual MLPs ขนาด 1,000 ชั้นบน MNIST (กว้าง 32, ReLU, 5 epochs) ยังคงมีผลลัพธ์ห่างจาก BP เพียงประมาณ 2 จุด และ PC-ALM มีการปรับปรุงที่ดีขึ้นกว่า PC ในทุกการทดสอบมาตรฐานที่ลอง รวมถึง ResNet-18 บน CIFAR-10 และ Tiny ImageNet

ประเด็นสำคัญ

• PC-ALM เพิ่ม Lagrange multiplier ต่อเลเยอร์ให้กับ predictive coding โดยการอัปเดตทุกอย่างยังคงเป็น layer-local
• ใน linear networks ตัวคูณจะลู่เข้าหา backprop gradients ที่แม่นยำ
• ให้ผลลัพธ์เทียบเท่า BP ในตารางความกว้าง-ความลึกตั้งแต่ 8 ถึง 128 ที่ T = 2L ในขณะที่ PC ล้มเหลวในเซลล์ที่ลึกและแคบ
• สามารถฝึก residual MLPs 1,000 ชั้นได้โดยมีผลลัพธ์ห่างจาก BP ประมาณ 2 จุดบน MNIST
• โค้ด JAX ภายใต้สัญญาอนุญาต MIT สามารถทำซ้ำผลลัพธ์ได้บน CPU