Sebuah neural network jarang mengakhiri pelatihan tepat di satu-satunya titik berguna dalam ruang parameter. Menjelang akhir pelatihan, stochastic gradient descent (SGD) dapat mengunjungi beberapa konfigurasi parameter yang berdekatan dan semuanya berkinerja cukup baik, sementara checkpoint terakhir hanya mewakili salah satunya.

Stochastic weight averaging (SWA) mengubah pengamatan tersebut menjadi teknik pelatihan sederhana: kumpulkan parameter model dari beberapa titik pada bagian akhir trajektori SGD lalu hitung rata-rata aritmetikanya. Hasilnya adalah satu model dengan bobot yang dirata-ratakan, sehingga inference tidak perlu menjalankan ensemble dari semua checkpoint yang dikumpulkan.

Idenya terdengar sangat sederhana, tetapi detailnya penting. Merata-ratakan checkpoint dari training run yang tidak berhubungan dapat gagal, merata-ratakan terlalu awal dapat mencampurkan solusi buruk ke hasil, dan model dengan batch normalization memerlukan penanganan khusus setelah parameternya dirata-ratakan.

Artikel ini membangun model mental praktis untuk SWA. Anda akan mempelajari apa yang sebenarnya dirata-ratakan, mengapa learning-rate schedule penting, bagaimana SWA berbeda dari checkpoint ensembling dan exponential moving average, bagaimana mengevaluasinya secara adil, serta kapan checkpoint akhir yang lebih sederhana menjadi pilihan yang lebih baik.

Mulai dari contoh terkecil yang berguna

Misalkan sebuah model kecil hanya memiliki dua parameter yang dapat dilatih. Menjelang akhir pelatihan, SGD menghasilkan checkpoint berikut:

checkpoint A: [2.0, 4.0]
checkpoint B: [2.4, 3.6]
checkpoint C: [2.2, 4.2]

Rata-rata dengan bobot setara adalah:

w_swa = ([2.0, 4.0] + [2.4, 3.6] + [2.2, 4.2]) / 3
      = [2.2, 3.9333...]

SWA menggunakan [2.2, 3.9333...] sebagai parameter satu model.

Pada network nyata, operasi yang sama diterapkan elemen demi elemen pada setiap tensor parameter yang dirata-ratakan. Anda tidak merata-ratakan prediksi saat inference dan tidak perlu mempertahankan tiga model aktif hanya untuk menghasilkan satu prediksi.

Perbedaan ini penting. Ensemble dapat mengevaluasi A, B, dan C secara terpisah lalu menggabungkan prediksinya. SWA justru membuat vektor parameter baru dari A, B, dan C lalu mengevaluasi satu network tersebut.

Model mental yang berguna adalah trajektori, bukan tiga model independen

Optimisasi neural network bersifat non-convex, sehingga merata-ratakan parameter secara sembarang tidak aman. Dua network yang dilatih secara independen dapat menerapkan fungsi serupa sambil mengatur hidden unit secara berbeda. Rata-rata per-koordinat di antara model tersebut dapat jatuh pada solusi yang buruk.

SWA lebih terbatas. Metode aslinya merata-ratakan titik yang diambil sepanjang satu trajektori SGD sambil menggunakan learning-rate schedule yang membuat optimizer terus menjelajahi suatu wilayah ruang parameter. Karena titik-titik itu berasal dari jalur optimisasi berkelanjutan yang sama, korespondensi parameter terjaga dengan cara yang tidak dijamin pada run yang tidak berhubungan.

Jika parameter yang dikumpulkan adalah w_1, w_2, ..., w_n, parameter SWA adalah:

w_swa = (1 / n) * sum(w_i for i = 1..n)

Anda tidak perlu menyimpan setiap checkpoint untuk menghitungnya. Pertahankan running average:

average_1 = w_1
average_n = average_(n-1) + (w_n - average_(n-1)) / n

Setelah setiap pembaruan, nilai running tersebut sama dengan rata-rata aritmetika semua parameter yang telah dikumpulkan. Karena itu, kebutuhan memori kira-kira hanya satu salinan parameter sebesar model tambahan, bukan arsip checkpoint yang terus bertambah.

Mengapa SWA mengubah bagian akhir pelatihan biasa

Schedule konvensional sering menurunkan learning rate menuju nilai kecil agar optimisasi menetap di sekitar satu titik. SWA membutuhkan perilaku akhir pelatihan yang agak berbeda karena merata-ratakan beberapa checkpoint yang hampir identik hanya memberi sedikit manfaat.

Prosedur SWA yang diajukan Izmailov dan rekan-rekannya menggunakan learning rate konstan atau siklik selama fase averaging. Tujuannya agar SGD terus bergerak melalui wilayah yang berguna alih-alih langsung mengerucut ke satu titik akhir. Parameter yang diambil sepanjang trajektori tersebut kemudian dirata-ratakan.

Rencana pelatihan sederhana terlihat seperti berikut:

1. Train normally for most of the budget.
2. Enter an SWA phase with an appropriate non-vanishing learning rate.
3. Periodically add the current parameters to the running average.
4. Finish training and prepare the averaged model for evaluation.
5. Recompute batch-normalization statistics when required.
6. Compare the SWA model with the normal checkpoint on held-out data.

Waktu mulai, learning rate, dan frekuensi pengumpulan adalah hyperparameter, bukan konstanta universal. Schedule yang bekerja untuk satu arsitektur dan dataset tidak otomatis menjadi jaminan untuk yang lain.

Mengapa averaging dapat membantu

Bayangkan wilayah low-loss di sekitar solusi berbentuk lembah lebar. SGD dapat bergerak di dalam lembah tersebut dan mengunjungi beberapa titik dengan training loss yang baik. Rata-rata aritmetika titik-titik itu dapat berada lebih dekat ke bagian tengah wilayah dibandingkan satu titik di dekat sisinya.

Inilah intuisi geometris di balik SWA. Riset SWA awal melaporkan bahwa merata-ratakan iterasi SGD dengan learning-rate schedule yang dimodifikasi dapat menemukan solusi yang terkait dengan optimum yang lebih lebar dan meningkatkan generalisasi pada konfigurasi klasifikasi gambar yang dievaluasi.

Intuisi tersebut tidak boleh diubah menjadi aturan universal bahwa parameter yang tampak lebih datar pasti melakukan generalisasi lebih baik. Geometri loss neural network bergantung pada parameterisasi, dan manfaat empiris SWA tetap harus diukur pada tugas yang benar-benar penting.

Klaim praktisnya lebih sempit: weight averaging pada akhir pelatihan adalah eksperimen berkompleksitas rendah yang dapat menghasilkan satu model berguna ketika trajektori optimizer menyediakan sampel parameter yang kompatibel dan berkualitas baik.

Jangan samakan SWA dengan checkpoint ensembling

Misalkan tiga checkpoint masing-masing menghasilkan vektor probabilitas untuk sebuah input. Prediction ensemble menggabungkan vektor output tersebut:

p_ensemble = (p_A + p_B + p_C) / 3

SWA merata-ratakan parameter terlebih dahulu:

w_swa = (w_A + w_B + w_C) / 3
p_swa = model(x; w_swa)

Pada neural network nonlinear, kedua ekspresi ini secara umum tidak sama:

model(x; average(weights)) != average(model(x; weights))

Ensemble dapat memperoleh robustness dari beberapa model yang dievaluasi secara independen, tetapi biaya inference biasanya meningkat karena memerlukan beberapa forward pass. SWA mempertahankan struktur inference satu model. Biaya tambahannya terutama muncul selama pelatihan dan persiapan model ketika salinan parameter untuk running average dipertahankan.

Jika latency inference memungkinkan ensemble dan kualitas prediktif maksimum adalah tujuan utama, bandingkan metode-metode tersebut secara empiris alih-alih menganggap SWA menggantikan ensembling.

SWA juga berbeda dari exponential moving average

Exponential moving average (EMA) juga mempertahankan set parameter model kedua, tetapi checkpoint yang lebih baru mendapat bobot lebih besar. Pembaruan umum berbentuk:

w_ema <- decay * w_ema + (1 - decay) * w_current

Dengan decay seperti 0.999, informasi lama memudar secara bertahap ketika parameter baru masuk. Sebaliknya, SWA memberikan bobot setara pada sampel parameter yang dimasukkan ke rata-rata aritmetikanya.

Tujuan pelatihannya juga dapat berbeda. SWA umumnya dipasangkan dengan learning-rate schedule akhir yang dirancang agar terus mengambil sampel dari suatu wilayah solusi. EMA sering diperbarui sepanjang pelatihan sebagai versi yang diperhalus dari model yang terus berkembang.

Tidak ada aturan yang secara inheren lebih unggul. Jika recipe pelatihan Anda sudah memiliki EMA yang teruji, menggantinya dengan SWA menambah variabel eksperimen baru. Bandingkan keduanya dengan data split, training budget, dan protokol evaluasi yang sama.

Batch normalization memerlukan perhatian khusus

Merata-ratakan parameter yang dapat dilatih bukan seluruh persoalan pada network yang menggunakan batch normalization.

Layer batch normalization mempertahankan statistik aktivasi berjalan seperti mean dan variance. Statistik tersebut diamati ketika model training biasa bergerak sepanjang trajektori parameternya. Setelah SWA membuat set bobot rata-rata baru, statistik berjalan lama mungkin tidak lagi cocok dengan aktivasi yang dihasilkan model hasil averaging.

Solusi umum adalah menjalankan satu pass pada data training menggunakan model SWA final untuk menghitung ulang statistik batch normalization sebelum evaluasi. PyTorch, misalnya, menyediakan torch.optim.swa_utils.update_bn() untuk tujuan ini.

Ini adalah detail implementasi yang berdampak pada kualitas model. Mengabaikannya dapat membuat parameter average yang benar tampak lebih buruk hanya karena statistik normalisasinya sudah stale.

Arsitektur tanpa batch normalization tidak memerlukan perhitungan ulang khusus ini. Layer stateful lain atau buffer khusus framework tetap perlu diperiksa berdasarkan semantik yang didokumentasikan, bukan diasumsikan berperilaku seperti parameter trainable biasa.

Pola implementasi praktis

Pseudo-code berikut menunjukkan transisi state penting tanpa mengikat konsep ke satu framework:

model = initialize_model()
average = empty_average()
count = 0

for epoch in training_epochs:
    train_one_epoch(model)

    if epoch >= swa_start:
        set_or_step_swa_learning_rate()
        average = running_parameter_average(average, model, count)
        count += 1
    else:
        step_normal_learning_rate_schedule()

swa_model = copy_model_with_parameters(average)
recompute_normalization_statistics_if_needed(swa_model)
evaluate(swa_model)

Dalam kode produksi, utamakan utility averaging yang disediakan dan didokumentasikan framework jika sesuai dengan setup pelatihan Anda. Dokumentasi PyTorch saat ini, misalnya, menyediakan AveragedModel untuk mempertahankan parameter hasil averaging dan SWALR untuk melakukan annealing menuju SWA learning rate.

Pseudo-code ini sengaja disederhanakan. Distributed training, mixed precision, optimizer state, checkpoint resumption, parameter freezing, dan custom model buffer dapat memengaruhi di mana serta bagaimana average harus dipertahankan.

Evaluasi perbandingan yang tepat

Eksperimen SWA hanya berguna jika baseline-nya adil. Minimal, catat model biasa dan model SWA yang dihasilkan oleh training run yang sama atau oleh run yang dicocokkan dengan cermat.

Bandingkan metrik yang sesuai dengan tugas deployment, seperti validation loss, accuracy, F1, ranking quality, atau ukuran khusus domain lainnya. Ukur juga properti operasional yang penting bagi sistem.

SWA biasanya mempertahankan arsitektur dan jumlah parameter yang sama dengan model dasar, sehingga tidak boleh dianggap sebagai model compression. SWA juga tidak secara inheren membuat forward pass lebih murah. Jika latency berubah, periksa runtime path yang sebenarnya alih-alih mengaitkan perbedaan itu dengan averaging.

Biaya pelatihan dapat meningkat karena fase SWA mungkin memperpanjang training atau mempertahankan learning rate yang terus melakukan eksplorasi. Memori juga meningkat selama training karena salinan parameter hasil averaging disimpan. Biaya ini mungkin kecil dibandingkan keseluruhan training job, tetapi tetap tidak nol.

Kesalahan umum

Merata-ratakan checkpoint yang tidak berhubungan

Jangan menganggap model yang dilatih secara independen dapat dirata-ratakan hanya karena memiliki arsitektur yang sama. Neural network dapat merepresentasikan fungsi serupa dengan susunan parameter berbeda. SWA didasarkan pada averaging titik kompatibel dari satu trajektori optimisasi, bukan file model sembarang.

Memulai sebelum model mencapai wilayah yang berguna

Jika parameter awal yang berkualitas buruk masuk ke equal-weight average, pengaruhnya tidak memudar seperti pada EMA. Pilih fase averaging dengan sengaja dan validasi titik mulainya.

Membiarkan learning rate menghilang selama pengumpulan

Jika checkpoint yang dikumpulkan hampir identik, SWA hanya memiliki sedikit trajektori untuk dirata-ratakan. Strategi learning rate merupakan bagian dari metode, bukan detail tuning yang tidak terkait.

Melupakan statistik normalisasi

Untuk model dengan batch normalization, statistik berjalan yang stale dapat merusak validitas perbandingan. Hitung ulang atau tangani statistik tersebut menggunakan prosedur yang didukung framework.

Hanya melaporkan model hasil averaging

Pertahankan checkpoint biasa sebagai baseline. Jika SWA tidak meningkatkan metrik held-out cukup banyak untuk membenarkan kompleksitas training tambahan, checkpoint normal adalah artefak deployment yang lebih sederhana.

Kapan SWA layak dicoba

SWA layak diuji ketika Anda sudah memiliki pipeline pelatihan bergaya SGD yang stabil, mampu membayar fase averaging di akhir, dan ingin mengeksplorasi apakah satu solusi hasil averaging melakukan generalisasi lebih baik daripada titik akhir trajektori. Teknik ini sangat praktis ketika deployment tetap harus menggunakan satu model, bukan ensemble beberapa checkpoint.

SWA kurang menarik ketika training sudah sangat mahal dan tidak dapat diperpanjang, ketika recipe optimizer tidak menyediakan sampel kompatibel yang berguna untuk dirata-ratakan, atau ketika EMA tervalidasi maupun metode pemilihan checkpoint lain sudah memenuhi kebutuhan. SWA juga bukan pengganti untuk memperbaiki data leakage, label buruk, objektif yang tidak sesuai, atau desain evaluasi yang lemah.

Kesimpulan

Stochastic weight averaging paling tepat dipahami sebagai perubahan pada bagian akhir optimisasi: biarkan SGD terus menjelajahi wilayah yang berguna, kumpulkan state parameter yang kompatibel, lalu ganti titik akhir dengan rata-rata berbobot setara dari state tersebut.

Aritmetikanya sederhana; konteks pelatihannya yang membuat teknik ini bermakna. Gunakan fase averaging yang disengaja, pertahankan baseline checkpoint normal, perbarui statistik batch normalization bila diperlukan, dan evaluasi satu model hasilnya pada metrik yang penting. Jika kondisi tersebut ditangani dengan cermat, SWA menjadi eksperimen praktis alih-alih sekadar trik checkpoint lain.

Kecerdasan Buatan Neural Network Pelatihan Optimisasi