Mallin tislaus: opi toisen mallin tiedosta

Mallin tislaus (knowledge distillation) kouluttaa oppilasmallia opettajamallin tuottamilla signaaleilla. Oppilas voi olla pienempi, edullisempi ajaa tai tietylle laitteelle sopivampi. Sen painot opitaan koulutuksessa; ne eivät ole vain lyhennetty kopio opettajan painoista.

Opettaja voi antaa enemmän tietoa kuin yhden oikean luokan. Sen todennäköisyysjakauma kertoo, mitä vaihtoehtoja se pitää uskottavina. Kielimallissa tämä voi tarkoittaa seuraavan tokenin sanastopisteitä. Muut menetelmät käyttävät tuotettuja vastauksia, välikerrosten esityksiä tai signaalien yhdistelmiä.

Opi koko jakaumasta

Tämä PyTorch-häviöfunktio ottaa oppilaan ja opettajan pisteet muodossa batch × classes sekä kokonaislukuiset tavoiteluokat. Mallien täytyy pisteyttää samat luokat samassa järjestyksessä. Opettajan tulos irrotetaan derivointiketjusta, joten häviö päivittää vain oppilaan puolta.

import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels,
                      temperature=2.0, mix=0.5):
    teacher = F.softmax(teacher_logits.detach() / temperature, dim=-1)
    student = F.log_softmax(student_logits / temperature, dim=-1)
    soft_loss = F.kl_div(student, teacher, reduction="batchmean")
    soft_loss *= temperature ** 2
    hard_loss = F.cross_entropy(student_logits, labels)
    return mix * soft_loss + (1 - mix) * hard_loss

Lämpötila pehmentää kumpaakin jakaumaa. KL-divergenssi vertaa niiden todennäköisyyksiä; kerroin temperature ** 2 kompensoi pehmennyksen vaikutusta gradienttien mittakaavaan. Ristientropia antaa tavallisen palautteen tavoiteluokista. mix määrää tavoitteiden painotuksen.

Tokenitason kielimallissa kelvolliset sijainnit voidaan yhdistää eräulottuvuuteen, kun täyte- ja muut ohitettavat tavoitesijainnit on poistettu. Erilaiset sanastot vaativat lisäksi tavan sovittaa luokat yhteen.

Pienenkin mallin täytyy oppia

Tislaus siirtää valittua toimintaa koulutussyötteiden avulla. Se ei takaa, että opettajan kaikki kyvyt mahtuvat oppilaaseen. Arvioi oppilasta tehtävillä ja laitteilla, joilla sitä käytetään.

Kvantisointi muuttaa lukujen tarkkuutta. Tislaus muuttaa mallin koulutustapaa. Niitä voidaan käyttää yhdessä.