Tukar sepenuhnya model bahasa: TTT seni bina baharu mengatasi Transformer dan model ML menggantikan keadaan tersembunyi RNN

Prestasi model besar telah dipertingkatkan daripada 125M kepada 1.3B.

Sukar dipercayai, ini akhirnya berlaku.

Sebuah seni bina model bahasa besar (LLM) baharu dijangka menggantikan Transformer, yang popular dalam bidang AI setakat ini, dan prestasinya lebih baik daripada Mamba. Pada hari Isnin, kertas kerja mengenai Latihan Masa Ujian (TTT) menjadi topik hangat dalam komuniti kecerdasan buatan.

Pautan kertas: https://arxiv.org/abs/2407.04620

Pengarang kajian ini adalah dari Universiti Stanford, Universiti California, Berkeley, Universiti California, San Diego dan Meta. Mereka mereka bentuk seni bina baharu, TTT, yang menggantikan keadaan tersembunyi RNN dengan model pembelajaran mesin. Model memampatkan konteks melalui penurunan kecerunan sebenar token input.

Karan Dalal, salah seorang pengarang kajian, berkata beliau percaya ini secara asasnya akan mengubah pendekatan model bahasa.

Dalam model pembelajaran mesin, lapisan TTT secara langsung menggantikan Attention dan membuka kunci seni bina kerumitan linear melalui ingatan ekspresif, membolehkan kami melatih LLM dengan berjuta-juta (kadang-kadang berbilion) token dalam konteks .

Pengarang menjalankan satu siri perbandingan pada model besar dengan saiz parameter dari 125M hingga 1.3B dan mendapati bahawa kedua-dua TTT-Linear dan TTT-MLP boleh memadankan atau mengalahkan kaedah seni bina Transformer dan Mamba yang paling berkuasa.

Sebagai pemampatan maklumat baharu dan mekanisme ingatan model, lapisan TTT boleh secara ringkas dan terus menggantikan lapisan perhatian kendiri dalam Transformer.

Berbanding dengan Mamba, TTT-Linear mempunyai kebingungan yang lebih rendah, FLOP yang lebih sedikit (kiri) dan penggunaan konteks yang panjang (kanan) yang lebih baik:

Ini bukan sahaja linear dalam kerumitan teori berjalan dan kerumitan larian sebenar. masa juga lebih pantas.

• Selepas kertas masuk dalam talian, pengarang membuat kod dan jax awam untuk orang ramai melatih dan menguji: https://github.com/test-time-training/ttt-lm-jax
• Juga kod inferens PyTorch: https://github.com/test-time-training/ttt-lm-pytorch

Pengenalan kaedah

konteks panjang adalah dalam

sifat lapisan RNN: Tidak seperti mekanisme perhatian diri, lapisan RNN mesti memampatkan konteks ke dalam keadaan tersembunyi bersaiz tetap dan peraturan kemas kini perlu menemui struktur dan hubungan asas antara beribu-ribu atau berjuta-juta token.

Pasukan penyelidik mula-mula memerhatikan bahawa pembelajaran penyeliaan kendiri boleh memampatkan set latihan yang besar kepada pemberat untuk model seperti LLM dan model LLM sering mempamerkan pemahaman yang mendalam tentang hubungan semantik antara data latihan mereka.

Diinspirasikan oleh pemerhatian ini, pasukan penyelidik mereka bentuk kelas lapisan pemodelan jujukan baharu, di mana keadaan tersembunyi ialah model dan peraturan kemas kini ialah langkah pembelajaran diselia sendiri. Memandangkan proses mengemas kini keadaan tersembunyi pada jujukan ujian adalah bersamaan dengan melatih model pada masa ujian, pasukan penyelidik memanggil lapisan baharu ini lapisan Latihan Masa Ujian (TTT).

Pasukan penyelidik memperkenalkan dua contoh mudah: TTT-Linear dan TTT-MLP, di mana keadaan tersembunyi masing-masing adalah model linear dan MLP dua lapisan. Lapisan TTT boleh disepadukan ke dalam mana-mana seni bina rangkaian dan dioptimumkan hujung ke hujung, serupa dengan lapisan RNN dan perhatian diri.

🎜

Untuk menjadikan lapisan TTT lebih cekap, kajian ini menggunakan beberapa helah untuk menambah baik lapisan TTT:

Pertama, sama seperti mengambil langkah kecerunan untuk urutan kelompok mini semasa latihan biasa untuk mendapatkan keselarian yang lebih baik, kajian itu Gunakan kumpulan kecil token semasa TTT.

Kedua, kajian membangunkan bentuk dwi untuk operasi dalam setiap kumpulan mini TTT untuk menggunakan GPU dan TPU moden dengan lebih baik. Output bentuk dwi adalah bersamaan dengan pelaksanaan mudah, tetapi latihan adalah lebih daripada 5 kali lebih cepat. Seperti yang ditunjukkan dalam Rajah 3, TTT-Linear adalah lebih pantas daripada Transformer dan setanding dengan Mamba dalam konteks 8k.

Pasukan penyelidik percaya bahawa semua lapisan pemodelan jujukan boleh dilihat sebagai menyimpan konteks sejarah ke dalam keadaan tersembunyi, seperti yang ditunjukkan dalam Rajah 4.

Contohnya, lapisan RNN seperti lapisan LSTM, RWKV dan Mamba memampatkan konteks ke dalam keadaan saiz tetap merentas masa. Mampatan ini mempunyai dua akibat: Di satu pihak, memetakan token input x_t kepada token output z_t adalah cekap kerana peraturan kemas kini dan peraturan keluaran untuk setiap token memerlukan masa yang tetap. Sebaliknya, prestasi lapisan RNN dalam konteks yang panjang dihadkan oleh ekspresi keadaan tersembunyinya s_t.

Perhatian diri juga boleh dilihat dari perspektif di atas, kecuali keadaan tersembunyinya (selalunya dipanggil cache Nilai-Kekunci) ialah senarai yang berkembang secara linear dengan t. Peraturan kemas kininya hanya menambahkan tuple KV semasa pada senarai ini, manakala peraturan keluarannya mengimbas semua tupel sebelum t untuk membentuk matriks perhatian. Keadaan tersembunyi secara eksplisit menyimpan semua konteks sejarah tanpa pemampatan, yang menjadikan perhatian diri lebih ekspresif daripada lapisan RNN untuk konteks yang panjang. Walau bagaimanapun, masa yang diperlukan untuk mengimbas keadaan tersembunyi yang tumbuh secara linear ini juga berkembang secara linear. Untuk memastikan konteks yang panjang cekap dan ekspresif, penyelidik memerlukan heuristik mampatan yang lebih baik. Khususnya, beribu-ribu atau mungkin berjuta-juta token perlu dimampatkan ke dalam keadaan tersembunyi yang secara berkesan menangkap struktur dan perhubungan asasnya. Ini mungkin terdengar sukar, tetapi ramai orang sebenarnya sangat mengenali heuristik ini.

Seni bina tulang belakang. Cara paling bersih untuk menyepadukan mana-mana lapisan RNN ke dalam seni bina yang lebih besar adalah dengan menggantikan secara langsung perhatian diri dalam Transformer, di sini dipanggil tulang belakang. Walau bagaimanapun, RNN sedia ada (seperti Mamba dan Griffin) menggunakan lapisan tulang belakang yang berbeza daripada Transformer. Paling ketara, lapisan tulang belakangnya mengandungi lilitan temporal sebelum lapisan RNN, yang boleh membantu mengumpul maklumat setempat merentas masa. Selepas bereksperimen dengan tulang belakang Mamba, para penyelidik mendapati bahawa ia juga boleh meningkatkan kekeliruan lapisan TTT, jadi ia dimasukkan ke dalam kaedah yang dicadangkan, seperti yang ditunjukkan dalam Rajah 16.

Hasil eksperimen

Dalam percubaan, penyelidik membandingkan TTT-Linear dan TTT-MLP dengan Transformer dan Mamba, dua garis dasar.

short Text

From Rajah 11 Kita dapat membuat kesimpulan berikut: konteks

• 2k, prestasi TTT-linear (M), Mamba dan Transformer adalah setanding kerana daripada garisan Kebanyakannya bertindih. TTT-MLP (M) berprestasi lebih teruk sedikit dengan belanjawan FLOP yang lebih besar. Walaupun TTT-MLP mempunyai kebingungan yang lebih baik daripada TTT-Linear pada pelbagai saiz model, kos tambahan FLOP mengimbangi kelebihan ini.
• Untuk konteks 8k, kedua-dua TTT-Linear (M) dan TTT-MLP (M) berprestasi jauh lebih baik daripada Mamba, yang agak berbeza daripada pemerhatian dalam konteks 2k. Malah TTT-MLP (T) menggunakan rangkaian tulang belakang Transformer adalah lebih baik sedikit daripada Mamba pada sekitar 1.3B. Fenomena penting ialah apabila panjang konteks meningkat, kelebihan lapisan TTT berbanding lapisan Mamba juga berkembang.
• Dengan panjang konteks mencecah 8k, Transformer masih berprestasi baik dalam kebingungan di bawah setiap saiz model, tetapi ia tidak lagi berdaya saing kerana kos FLOP.

Keputusan di atas menunjukkan kesan penukaran lapisan TTT daripada rangkaian tulang belakang Mamba kepada rangkaian tulang belakang Transformer. Para penyelidik membuat hipotesis bahawa konvolusi temporal dalam rangkaian tulang belakang Mamba akan lebih membantu apabila keadaan tersembunyi lapisan pemodelan jujukan kurang ekspresif. Model linear kurang ekspresif daripada MLP dan oleh itu mendapat lebih banyak manfaat daripada konvolusi. . Kaedah latihan di sini adalah sama seperti Pile, dan semua eksperimen untuk lapisan TTT dilakukan dalam satu larian latihan. Daripada subset keputusan dalam Rajah 12, mereka membuat pemerhatian berikut:

Dalam konteks Buku 2k, semua pemerhatian untuk Pile 2k masih kekal, kecuali Mamba kini berprestasi lebih baik sedikit daripada TTT-Linear ( Dan baris mereka secara kasar bertindih dalam Longgokan 2k).

Dalam konteks 32k, kedua-dua TTT-Linear (M) dan TTT-MLP (M) berprestasi lebih baik daripada Mamba, serupa dengan pemerhatian untuk Pile 8k. Malah TTT-MLP (T) dengan tulang belakang Transformer berprestasi lebih baik sedikit daripada Mamba dalam konteks 32k.

TTT-MLP (T) hanya lebih teruk sedikit daripada TTT-MLP (M) pada skala 1.3B. Seperti yang dinyatakan di atas, adalah sukar untuk mendapatkan undang-undang penskalaan empirikal kerana kekurangan padanan linear yang jelas. Walau bagaimanapun, aliran kukuh dalam TTT-MLP (T) menunjukkan bahawa tulang belakang Transformer mungkin lebih sesuai untuk model yang lebih besar dan konteks yang lebih panjang, di luar skop penilaian kami.

Clock Time

Latihan dan inferens LLM boleh diuraikan kepada ke hadapan, ke belakang dan generasi. Pemprosesan perkataan isyarat semasa inferens (juga dipanggil pra-populasi) adalah sama seperti operasi ke hadapan semasa latihan, kecuali operasi ke belakang tidak memerlukan penyimpanan nilai pengaktifan perantaraan.

Memandangkan kedua-dua ke hadapan (semasa latihan dan inferens) dan ke belakang boleh diproses secara selari, bentuk dwi digunakan di sini. Menjana token baharu (juga dipanggil penyahkodan) bersifat berurutan, jadi bentuk mentah digunakan di sini.

Pengkaji menyebut bahawa disebabkan oleh keterbatasan sumber, percubaan dalam artikel ini ditulis dalam JAX dan dijalankan pada TPU. Pada pod TPU v5e-256, garis dasar Transformer mengambil masa 0.30 saat setiap lelaran untuk melatih dengan konteks 2k, manakala TTT-Linear mengambil masa 0.27 saat setiap lelaran, iaitu 10% lebih pantas tanpa sebarang pengoptimuman sistem. Memandangkan Mamba (dilaksanakan dengan PyTorch, Triton dan CUDA) hanya boleh dijalankan pada GPU, untuk membuat perbandingan yang adil, penyelidik menjalankan pengoptimuman sistem awal kaedah ini supaya ia boleh berjalan pada GPU.

Sebelah kiri Rajah 15 menunjukkan kependaman kernel hadapan untuk setiap model pada saiz kelompok 16. Semua model ialah 1.3B (Mamba ialah 1.4B). Perlu diingat bahawa garis dasar Transformer di sini adalah lebih pantas daripada yang ada dalam kertas Mamba kerana vLLM digunakan di sini dan bukannya HuggingFace Transformer.

Selain itu, para penyelidik juga menulis satu lagi kernel GPU untuk penjanaan dan menanda aras kelajuannya dengan saiz kelompok 512 di sebelah kanan Rajah 15. Satu lagi metrik masa jam dinding yang biasa digunakan ialah daya pemprosesan, yang mengambil kira potensi manfaat menggunakan saiz kelompok yang lebih besar. Untuk throughput, semua pemerhatian di atas dan susunan antara kaedah masih dipegang.

Pengarang utama

Selepas kajian TTT diserahkan, salah seorang pengarang kertas kerja, Penolong Profesor UCSD Xiaolong Wang, menulis ucapan tahniah di Twitter. Beliau berkata, penyelidikan mengenai TTT berlangsung selama setahun setengah, tetapi sebenarnya sudah lima tahun idea Latihan Masa Ujian (TTT) dilahirkan. Walaupun idea asal dan hasil semasa adalah berbeza sama sekali.

Tiga pengarang utama kertas TTT masing-masing adalah dari Stanford, UC Berkeley dan UCSD.

Antaranya, Yu Sun ialah seorang felo pasca doktoral di Universiti Stanford Beliau lulus dari UC Berkeley EECS dengan Ph.D., dan hala tuju penyelidikan jangka panjangnya ialah TTT.

Xinhao Li ialah calon PhD di UCSD Dia lulus dari Universiti Sains dan Teknologi Elektronik China.

Karan Dalal ialah calon PhD di UC Berkeley yang mengasaskan syarikat permulaan teleperubatan veterinar yang dipanggil Otto semasa di sekolah menengah.

Tiga orang di atas semuanya menulis latihan masa ujian di baris pertama tapak web peribadi mereka yang memperkenalkan arah penyelidikan.

Untuk butiran penyelidikan lanjut, sila rujuk kertas asal.

Atas ialah kandungan terperinci Tukar sepenuhnya model bahasa: TTT seni bina baharu mengatasi Transformer dan model ML menggantikan keadaan tersembunyi RNN. Untuk maklumat lanjut, sila ikut artikel berkaitan lain di laman web China PHP!