ThinkBit: ricerca delle parole di codice, ricalibrazione delle scale e selezione in base al ragionamento

Abstract degli autori

I modelli linguistici di ragionamento generano migliaia di token per risposta e la velocità di generazione è limitata dalla larghezza di banda della memoria. Memorizzare i pesi con due bit anziché in bfloat16 potrebbe quindi rendere l’inferenza 8$\times$ meno costosa. I metodi a due bit vengono di solito valutati in base alla perplessità e all’accuratezza nelle domande a scelta multipla, due misure che dicono poco sulla capacità di un modello di ragionare ancora su problemi più difficili. Presentiamo ThinkBit, un framework rigorosamente scalare a due bit progettato per preservare il ragionamento, nel quale ogni peso è dato da una parola di codice intera moltiplicata per una scala bfloat16 condivisa da un gruppo di pesi. Per prima cosa usiamo una metrica di sensibilità per assegnare a ogni livello una categoria ternaria, a 2 bit o a 3 bit, così da distribuire in modo ottimale il numero di bit tra i livelli. Segue un algoritmo di quantizzazione scalare basato su Gumbel-softmax che, livello per livello, cerca le parole di codice ottimali blocco per blocco. Poiché questo approccio non coglie le interazioni tra i livelli quantizzati, ricalibriamo le scale, mantenendo fisse le parole di codice, mediante una funzione di perdita basata sulla distillazione della conoscenza rispetto al modello insegnante. Eseguiamo poi una breve fase di addestramento consapevole della quantizzazione (QAT), che permette al modello di imitare la distribuzione del modello di base riassegnando le parole di codice. Infine, un test di ragionamento appaiato sui problemi di addestramento seleziona il checkpoint finale tra quello della fase di ricalibrazione e quello della fase QAT. Su Qwen3-1.7B, a parità di dimensioni, il nostro checkpoint ricalibrato ottiene 46,7 su MATH-500, contro 30,0 per EfficientQAT addestrato sugli stessi dati, e raggiunge 73,0 dopo il QAT sulle tracce di ragionamento. Con 0,52 GB, è più piccolo di ogni versione GGUF a 2 bit del modello pubblicata, eppure è più accurato su MATH-500. Mostriamo che il framework è scalabile fino a 14B: un modello da 3,74 GB ottiene 94,8 su MATH-500, contro 97,5 per il modello bfloat16 da 29,5 GB.