Mixture-of-Depths: teknik untuk mengalokasikan sumber daya komputasi secara dinamis di Transformer

• Mixture-of-Depths (MoD) dari Google DeepMind adalah pendekatan yang membuat model bahasa Transformer tidak menggunakan FLOPs yang sama untuk semua token, melainkan hanya mengikutsertakan token yang diperlukan di tiap layer dalam komputasi attention dan MLP
• Router per layer membuat bobot skalar untuk setiap token, lalu hanya token dalam kapasitas top-k yang sudah ditentukan yang melewati blok, sementara sisanya melewati jalur pintas melalui residual connection
• Dengan menetapkan k sejak awal, computation graph dan ukuran tensor tetap statis, tetapi alokasi komputasi per token berubah secara dinamis sesuai konteks
• Dalam eksperimen, berdasarkan isoFLOP, MoD dapat menghasilkan loss lebih rendah dibanding transformer dasar pada waktu pelatihan yang sama, atau mengurangi FLOPs per forward pass pada performa yang sama sehingga step pelatihan dan sampling menjadi lebih cepat
• Konfigurasi terbaik adalah menerapkan blok berkapasitas 12,5% pada blok berselang-seling (every other block), dan pada auto-regressive sampling masalah non-kausalitas top-k dikurangi dengan routing berbasis predictor

---

Pemborosan komputasi yang ditargetkan MoD

• Model bahasa transformer umum menggunakan jumlah komputasi yang sama untuk semua token dalam forward pass
• MoD berangkat dari premis bahwa tidak semua token dan sequence membutuhkan waktu serta upaya yang sama untuk prediksi, lalu mengalokasikan komputasi ke token yang diperlukan sambil mengurangi total anggaran komputasi
• Conditional computation yang sudah ada dapat membuat computation graph dinamis, sehingga mungkin tidak cocok dengan computation graph statis yang disukai hardware saat ini
• Jika pengguna menentukan total anggaran komputasi sebelum pelatihan, model belajar di mana harus memakai komputasi per token dan per layer dalam batas tersebut
  • Jumlah komputasi, penggunaan memori, dan FLOPs per forward pass dapat diprediksi sejak awal
  • Token yang menjadi target routing berubah sesuai konteks input

Struktur routing dan cara pengurangan komputasi

• MoD menggunakan router seperti MoE transformer, tetapi alih-alih memilih salah satu dari beberapa expert, ia mengirim token ke salah satu dari dua jalur
  • Komputasi blok standar yang mencakup attention dan MLP
  • Residual connection yang meneruskan nilai apa adanya
• Residual connection memiliki biaya komputasi rendah, dan output blok ditentukan oleh nilai input
• Jika capacity blok ditetapkan lebih kecil dari jumlah total token T, FLOPs per forward pass menjadi lebih rendah dibanding transformer dasar
  • Misalnya, jika capacity dikurangi menjadi T/2, perkalian matriks query-key pada attention menjadi (T/2)², bukan T², sehingga FLOPs menjadi sekitar 25%
• Jika capacity dibuat 0, model menjadi sangat cepat tetapi tidak memakai sebagian besar parameter transformer sehingga performanya memburuk; jika capacity dibuat T, hasilnya sama dengan transformer dasar
• Tujuannya adalah menemukan titik di antara dua ekstrem ini yang dapat meningkatkan kecepatan step sambil menghasilkan performa yang sama atau lebih baik daripada transformer dasar

Routing top-k expert-choice

• Pada setiap layer l, router menghasilkan bobot skalar untuk setiap embedding token
• Sesuai capacity C yang ditentukan pengguna, hanya token yang termasuk top-k di antara token sequence yang ikut dalam komputasi attention dan MLP
• Token yang tidak dipilih melewati komputasi blok dan diteruskan ke layer berikutnya melalui residual connection
• Pendekatan ini berbasis expert-choice routing
  • Pada token-choice routing, token memilih jalurnya sendiri sehingga dapat timbul masalah load balancing
  • Pada expert-choice routing, setiap jalur memilih token top-k sehingga capacity terisi tepat
• Dalam MoD, jalurnya pada dasarnya hanya dua: “komputasi blok” dan “bypass”, sehingga satu operasi top-k dapat membagi token menjadi dua set
• Bobot router dikalikan ke output blok sehingga berada di jalur gradient, dan dipengaruhi oleh gradient descent selama proses pelatihan language modeling

Masalah non-kausalitas pada sampling

• Operasi top-k pada expert-choice routing bersifat non-kausal (non-causal)
  • Untuk menentukan apakah suatu token masuk top-k, bobot router dari token-token setelahnya juga harus diketahui
  • Pada auto-regressive sampling, informasi token masa depan tidak dapat digunakan
• Untuk menghindari hal ini, dua metode diuji
  • Auxiliary loss binary cross-entropy yang membuat sigmoid dari output router memprediksi apakah token akan dipilih top-k
  • MLP predictor kecil tambahan yang menerima input yang sama dengan router tetapi menerapkan stop-gradient
• Metode pertama memengaruhi objective utama language modeling sekitar 0,2–0,3%
• Metode kedua yang berbasis predictor tidak memengaruhi objective language modeling dan tidak berdampak signifikan pada step speed
• Tugas prediksi tambahan secara empiris cepat mencapai akurasi 99%

Hasil eksperimen dan observasi

• Eksperimen awal dilakukan dengan mencari hiperparameter MoD pada anggaran total 6e18 FLOPs
• Konfigurasi MoD optimal adalah menerapkan routing pada setiap blok berselang-seling dan menetapkan top-k ke 256
  • Dengan panjang sequence 2048, hanya 256 token, yaitu 12,5%, yang melewati attention dan MLP
  • Sisa 1792 token, yaitu 87,5%, melewati blok melalui jalur bypass
• Varian MoD dengan 220M parameter menunjukkan performa sedikit lebih baik daripada baseline optimal isoFLOP, sementara step speed selama pelatihan lebih dari 60% lebih cepat, dan total wall-clock time pelatihan pada hardware yang sama hampir sama
• Dalam analisis isoFLOP pada 6e18, 2e19, dan 1e20 FLOPs, model FLOP-optimal MoD juga cenderung memiliki lebih banyak parameter daripada baseline
• Ada varian-varian MoD yang menghasilkan performa sama atau loss lebih rendah dengan FLOPs per forward pass yang lebih kecil
• Peningkatan step speed berasal dari dua faktor
  • Karena sebagian token melewati blok melalui bypass, FLOPs per parameter berkurang pada ukuran model yang sama
  • Dapat dipilih varian MoD yang lebih kecil tetapi menghasilkan performa sama dengan baseline
• MoD yang menggunakan stochastic routing menunjukkan performa jauh lebih rendah dibanding transformer dasar dan MoD umum, sehingga routing yang dipelajari terbukti penting
• Dalam evaluasi auto-regressive, digunakan held-out data berisi 256.000 sequence dan 500M token, dan penurunan performa tetap kecil meski diganti ke routing berbasis predictor
• MoD juga dapat digabungkan dengan MoE menjadi Mixture-of-Depths-and-Experts (MoDE)
  • Staged MoDE menentukan sebelum attention apakah token akan dikirim ke blok atau dibypass
  • Integrated MoDE mengintegrasikan routing MoD dengan menambahkan expert “no-op” di antara expert MLP yang sudah ada
• Dalam analisis routing, diamati pola bahwa sebagian token sering melewati beberapa blok, sementara token lain melakukan bypass kapan pun memungkinkan
• Analisis awal menunjukkan token yang sering melewati blok berkorelasi dengan kasus ketika entropy prediksi output lebih tinggi, yang terkait dengan kemungkinan bahwa token tersebut lebih sulit diprediksi

[GN⁺]

Komentar Hacker News

• Routing yang lebih kompleks tampaknya pasti akan makin umum
  Khususnya, suatu saat nanti saya pikir arahnya akan ke routing rekursif yang melewatkan kembali sekumpulan pakar. Ke depannya, 'chain-of-thought' sepertinya akan terjadi secara rekursif di dalam model

  • Objek hipotetis seperti ini bisa disebut Recursive Neural Networks
  • Yang digambarkan di sini terlihat agak mirip dengan rangkaian riset Universal Transformers. Caranya adalah melewatkan embedding input beberapa kali melalui satu blok transformer, lalu modul terpisah menilai apakah embedding itu sudah cukup “matang” untuk dikeluarkan
    Yang lebih dekat dengan gagasan “pakar” adalah makalah Sparse Universal Transformers tahun lalu, yang menggabungkan Universal Transformer dengan sparse mixture of experts, sehingga mekanisme gating menentukan blok transformer mana yang dipakai dan dalam urutan apa
    Ini bukan bidang keahlian saya, tetapi setahu saya cukup sulit dilatih dengan benar, dan untuk menghasilkan keluaran yang mirip transformer biasa, saat inferensi dibutuhkan total komputasi yang lebih besar. Meski begitu, ini arah yang menarik, dan menurut saya batas atas jumlah langkah komputasi per token adalah salah satu kelemahan besar dari arsitektur transformer klasik
  • Menurut saya alasan ini belum berhasil adalah karena pada saat pelatihan tidak ada cara untuk menentukan berapa kali rekursi harus dilakukan
    Jika memilih jumlah acak atau mencoba beberapa kedalaman rekursi, keluarannya menjadi “kabur”. Artinya, menjadi tidak jelas apakah keluaran suatu layer harus memberikan informasi penting untuk hasil akhir, atau harus memberikan input terbaik untuk putaran rekursi berikutnya
  • Attention pada dasarnya adalah routing, dan metode routing lain seperti ini memberi model opsi yang lebih kasar sehingga berpotensi membuat pelatihan lebih mudah
  • Arahnya jelas menuju peningkatan routing dinamis, tetapi menurut saya MoE/MoD/MoDE lebih cenderung memungkinkan penyimpanan fakta tambahan dengan mengurangi tumpang tindih di dalam bobot, bukan membuat penalaran yang lebih dalam
    Penalaran yang lebih dalam sepertinya akan datang dari kedinamisan tingkat token, bukan tingkat layer. Misalnya ada makalah Quiet-STaR terbaru yang menghasilkan token alasan yang kemudian dibuang oleh model: https://arxiv.org/abs/2403.09629

• Ini mungkin makalah terpenting tahun 2024
  Gagasan bahwa kita menginginkan model yang tidak perlu memakai jumlah komputasi yang sama untuk semua token sudah ada sejak lama, tetapi ini pertama kalinya saya melihat mekanisme yang meyakinkan untuk itu

  Equipped with these new methods, we can sample autoregressively by choosing to route tokens to or around a block based on the router’s output, which does not depend on any information from future tokens. We provide empirical evidence that this is a relatively easy auxiliary task that quickly achieves 99% accuracy.
  Bukankah ini agak mengejutkan?

  • Sparse Universal Transformer lebih lama, dan sudah melakukan early exit berbasis routing
  • Sampai dibilang paling penting? Bahwa tidak semua token membutuhkan seluruh jendela konteks seharusnya merupakan optimisasi yang obvious

• Versi pengantar sederhana:
  Bayangkan ada asisten pintar yang bisa memahami dan memproses ucapan. Biasanya asisten ini memberi perhatian yang sama pada semua kata, terlepas dari seberapa penting tiap kata bagi makna keseluruhan
  Sekarang bayangkan kita menemukan cara mengajarkan asisten itu memakai “sumber daya otak” dengan lebih cerdas. Alih-alih memberi perhatian yang sama pada semua kata, ia lebih fokus pada kata-kata yang paling penting untuk memahami makna, dan menyesuaikan fokus itu secara langsung sesuai konteks
  Agar asisten tidak kelebihan beban, kita juga membatasi total “sumber daya otak” yang bisa dipakai pada saat tertentu. Ibaratnya diberi anggaran dan diberi tahu, “kamu hanya bisa memakai sumber daya untuk sejumlah kata tertentu sekaligus.” Maka asisten harus memutuskan kata mana yang paling penting
  Meski ada batasan ini, asisten tetap menyesuaikan penggunaan sumber dayanya secara fleksibel. Ia memakai lebih banyak untuk beberapa kata, lebih sedikit untuk kata lain, sambil memprioritaskan sesuai situasi
  Dengan pelatihan seperti ini, asisten bekerja efisien sambil tetap memberi perhatian secara cerdas. Ia memahami sama baiknya dengan asisten yang memberi perhatian sama pada semua kata, tetapi memakai total sumber daya otak lebih sedikit, sehingga respons dan pemrosesan informasi baru menjadi jauh lebih cepat

  • Saya tahu ini ELI5, tetapi bukankah cara yang dijelaskan itu sudah dilakukan oleh attention? Karena ia memang secara spesifik berfokus pada kata-kata yang penting secara kontekstual dalam sekuens sebelumnya

• Berdasarkan pemahaman saya, saya menuliskan sedikit rangkuman di sini:
  https://lifeinthesingularity.com/p/googles-breakthroughs-in-...

  • Tulisannya bagus. Mengingatkan pada gaya New Scientist. Rasanya seperti “menjelaskan sambil berjalan tetapi tetap singkat”, jadi bagus untuk mendapatkan gambaran bidang ini

• Sangat mirip dengan mixture of experts. Bedanya, alih-alih merutekan token ke beberapa pakar, ini seperti “menempatkannya pada satu pakar yang bisa dilewati secara dinamis”

  • Menggabungkan keduanya sepertinya akan cukup keren. Bisa makin mengurangi komputasi MoE sambil mempertahankan performa

• “This is more computationally efficient than performing a full content-based lookup across an entire memory buffer for each step in the future, and could be one step towards drastically increasing the context-length available for making a prediction.”
  Apakah ini cara yang memungkinkan jendela konteks 10 juta token? Atau maksudnya jendela konteks yang lebih panjang di masa depan?

• Kesimpulan yang saya dapat setelah mencoba memahami dan mengimplementasikan beberapa algoritma RASP adalah bahwa fungsi tertentu membutuhkan sejumlah layer transformer tertentu agar bisa bekerja
  Jika mengikuti logika ini, terlihat bahwa fungsi-fungsi yang dipelajari transformer bisa tersebar di banyak head. Mengulangi fungsi seperti ini bisa sangat bernilai untuk memahami dan memecahkan masalah, tetapi cara inferensi saat ini tidak bisa menjalankan ulang sekumpulan head berurutan. Makalah ini benar-benar terlihat sebagai arah yang menjanjikan
  [1] https://arxiv.org/pdf/2106.06981.pdf
  [2] https://www.youtube.com/watch?v=t5LjgczaS80

• Satu-satunya kekurangan dari perkembangan LLM yang terlalu cepat mungkin adalah makalah keluar lebih cepat daripada kemampuan siapa pun—setidaknya di luar Google—untuk mempelajari dan menguji peningkatannya
  Ketika saya mulai belajar deep learning, ReLU dan dropout sedang populer, dan di GPU konsumen 1080 kita bisa mengubah satu-dua baris kode lalu menguji apakah ada peningkatan dalam beberapa jam. Sekarang rasanya harus menunggu beberapa minggu sampai tempat seperti Mistral mencobanya

  • Selamat datang di kaum miskin GPU
    Saya berfokus pada pendekatan kuantisasi dan mengujinya dengan GPU lama generasi sebelumnya

• hu-po mengulas makalah AI secara live streaming mendalam
  Sangat direkomendasikan, dan di sini ia membahas makalah mixture-of-depths yang sedang dibicarakan: https://www.youtube.com/watch?v=Teru_qIdB8Y