Scalare i minibatch di trasporto ottimale online per il flow matching

Abstract degli autori

Una sfida centrale nel portare i modelli di flusso su scala più ampia è mantenere la stabilità durante l’addestramento e l’integrazione. Nel flow matching, questa stabilità dipende in modo cruciale da come i campioni sorgente vengono abbinati ai dati. L’abbinamento casuale indipendente introduce una notevole componente stocastica: per ridurla, si ricorre al trasporto ottimale (OT) su minibatch, che costruisce online accoppiamenti tra sorgente e destinazione (Tong et al., 2024). Tuttavia, quando le applicazioni si estendono a dataset più grandi e di dimensionalità più elevata, i minibatch piccoli possono limitare questi benefici, mentre estendere l’accoppiamento a minibatch molto più grandi a ogni iterazione ha un costo computazionale proibitivo. Inoltre, resta irrisolto sul piano teorico il modo in cui l’accoppiamento OT influisce sull’intero gradiente stocastico e sull’integrazione. In questo lavoro sviluppiamo un’approssimazione scalabile del trasporto ottimale entropico online, che consente accoppiamenti tra sorgente e destinazione molto più ampi preservando i benefici dell’OT. Colmiamo la lacuna teorica sia per l’OT non regolarizzato sia per quello entropico, dimostrando che, rispetto all’abbinamento indipendente, riducono in senso stretto la varianza dell’intero gradiente stocastico e l’errore di integrazione. Descriviamo inoltre come questi miglioramenti si accentuino all’aumentare della dimensione del minibatch usato per l’accoppiamento online. Esperimenti sull’addestramento di flussi per l’interpolazione genica a livello di singola cellula e di flussi per immagini confermano queste previsioni: accoppiamenti online più ampi migliorano sistematicamente la stabilità dell’addestramento e le prestazioni finali nella generazione.