Menjelaskan Attention Residuals: Menata Ulang Kedalaman Transformer

LLM modern menjadi semakin dalam, lebar, dan haus komputasi, namun menumpuk lebih banyak lapisan Transformer tidak selalu berbanding lurus dengan peningkatan kinerja. Salah satu alasannya adalah karena koneksi residual standar menggabungkan keluaran lapisan dengan bobot unit tetap, sehingga setiap lapisan mewarisi penjumlahan seragam dari semua yang mendahuluinya. Ini dapat mengencerkan representasi awal, memperbesar magnitudo state tersembunyi, dan menyulitkan jaringan untuk kembali menggunakan fitur perantara yang paling bermanfaat secara selektif.

Alih-alih memperlakukan kedalaman sebagai rekuren aditif tetap, makalah dari tim Kimi tentang Attention Residuals memungkinkan setiap lapisan melakukan attention terhadap keluaran lapisan sebelumnya menggunakan bobot softmax yang dipelajari. Dalam artikel ini, kita akan memahami mengapa agregasi residual standar menjadi hambatan, bagaimana Attention Residuals bekerja, mengapa varian blok penting, dan apa yang sebenarnya disarankan hasilnya tentang penskalaan model bahasa yang lebih dalam.

Jika Anda ingin mempelajari lebih lanjut tentang beberapa gagasan di balik arsitektur Transformer modern, saya sarankan mengikuti kursus DataCamp Tutorial Model Transformer di PyTorch.

Masalah: Sinyal yang Tenggelam

Koneksi residual telah menjadi fondasi pembelajaran mendalam selama bertahun-tahun. Pada Transformer, pembaruan residual standar adalah:

hl=hl−1​+fl−1​(hl−1​)

Ini membantu aliran gradien melalui jaringan yang sangat dalam. Namun, residual bukan hanya trik gradien. Residual juga menentukan bagaimana informasi diagregasikan sepanjang kedalaman. Jika kita uraikan rekuren ini, kita dapatkan:

hl​=h1​+∑i=1l−1​fi​(hi​)

Ini berarti state tersembunyi pada lapisan l hanyalah embedding dan merupakan penjumlahan berbobot seragam dari semua keluaran lapisan sebelumnya. Jadi, setiap kontribusi pada dasarnya diberi bobot yang sama.

Ini menjadi masalah saat skala membesar. Makalah ini menunjukkan bahwa pada arsitektur PreNorm, akumulasi residual tak berbobot menyebabkan magnitudo state tersembunyi tumbuh seiring kedalaman, kira-kira sebagai O(L). 

Seiring aliran residual membesar, keluaran lapisan awal menjadi semakin encer di dalam penjumlahan berjalan yang lebih besar. Setelah sinyal lapisan awal tercampur dalam akumulasi ini, lapisan yang lebih dalam tidak dapat memulihkannya secara selektif. Akibatnya, mereka hanya beroperasi pada state teragregasi.

Ini mengarah pada apa yang disebut makalah tersebut sebagai efek “sinyal yang tenggelam”. Bukti empiris kuat dari inefisiensi ini datang dari studi pemangkasan lapisan, yaitu sebagian besar lapisan pada model terlatih sering kali dapat dihapus dengan dampak minimal pada kinerja. Ini menunjukkan bahwa saat kita terus menambah kedalaman, model kurang memiliki mekanisme efektif untuk memanfaatkannya sepenuhnya. Alih-alih membentuk rantai penalaran hierarkis, jaringan berperilaku lebih seperti relai berlebih di mana sinyal awal semakin diencerkan.

Ada juga trade-off pada lintasan maju. Ketika aliran residual membesar, lapisan selanjutnya mungkin perlu menghasilkan keluaran bermagnitudo lebih besar agar dapat memengaruhi state terakumulasi secara bermakna. Makalah ini menghubungkannya dengan perilaku PreNorm, di mana magnitudo state tersembunyi meningkat secara monoton seiring kedalaman di bawah agregasi residual standar.

Koneksi residual membantu aliran gradien, dan gradien terhadap state tersembunyi perantara dapat ditulis sebagai:

Suku identitas di sini mempertahankan jalur gradien langsung. Namun residual tetap memaksa jalur agregasi maju untuk memperlakukan setiap lapisan sebelumnya dengan bobot tetap 1,0. Itulah keterbatasan struktural yang coba diperbaiki oleh AttnRes.

Dari Residual Aditif ke Attention Sepanjang Kedalaman

Langkah konseptual paling menarik dalam makalah ini adalah dualitas waktu–kedalaman. Koneksi residual memampatkan informasi sepanjang kedalaman, sebagaimana jaringan rekuren memampatkan informasi sepanjang waktu. Dalam pemodelan sekuens, attention menggantikan rekuren dengan memungkinkan setiap posisi mengakses posisi sebelumnya secara selektif. AttnRes menerapkan transisi yang sama pada kedalaman jaringan. Alih-alih mendefinisikan state tersembunyi berikutnya sebagai jumlah tetap atas lapisan sebelumnya, AttnRes memungkinkan setiap lapisan melakukan attention terhadap keluaran lapisan terdahulu:

di mana bobot αi→l adalah bobot attention softmax sepanjang kedalaman dan dijumlahkan menjadi 1. Bobot-bobot tersebut dihitung sebagai:

dengan makalah menggunakan:

Setiap lapisan memperoleh vektor pseudo-query terpelajar wl, dan query tersebut melakukan attention terhadap key dan value yang dibangun dari keluaran lapisan sebelumnya. RMSNorm diterapkan pada key agar lapisan dengan magnitudo keluaran yang secara alami lebih besar tidak mendominasi softmax hanya karena skalanya lebih besar.

Perubahan kunci dalam makalah ini adalah residual standar diperlakukan sebagai semacam attention linear sepanjang kedalaman, sementara AttnRes meningkatkannya menjadi attention softmax sepanjang kedalaman. Alih-alih akumulasi seragam, kita mendapatkan pengambilan selektif di sepanjang kedalaman.

Detail implementasi kecil namun penting adalah inisialisasi. Para penulis menyarankan semua vektor pseudo-query diinisialisasi ke nol, yang memastikan bobot attention awal bersifat seragam, sehingga model memulai pelatihan sebagai rata-rata bobot sama alih-alih mekanisme attention yang bias secara acak.

Gambar 1: Gambaran umum Attention Residuals (makalah Attention Residuals)

Full Attention Residuals

Dalam Full Attention Residuals, setiap lapisan melakukan attention terhadap semua keluaran lapisan sebelumnya, memungkinkannya mengagregasikan informasi secara selektif dari seluruh kedalaman. Ini memberi model fleksibilitas maksimum, sehingga lapisan yang lebih dalam dapat menekankan pendahulu terdekatnya, embedding awal, atau lapisan lain mana pun jika di sanalah sinyal paling berguna berada.

Dalam pelatihan standar, overhead memori Full AttnRes lebih kecil daripada yang tampak, karena banyak keluaran lapisan sudah disimpan untuk backpropagation. Namun pelatihan skala besar mengubah gambaran. Setelah Anda memperkenalkan komputasi ulang aktivasi dan paralelisme pipeline, keluaran awal tersebut harus disimpan dan dikomunikasikan secara eksplisit agar lapisan selanjutnya dapat melakukan attention terhadapnya, dan itu menjadi mahal.

Jadi Full AttnRes adalah cara terbaik untuk memahami gagasan inti, tetapi bukan versi yang paling ingin digunakan tim pada skala besar.

Block Attention Residuals

Makalah ini memperkenalkan gagasan Block Attention Residuals untuk mengurangi overhead memori dan komunikasi. Lapisan-lapisan model dipartisi menjadi blok sehingga di dalam setiap blok, keluaran digabung menggunakan akumulasi residual aditif standar, tetapi lintas blok, model melakukan attention terhadap ringkasan tingkat blok alih-alih setiap lapisan sebelumnya satu per satu.

Jika Bn adalah himpunan lapisan dalam blok n, maka representasi blok diberikan oleh:

Model kemudian melakukan attention terhadap embedding b0=h1 ringkasan blok sebelumnya, dan jumlah parsial blok saat ini seiring kemajuan komputasi. Ini mengurangi memori dan komunikasi dari O(Ld) menjadi O(Nd), di mana N adalah jumlah blok.

Makalah ini melaporkan bahwa menggunakan sekitar delapan blok memulihkan sebagian besar manfaat dari versi penuh, dan kesenjangan kinerja antara Full AttnRes dan Block AttnRes menyempit seiring peningkatan skala. Ukuran blok seperti S=2,4,8 tetap mendekati versi penuh, sementara pengelompokan yang jauh lebih kasar cenderung kembali ke perilaku baseline.

Inti temuan paling berguna dari makalah ini adalah Anda tidak memerlukan attention sepanjang kedalaman penuh terhadap setiap lapisan untuk memperoleh sebagian besar keuntungan. Sementara Block AttnRes membuat attention sepanjang kedalaman menjadi layak secara komputasional, menerapkannya secara efisien pada skala besar tetap memerlukan perancangan sistem yang cermat. Makalah ini juga memperkenalkan beberapa optimisasi yang membuat AttnRes praktis untuk pelatihan dan inferensi dunia nyata.

Gambar 2: Contoh komunikasi pipeline berbasis cache dengan 4 peringkat fisik dan 2 tahap virtual per peringkat, di mana kotak arsir menandakan akhir blok AttnRes(makalah Attention Residuals)

Berikut tiga optimalisasi sistem kunci yang membuat AttnRes dapat menjadi pengganti langsung residual standar:

• Komputasi dua fase: Dengan memisahkan pseudo-query dari state tersembunyi, kita dapat membatch query untuk semua lapisan dalam satu blok. Fase 1 menghitung attention antarblok secara paralel, sementara Fase 2 menangani dependensi intrablok secara berurutan menggunakan Online Softmax Merge. Ini secara signifikan mengurangi I/O per lapisan dibandingkan metode multi-stream sebelumnya seperti mHC.
• Caching lintas tahap: Dalam paralelisme pipeline, kami menghilangkan transfer data yang berulang dengan melakukan cache blok yang sebelumnya diterima secara lokal. Ini memberikan peningkatan V× (di mana V adalah jumlah tahap virtual seperti terlihat pada Gambar 2) pada biaya puncak per transisi, memungkinkan komunikasi ditumpangtindihkan sepenuhnya dengan komputasi.
• Prefilling hemat memori: Dengan membagi representasi blok ke perangkat paralel-tensor, jejak memori untuk sekuens konteks panjang (128K+) berkurang secara signifikan (misalnya, dari 15GB menjadi 0,3GB per perangkat).

Hasil dan Analisis Kinerja

Makalah ini memvalidasi AttnRes pada beberapa tingkat, seperti hukum penskalaan, ablation, dinamika pelatihan, dan tolok ukur hilir.

Gambar 3: Kurva hukum penskalaan untuk Attention Residuals. (makalah Attention Residuals)

Namun, hasil utama berasal dari eksperimen hukum penskalaan. Di lima ukuran model, baik Full AttnRes maupun Block AttnRes secara konsisten mencapai loss validasi yang lebih rendah daripada baseline. Berdasarkan kurva terpasang, Block AttnRes mencapai loss yang sama dengan baseline yang dilatih dengan sekitar 1,25× komputasi lebih banyak. Ini menunjukkan bahwa AttnRes tidak hanya lebih baik secara teori, tetapi juga lebih efisien secara komputasi.

Pada tugas hilir, AttnRes meningkat dibanding baseline di semua tolok ukur yang dievaluasi. Beberapa peningkatan terbesar terdapat pada tugas yang intensif penalaran, seperti ditunjukkan pada Tabel 1:

Tabel 1: Perbandingan kinerja AttnRes dengan baseline, keduanya setelah resep pra-pelatihan yang sama (makalah Attention Residuals)

Keuntungan sangat kuat pada penalaran multi-langkah, matematika, dan kode, yang sejalan dengan hipotesis inti makalah bahwa jika lapisan selanjutnya dapat mengambil kembali representasi awal secara selektif alih-alih mewarisi agregat yang kabur, penalaran komposisional akan membaik.

Dibandingkan baseline, AttnRes menunjukkan loss validasi yang lebih rendah sepanjang pelatihan, magnitudo keluaran yang lebih terkendali di seluruh kedalaman, dan magnitudo gradien yang lebih seragam di seluruh lapisan. Baseline menderita pola pengenceran PreNorm yang umum, di mana magnitudo state tersembunyi meningkat secara monoton seiring kedalaman, sementara Block AttnRes menghasilkan pola yang lebih terkendali dan terbatas berkat agregasi selektif di batas blok.

Ada juga temuan tingkat arsitektur yang lebih luas bahwa, dalam sapuan komputasi dan parameter tetap, konfigurasi optimal untuk AttnRes bergeser ke model yang lebih dalam dan lebih sempit dibanding baseline. Ini menunjukkan bahwa AttnRes dapat membuat penambahan kedalaman menjadi lebih bermanfaat daripada di bawah agregasi residual standar.

Kesimpulan

Residual sering diperlakukan sebagai trik pelatihan yang diperlukan, yakni shortcut identitas yang menjaga gradien tetap hidup. AttnRes menunjukkan bahwa pandangan ini terlalu sempit. Jalur residual juga merupakan mekanisme penyaluran informasi di sepanjang kedalaman, dan akumulasi aditif tetap mungkin terlalu sederhana untuk model yang sangat dalam.

Attention Residuals bukan tentang menambal masalah yang diketahui, melainkan meningkatkan pilihan desain yang selama ini terabaikan. Pemodelan sekuens berevolusi dari rekuren ke attention karena rekuren tetap terlalu membatasi. AttnRes berargumen bahwa kedalaman kini siap untuk transisi yang sama.

Model tetap menunjukkan lokalitas yang kuat, dengan lapisan sering kali paling memperhatikan pendahulu terdekat, namun juga mempelajari pola lompatan yang tidak sepele, mempertahankan bobot persisten pada embedding, dan menjaga perilaku yang berbeda antara lapisan pra-attention dan pra-MLP. 

Inti praktisnya bukan bahwa semua model harus mengganti residual dengan attention sepanjang kedalaman penuh. Alih-alih, ini membingkai ulang agregasi residual sebagai ruang desain arsitektural inti.