Sebuah neural network jarang mengakhiri training pada satu-satunya titik yang berguna di parameter space. Menjelang akhir training, stochastic gradient descent (SGD) dapat mengunjungi beberapa konfigurasi parameter berdekatan yang semuanya bekerja cukup baik, sementara checkpoint terakhir hanya merepresentasikan salah satunya.

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

Idenya terdengar sangat sederhana, tetapi detailnya penting. Averaging checkpoint dari training run yang tidak berkaitan dapat gagal, averaging terlalu dini dapat mencampurkan solusi buruk, dan model dengan batch normalization memerlukan penanganan khusus setelah parameternya dirata-ratakan.

Artikel ini membangun mental model praktis untuk SWA: apa yang sebenarnya dirata-ratakan, mengapa learning-rate schedule penting, perbedaannya dengan checkpoint ensembling dan exponential moving average, cara mengevaluasinya secara adil, serta kapan checkpoint akhir biasa merupakan pilihan yang lebih baik.

Mulai dengan contoh terkecil yang berguna

Misalkan model kecil hanya memiliki dua parameter trainable. Menjelang akhir training, 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 sama adalah:

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

SWA memakai [2.2, 3.9333...] sebagai parameter satu model. Pada network nyata, operasi yang sama diterapkan elemen demi elemen pada setiap tensor parameter yang di-average. Anda tidak merata-ratakan prediksi saat inference dan tidak 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 membuat parameter vector baru dari A, B, dan C lalu mengevaluasi satu network tersebut.

Mental model yang berguna adalah trajectory, bukan tiga model independen

Optimisasi neural network bersifat non-convex, sehingga averaging parameter arbitrer tidak aman. Dua network yang dilatih secara independen dapat merepresentasikan fungsi serupa dengan susunan hidden unit berbeda. Rata-rata per-koordinat dari kedua model dapat jatuh pada solusi buruk.

SWA lebih terbatas. Metode aslinya merata-ratakan titik yang diambil sepanjang satu trajectory SGD dengan learning-rate schedule yang menjaga optimizer tetap menjelajahi suatu region parameter space. Karena titik tersebut berasal dari jalur optimisasi yang berkelanjutan, correspondence antarparameter terjaga dengan cara yang tidak dijamin pada run yang tidak terkait.

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

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

Tidak perlu menyimpan setiap checkpoint. Gunakan running average:

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

Setelah setiap update, nilai tersebut sama dengan rata-rata aritmetika semua parameter yang sudah dikumpulkan. Kebutuhan memori kira-kira satu salinan parameter seukuran model, bukan arsip checkpoint yang terus bertambah.

Mengapa SWA mengubah bagian akhir training

Schedule konvensional sering menurunkan learning rate menuju nilai kecil agar optimisasi menetap dekat satu titik. SWA membutuhkan perilaku late-training yang berbeda karena averaging beberapa checkpoint yang hampir identik memberi sedikit manfaat.

Prosedur SWA yang diajukan Izmailov dan kolega menggunakan learning rate konstan atau cyclical selama fase averaging. Tujuannya agar SGD terus bergerak melalui region yang berguna, bukan langsung runtuh ke satu titik akhir. Parameter sepanjang trajectory tersebut kemudian dirata-ratakan.

Rencana training sederhana:

1. Train secara normal untuk sebagian besar budget.
2. Masuk ke fase SWA dengan learning rate non-vanishing yang sesuai.
3. Secara berkala tambahkan parameter saat ini ke running average.
4. Selesaikan training dan siapkan model hasil averaging untuk evaluasi.
5. Hitung ulang statistik batch normalization bila diperlukan.
6. Bandingkan model SWA dengan checkpoint normal pada held-out data.

Waktu mulai, learning rate, dan frekuensi pengumpulan adalah hyperparameter, bukan konstanta universal.

Mengapa averaging dapat membantu

Bayangkan low-loss region di sekitar solusi berbentuk lembah lebar. SGD dapat bergerak di dalam lembah dan mengunjungi beberapa titik dengan training loss yang baik. Rata-rata aritmetika titik-titik tersebut dapat berada lebih sentral daripada satu titik individual di sisi lembah.

Ini adalah intuisi geometris SWA. Pekerjaan SWA awal melaporkan bahwa averaging iterate SGD dengan learning-rate schedule yang dimodifikasi dapat menemukan solusi yang terkait dengan optima lebih lebar dan meningkatkan generalisasi pada setting image classification yang dievaluasi.

Intuisi tersebut bukan aturan universal bahwa parameter yang tampak lebih flat pasti memiliki generalisasi lebih baik. Loss geometry neural network bergantung pada parameterisasi, dan manfaat SWA tetap harus diukur pada tugas yang relevan.

Klaim praktisnya lebih sempit: late-training weight averaging adalah eksperimen berkompleksitas rendah yang dapat menghasilkan satu model berguna ketika trajectory optimizer menyediakan sampel parameter yang kompatibel dan berkualitas baik.

Jangan samakan SWA dengan checkpoint ensembling

Ensemble menggabungkan output:

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, keduanya umumnya tidak sama:

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

Ensemble dapat memperoleh robustness dari beberapa model yang dievaluasi secara independen, tetapi inference cost biasanya meningkat karena membutuhkan beberapa forward pass. SWA mempertahankan struktur inference satu model. Biaya tambahannya terutama selama training dan persiapan model untuk memelihara salinan parameter rata-rata.

SWA juga berbeda dari exponential moving average

Exponential moving average (EMA) juga memelihara set parameter kedua, tetapi checkpoint terbaru mendapat bobot lebih besar. Update umum:

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

Dengan decay seperti 0.999, informasi lama memudar secara bertahap. SWA sebaliknya memberi bobot sama kepada sampel parameter yang masuk dalam rata-rata aritmetika.

Tujuan training juga dapat berbeda. SWA lazim dipasangkan dengan late learning-rate schedule untuk terus mengambil sampel dari suatu region solusi, sedangkan EMA sering diperbarui sepanjang training sebagai versi yang dihaluskan dari model yang terus berubah.

Tidak ada aturan yang secara inheren lebih unggul. Jika recipe training sudah memiliki EMA yang tervalidasi, bandingkan keduanya dengan data split, training budget, dan protokol evaluasi yang sama.

Batch normalization memerlukan perhatian eksplisit

Averaging parameter trainable bukan seluruh cerita untuk network yang menggunakan batch normalization.

Layer batch normalization menyimpan running activation statistics seperti mean dan variance. Statistik tersebut diamati ketika model training biasa bergerak sepanjang trajectory parameternya. Setelah SWA membuat set bobot rata-rata baru, statistik lama mungkin tidak lagi sesuai dengan activation yang dihasilkan model tersebut.

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

Detail implementasi ini berdampak pada kualitas model. Melupakannya dapat membuat parameter average yang benar tampak buruk hanya karena statistik normalization sudah stale.

Arsitektur tanpa batch normalization tidak membutuhkan kalkulasi ulang khusus ini. Buffer stateful lain tetap perlu diperiksa berdasarkan semantik framework masing-masing.

Pola implementasi praktis

Pseudo-code berikut menunjukkan transisi state penting tanpa terikat pada framework tertentu:

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 production code, gunakan averaging utility resmi framework jika sesuai. Dokumentasi PyTorch saat ini, misalnya, menyediakan AveragedModel untuk memelihara parameter rata-rata dan SWALR untuk 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 cara average dipelihara.

Evaluasi perbandingan yang tepat

Eksperimen SWA hanya berguna jika baseline-nya adil. Minimal, simpan model normal dan model SWA dari training run yang sama atau run yang dicocokkan dengan hati-hati.

Bandingkan metrik yang sesuai dengan deployment, seperti validation loss, accuracy, F1, ranking quality, atau ukuran domain-specific lain. Ukur pula properti operasional yang relevan.

SWA biasanya mempertahankan arsitektur dan jumlah parameter yang sama, sehingga bukan model compression dan tidak secara inheren membuat forward pass lebih murah. Jika latency berubah, periksa runtime path aktual.

Training cost dapat meningkat karena fase SWA memperpanjang training atau mempertahankan learning rate yang terus menjelajah. Memori juga bertambah selama training karena salinan parameter rata-rata disimpan. Biaya ini mungkin kecil dibanding seluruh training job, tetapi tidak nol.

Kesalahan umum

Merata-ratakan checkpoint yang tidak terkait

Jangan berasumsi model independen dapat dirata-ratakan hanya karena arsitekturnya sama. SWA bergantung pada titik kompatibel dari satu trajectory optimisasi, bukan file model arbitrer.

Memulai sebelum model mencapai region yang berguna

Jika parameter awal yang 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 memiliki sedikit trajectory untuk dirata-ratakan. Strategi learning rate adalah bagian metode.

Melupakan statistik normalization

Untuk model dengan batch normalization, running statistics yang stale dapat merusak perbandingan. Hitung ulang atau tangani dengan prosedur yang didukung framework.

Hanya melaporkan model hasil averaging

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

Kapan SWA layak dicoba

SWA layak diuji ketika Anda sudah memiliki pipeline training bergaya SGD yang stabil, mampu membayar fase averaging di akhir, dan ingin mengetahui apakah satu solusi rata-rata dapat menggeneralisasi lebih baik daripada titik akhir trajectory. Metode ini terutama menarik ketika deployment tetap harus memakai satu model, bukan multi-checkpoint ensemble.

SWA kurang menarik ketika training sudah sangat mahal dan tidak dapat diperpanjang, recipe optimizer tidak menyediakan sampel kompatibel yang berguna, atau EMA maupun metode checkpoint-selection lain yang tervalidasi sudah memenuhi kebutuhan. SWA juga bukan pengganti perbaikan data leakage, label buruk, objective yang tidak cocok, atau desain evaluasi lemah.

Kesimpulan

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

Aritmetikanya sederhana; konteks training-lah yang membuatnya bermakna. Gunakan fase averaging yang disengaja, pertahankan baseline checkpoint normal, refresh statistik batch normalization bila perlu, dan evaluasi model tunggal hasilnya pada metrik yang benar-benar penting. Dengan kondisi tersebut, SWA menjadi eksperimen praktis, bukan sekadar trik checkpoint.

Kecerdasan Buatan Neural Network Training Optimisasi