ลดขนาด Neural Network ยังไงให้เหมือนเล่นหวย: Lottery Ticket Hypothesis

Neural network นั้นเป็นคำศัพท์ที่เหล่า deep learning practitioner และ data enthusiasts นั้นคุ้นชินเป็นอย่างดี แม้จะเป็น model ที่ใช้กันอย่างแพร่หลายในงาน computer vision, image processing, natural language processing แต่ model เหล่านั้น ยิ่งพัฒนามากขึ้น จำนวน parameter ก็จะเพิ่มขึ้นตามอย่างมหาศาล (ตัวอย่างเช่น OpenAI GPT-3 ที่มี parameter กว่า 175 พันล้าน parameters!) ทำให้หลายคนเริ่มสงสัยว่า parameters เหล่านี้มีจำนวนมากเกินไปรึเปล่า (overparameterized)?

จำนวน parameters ของ model ในแต่ละปี

ต่อมาได้มีงานวิจัยที่ได้ตีพิมพ์ใน conference ICLR 2019 โดยคุณ Jonathan Frankle และ Michael Carbin หัวข้อ THE LOTTERY TICKET HYPOTHESIS: FINDING SPARSE, TRAINABLE NEURAL NETWORKS. โดยงานวิจัยนี้ได้กล่าวถึงสิ่งที่ค้นพบไว้ดังนี้

The Lottery Ticket Hypothesis. A randomly-initialized, dense neural network contains a subnetwork that is initialized such that — when trained in isolation — it can match the test accuracy of the original network after training for at most the same number of iterations.

กล่าวคือ neural network ที่เพิ่งถูก random weight มาจะมี neural network ขนาดเล็กกว่า (subnetwork) ที่อยู่ภายใน neural network นั้น และหากแยก subnetwork ออกมา train จะทำให้ได้ test accuracy เท่ากับ model ขนาดเต็มของมันโดยใช้เวลาในการ train น้อยกว่าหรือเท่ากับเวลา train model เดิม

ตัวอย่างเช่นในรูปนี้ จะเห็นได้ว่าเราสามารถทำการตัด weight ที่ไม่จำเป็น (prune) ออกไปได้เกิดเป็น subnetwork อีกตัว ซึ่งหากเรา train subnetwork ตัวนี้แล้วได้ test accuracy ใกล้กับก่อนทำ pruning เราจะถือว่า subnetwork ตัวนี้เป็น winning ticket หรือว่าถูกหวยนั่นเอง

Pruning Neural Network

แต่คำถามคือเราจะหา subnetwork ตัวนี้เจอได้อย่างไร? Frankle & Carbin ได้เสนอ algorithm ในการหา winning ticket ไว้ซึ่งมีชื่อว่า Iterative Magnitude Pruning นั่นเอง

Iterative Magnitude Pruning (IMP)

IMP เป็นหนึ่งวิธีในการ pruning ที่ประสบความสำเร็จในการหา winning ticket ใน neural network โดย algorithm จะเป็นไปตามนี้

1. สร้าง neural network ที่มี parameter เริ่มต้น เป็น W_0
2. Train neural network จนได้ W_1
3. สร้าง mask (m) ซึ่งเป็น tensor ขนาดเท่า W โดยจะเป็นตัวกำหนดว่า ค่า weight ตัวไหนใน W ควรถูก set เป็น 0
4. สร้าง neural network ใหม่ (ด้วย architecture เดิม) และเริ่มต้นค่า weight เหมือนเดิมเป็น W_0 แต่ถูก mask ด้วย m (W_0 * m)
5. วน step 2-step 4 ไปเรื่อยๆจนได้ขนาด model ที่เราต้องการ

Iterative Magniture Pruning ภาพจาก Robert Lange

จะสังเกตได้ว่า algorithm นี้เป็น algorithm แบบ iterative กล่าวคือต้องทำซ้ำหลายๆรอบ จึงอาจจะทำให้หลายคนสงสัยได้ว่าทำไมเราไม่ทำแบบ one-shot (prune และ train ใหม่แบบรอบเดียวจบ)? จำเป็นไหมที่ weight ต้องเป็น W0 เท่านั้น? ทำไมไม่ต้อง random ค่าใหม่?

ในงานวิจัยได้มีการกล่าวถึงผลการทดลองว่าการใช้ weight เริ่มต้นของ subnetwork เป็นค่าเดิม และการ prune แบบ iterative นั้นจำเป็นเพื่อที่จะทำให้การ pruning ประสบผลสำเร็จอ้างอิงได้จากผลการทดลองดังนี้

ผลการทดลองโดยการใช้ LeNet บน MNIST ภาพจาก Frankle & Carbin 2018

• Winning Ticket หมายถึงการให้ subnetwork เริ่มต้นค่า weight เป็น W0 เหมือนเดิม
• Random Reinit หมายถึงการให้ subnetwork เริ่มต้น weight ใหม่
• Oneshot คือ prune รอบเดียวละ train ใหม่รอบเดียว
• Iteartive คือ prune ด้วย IMP

ข้อสังเกต

• การเริ่มต้น weight ของ subnetwork ด้วยค่าเดิมจาก original network ให้ accuracy ดีกว่า random ค่า weight ใหม่ทั้งใน IMP และ one-shot
• การ prune แบบ IMP นั้นจำเป็นต้องให้ subnetwork เริ่มต้นค่า weight ด้วย W0 เหมือน original network (สังเกตจากการที่ IMP ด้วย random reinit ให้ผลที่แย่ที่สุด)

Pruning with Rewinding

ในปีถัดมา Frankle ได้ทำงานวิจัยในเรื่อง Lottery Tickety Hypothesis ต่ออีกครั้งหนึ่งในหัวข้อ Stabilizing the Lottery Ticket Hypothesis โดยงานวิจัยนี้ได้กล่าวไว้ว่า การ prune ด้วย IMP ใน neural network ที่มีจำนวน layer เยอะจะทำให้ผลที่ไม่ดีนัก เว้นแต่ จะทำการ schedule learning rate ให้ดีด้วยการทำ warmup learning rate สังเกตได้จากผลการทดลองดังนี้

IMP ด้วยการ schedule และ ไม่ schedule learning rate บน CIFAR-10 ภาพจาก Frankle et al., 2019

จากกราฟ จะเห็นได้ว่าการ prune ด้วย IMP และใช้ learning rate schedule ด้วยการ warmup จะทำให้ winning ticket ได้ accuracy ที่สูงกว่า fix learning rate มาก จากข้อสังเกตตรงนี้ทำให้เกิดสมมติฐานว่า การ train subnetwork จะ sensitive กับ noise ในช่วงแรกๆของการ training เท่านั้น กล่าวอีกนัยนึงคือ subnetwork ต้อง train ไปซักพักนึงก่อนจะสามารถ train บน learning rate ที่สูงได้

จากสมมติฐานนี้ Frankle ได้เสนอว่าแทนที่เราจะเริ่มต้นค่า weight ของ subnetwork ด้วยการ mask ค่าเริ่มต้นของ original network เราจะเปลี่ยนมา mask weight หลังจากถูก train ไป k iteration แทน

1. สร้าง neural network ที่มี parameter เริ่มต้น เป็น W_0
2. Train neural network จำถึง iteration k ได้ weight W_k
3. Train neural network จนครบ iteration ได้ W_T
4. สร้าง mask (m) ซึ่งเป็น tensor ขนาดเท่า W โดยจะเป็นตัวกำหนดว่า ค่า weight ตัวไหนใน W ควรถูก set เป็น 0
5. สร้าง neural network ใหม่ (ด้วย architecture เดิม) และเริ่มต้นค่า weight เหมือนเดิมเป็น W_k แต่ถูก mask ด้วย m (W_k * m)
6. วน step 2-step 4 ไปเรื่อยๆจนได้ขนาด model ที่เราต้องการ

Rewinding IMP ภาพจาก Robert Lange

จะสังเกตได้ว่า algorithm นี้ไม่ต่างกับ IMP ปกติเพียงแต่แทนที่เราจะเริ่มต้นค่า weight ของ subnetwork ด้วย W_0*m เราจะเริ่มที่ W_k*m แทน

เพื่อยืนยันสมมติฐาน งานวิจัยนี้ได้ทำการทดลองด้วยการ train ResNet-20 บน CIFAR-10 และ ResNet-50 บน ImageNet เพื่อเปรียบเทียบประสิทธิภาพของการทำ rewinding ไปยัง iteration k เทียบกับ rewinding to iteration 0 (IMP แบบปกติ) ผลการทดลองออกมาดังนี้

รูปจาก Robert Lange

ข้อสังเกต

• การเริ่มต้นค่า weight ของ subnetwork นั้นไม่จำเป็นที่จะต้องเริ่มที่ W_0 * m เพื่อที่จะได้ winning ticket
• การเริ่มต้นค่า weight ของ subnetwork นั้นควรเริ่มในจุดที่ subnetwork นั้นทนต่อ learning rate ที่สูงได้ ซึ่งไม่ใช่ W_0 * m

อย่างไรก็ตาม การ pruning ด้วย IMP นั้นยังไม่เป็นที่ยอมรับมากนัก เพราะเราจำเป็นต้อง train subnetwork หลายรอบ ทำให้วิธีการนี้ค่อนข้างเปลือง computational resource เป็นไปได้ไหมที่เราจะสามารถหา winning ticket ได้อย่างมีประสิทธิภาพกว่านี้?

ตามหา ​Winning Tickets ด้วย Early Bird Tickets

จากปัญหาที่ว่า algorithm IMP นั้นค่อนข้างกินเวลาในการ train เพื่อที่จะได้ winning ticket มา ก็ได้มีงานวิจัยที่ชื่อว่า DRAWING EARLY-BIRD TICKETS: TOWARDS MORE EFFICIENT TRAINING OF DEEP NETWORKS ได้พยายามจะแก้ไขปัญหาในส่วนนี้ด้วยการตั้งคำถามว่า เราจำเป็นจะต้อง train neural network ให้ได้ accuracy สูงสุดก่อนเริ่ม prune จริงรึเปล่า?

งานวิจัยนี้ได้ทำการทดลองเพื่อที่จะพิสูจน์ว่า winning ticket สามารถหาได้ตั้งแต่ใน iteration แรกๆของการ train โดยไม่จำเป็นต้อง train จนได้ accuracy สูงสุดก่อน prune

Hamming Distance ของ mask ที่ได้จากการ prune ใน iteration ที่ต่างกัน ภาพจาก You et al., 2020

heatmap นี้แสดงถึงความแตกต่างของ mask ที่ได้จากการ prune model ในแต่ละ epoch โดย พิกัด (x, y) ในแต่ละจุดของ heatmap แสดงถึงความต่างกันระหว่าง mask ที่ได้จากการ prune model ที่ train ไป x epoch และ y epoch ซึ่งหากยิ่งพิกัดนั้นสีเข้ม หมายความว่า mask นั้นต่างกันมาก

ข้อสังเกต

• mask ที่ได้จากใน iteration แรกๆนั้นมีความต่างกันมากสังเกตได้ว่าในช่วง ~10 epochs heatmap จะมีสีเข้ม บ่งบอกว่าหาก prune neural network ในช่วง epoch แรกๆจะได้ mask ที่ไม่เหมือนกันเลย
• ในช่วง epoch ราวๆ 10–80 mask ที่ได้เริ่มจะไม่ค่อยต่างกันเท่าไร สังเกตได้จากการที่ heatmap มีสีเขียวอ่อนใกล้เหลืองขึ้นเรื่อยๆ
• หลัง epoch 80 mask ที่ได้ แทบจะไม่มีความต่างกันอีกต่อไป

ข้อสังเกตเหล่านี้เป็นตัวยืนยันว่าเราไม่จำเป็นต้อง prune neural network หลังจาก train network ไปจนถึง epoch ที่ accuracy สูงสุด เราสามารถ prune ได้ตั้งแต่ epoch แรกๆเพราะ mask ที่ได้จะไม่ต่างกันมากนัก โดย mask ที่เริ่มคงที่ใน iteration แรกๆ จะชื่อว่า Early-bird ticket นั่นเอง และงานวิจัยนี้ก็ได้เสนอ algorithm สำหรับการหา early-bird ticket ไว้ดังนี้

Early Bird Ticket Search Algorithm (You et al., 2019)

นอกจากนี้ข้อแตกต่างที่สำคัญของการ prune ด้วย IMP และ การ prune เพื่อหา Early Bird Ticket คือ prune เพื่อหา Early Bird Ticket จะ prune CNN filters ตามค่า Batch Normalization scaling factor แทน (รายละเอียดเพิ่มเติมสามารถหาได้จากงานวิจัยนี้)

ด้วยเทคนิคนี้ทำให้การหา winning ticket นั้นเร็วขึ้นมากโดยไม่ต้อง train หลายรอบอีกต่อไป! ซึ่งผลการทดลองนั้นค่อนข้างเป็นที่น่าสนใจทีเดียว

ผลการทดลองโดย paper จาก You et al., 2019

ข้อสังเกต

• การหา winning ticket ด้วย IMP (จุดสีดำ) ใช้ resource เยอะกว่า EB ticket ถึงสองเท่าโดยที่ Early bird ticket มี accuracy ที่ใกล้กันหรือดีกว่า

ส่งท้าย

medium นี้ผมได้เขียนขึ้นมาโดยสรุปจาก article ของคุณ Robert Tjarko Lange ในหัวข้อ The Lottery Ticket Hypothesis: A Survey หากใครสนใจและอยากหาข้อมูลเกี่ยวกับการ pruning และ Lottery Ticket Hypothesis ผมแนะนำให้อ่านบทความนี้ต่อได้เลยนะครับ

ขอบคุณที่ตามอ่านมาจนจบนะครับ หากมีข้อผิดพลาดหรือมีข้อเสนอแนะ ก็สามารถแสดงความคิดเห็นมาได้เลยนะครับ ยินดีที่จะรับฟังและนำไปแก้ไขครับ : )

บทความนี้จัดทำโดยสถาบันวิจัยปัญญาประดิษฐ์ประเทศไทย (AIResearch)

Facebook & Medium: AIResearch.in.th

AIResearch.in.th

Reference

1. Frankle, J., and M. Carbin. (2019): “The lottery ticket hypothesis: Finding sparse, trainable neural networks,” arXiv preprint arXiv:1803.03635, .
2. Frankle, J., G. K. Dziugaite, D. M. Roy, and M. Carbin. (2019): “Stabilizing the lottery ticket hypothesis,” arXiv preprint arXiv:1903.01611, .
3. Lange, R. T. (2020): “The Lottery Ticket Hypotheses: A Survey” https://roberttlange.github.io/year-archive/posts/2020/06/lottery-ticket-hypothesis/.
4. Lapthawan, T. (2019): “Lottery Ticket Hypothesis: จุดเปลี่ยนความคิดเกี่ยวกับ neural network” https://medium.com/@natlp/lottery-ticket-hypothesis-จุดเปลี่ยนความคิดเกี่ยวกับ-neural-network-43782c6bdcc3
5. You, H., C. Li, P. Xu, Y. Fu, Y. Wang, X. Chen, Y. Lin, Z. Wang, and R. G. Baraniuk. (2020): “Drawing early-bird tickets: Towards more efficient training of deep networks,” arXiv preprint arXiv:1909.11957, .