Instance Awareness of Multi-class Semantic Segmentation Loss Functions

1. Enunciato del Problema

Il documento affronta la sfida dello sbilanciamento simultaneo di istanze e classi nella segmentazione semantica multiclasse, specificamente nel contesto della segmentazione delle metastasi cerebrali (BraTS-METS 2025).

• Sbilanciamento delle Classi: Le classi rare (ad esempio, Cavità di Resezione - RC) hanno significativamente meno voxel e istanze rispetto alle classi comuni (ad esempio, Tumore Totale - WT). Le funzioni di perdita standard (come Dice + Cross-Entropy) mediando i gradienti su tutti i voxel di primo piano, causano il fatto che le classi rare ricevano segnali di addestramento sproporzionatamente piccoli.
• Sbilanciamento delle Istanze: All'interno di una singola classe, le lesioni piccole (micro-metastasi) sono spesso sottorappresentate rispetto alle lesioni grandi perché le funzioni di perdita standard a livello di voxel sono dominate dal volume maggiore delle grandi istanze.
• Il Divario: Le funzioni di perdita "sensibili alle istanze" esistenti (Blob Loss, CC Loss) sono state progettate per la segmentazione a singola classe. Esse eguagliano i gradienti tra le istanze della stessa classe, ma non riescono a risolvere lo sbilanciamento tra classi diverse in un contesto multiclasse. Al contrario, le funzioni di perdita multiclasse standard gestiscono lo sbilanciamento delle classi tramite una media uniforme, ma ignorano le disparità di dimensioni a livello di istanza.

2. Metodologia

Gli autori propongono un framework che estende le funzioni di perdita sensibili alle istanze alla segmentazione multiclasse e integra una ponderazione inversa alla dimensione in modo innovativo e localizzato.

A. Estensione Multiclasse tramite Decomposizione Uno-vs-Rest

Gli autori adattano Blob Loss e CC Loss (Connected Component) per contesti multiclasse utilizzando una strategia uno-vs-rest:

1. Decomposizione: Per ogni classe di primo piano $c$, la mappa di etichette multiclasse viene isolata in una mappa binaria $y_c$ (dove la classe $c$ è primo piano e tutte le altre sono sfondo).
2. Estrazione delle Istanze: Vengono identificati i componenti connessi 3D all'interno di $y_c$ per definire le singole istanze $\{S^c_1, \dots, S^c_{K_c}\}$.
3. Calcolo della Perdita:
   • CC Loss: Partiziona il volume in celle di Voronoi $V^c_k$ basate sulla distanza euclidea dal componente più vicino. La perdita viene calcolata solo all'interno di queste regioni locali per impedire che i grandi vicini dominino quelli piccoli.
   • Blob Loss: Costruisce un dominio specifico per l'istanza $\Omega^c_k$ rimuovendo tutti gli altri componenti di primo piano della stessa classe, lasciando solo il componente target e lo sfondo.
4. Aggregazione: La perdita finale per istanza ($L_{instance}$) è la media uniforme delle perdite su tutti i componenti e su tutte le classi:
   $$L_{instance} = \frac{1}{|C|} \sum_{c \in C} \frac{1}{K_c} \sum_{k=1}^{K_c} \ell^c_k$$
   Insight Chiave: Mediando uniformemente sulle classi ($\frac{1}{|C|}$), una classe rara con poche istanze (ad esempio RC) contribuisce in modo uguale al gradiente rispetto a una classe comune con molte istanze (ad esempio ET), risolvendo efficacemente lo sbilanciamento delle classi, mentre la media per componente risolve lo sbilanciamento delle istanze.

B. Ponderazione Inversa alla Dimensione Localizzata (IWL)

La ponderazione inversa alla dimensione standard (assegnare pesi inversamente proporzionali alla dimensione del componente) fallisce nei contesti multiclasse se applicata globalmente, perché i pesi variano di ordini di grandezza tra classi rare e comuni, destabilizzando l'addestramento.

• Soluzione Proposta: Integrare la ponderazione all'interno della perdita per componente.
• Meccanismo: Una mappa di pesi $w^c_k$ viene applicata solo ai voxel all'interno del dominio specifico del componente ($\Omega^c_k$ o $V^c_k$).
  $$w^c_k(v) = \min(\max(|\Omega^c_k|/|S^c_k|, 1), 2 \times 10^5)$$
• Effetto: Questo confina la riponderazione al contesto spaziale di una singola istanza, amplificando i gradienti per le istanze piccole senza creare squilibri globali dei pesi tra le classi.

C. Funzione di Perdita Combinata

La perdita totale combina la perdita globale standard (per penalizzare i falsi positivi nelle regioni di sfondo) con la nuova perdita per istanza:
$$L = \alpha \cdot L_{global} + \beta \cdot L_{instance}$$
Gli autori valutano $\alpha=1, \beta=1$ e $\alpha=1, \beta=2$.

3. Contributi Chiave

1. Riproposizione delle Perdite per Istanza: Dimostrato che estendere le funzioni di perdita sensibili alle istanze (Blob/CC) alla segmentazione multiclasse tramite decomposizione uno-vs-rest risolve simultaneamente lo sbilanciamento delle classi (tramite media uniforme delle classi) e lo sbilanciamento delle istanze.
2. Riponderazione Localizzata: Mostrato che la ponderazione inversa alla dimensione è efficace solo se integrata localmente all'interno delle perdite per componente. L'applicazione globale destabilizza l'addestramento, mentre l'applicazione locale migliora significativamente le prestazioni delle classi rare (ad esempio, il Dice per RC è aumentato da 0,22 a 0,44).
3. Analisi del Trade-off: Identificato un trade-off sistematico tra sovrapposizione a livello di voxel (Dice) e rilevamento a livello di istanza (Qualità Panottica). I metodi che aumentano aggressivamente il peso delle classi rare migliorano il Dice ma spesso degradano la qualità del rilevamento (PQ).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti sul dataset BraTS-METS 2025 (260 casi di test), focalizzandosi su quattro regioni: WT, TC, ET e la rara Cavità di Resezione (RC).

• Baseline: nnU-Net standard (Dice + CE) ha raggiunto un Dice per RC di 0,15 e una Media PQ di Primo Piano (a $\tau=0,5$) di 0,38.
• CC Loss Multiclasse:
  • Ha migliorato il Dice per RC a 0,31.
  • Ha mantenuto la Media PQ di Primo Piano a 0,39 (comparabile alla baseline).
  • Ha raggiunto la migliore Media Dice di Primo Piano complessiva (0,64).
• Blob Loss Multiclasse:
  • Ha raggiunto la migliore Media PQ di Primo Piano (0,40) e Qualità di Riconoscimento (0,53).
  • Ha migliorato il Dice per RC a 0,22.
• Con Ponderazione Inversa alla Dimensione Localizzata (IWL):
  • Blob + IWL: Ha raggiunto il Dice per RC più alto (0,44) e la Media Dice di Primo Piano (0,65).
  • Costo: Questo è avvenuto a scapito della qualità del rilevamento, facendo scendere la Media PQ di Primo Piano a 0,31.
• Parametro di Ponderazione ($\beta$): A differenza dei contesti a singola classe dove è raccomandato $\beta=2$, il contesto multiclasse ha funzionato meglio con $\beta=1$. La decomposizione uno-vs-rest fornisce già un bilanciamento delle classi sufficiente, e $\beta=2$ porta a una sovrasegmentazione delle classi rare.

5. Significato e Conclusione

• Impatto Clinico: Una segmentazione accurata della cavità di resezione (RC) è fondamentale per la pianificazione della Radioterapia Stereotassica (SRS) post-operatoria. Migliorare il Dice per RC da 0,15 a 0,44 migliora direttamente l'accuratezza della pianificazione del trattamento.
• Insight Metodologico: Il documento stabilisce che una riponderazione strutturata e con ambito spaziale definito (decomposizione uno-vs-rest + ponderazione inversa alla dimensione locale) è superiore alla riponderazione globale per la segmentazione multiclasse con dati sbilanciati.
• Direzioni Future: Gli autori evidenziano il trade-off intrinseco tra Dice e Qualità Panottica. Il lavoro futuro dovrebbe mirare a risolvere questo trade-off per raggiungere simultaneamente una sovrapposizione di voxel elevata e un rilevamento preciso delle istanze.

In sintesi, questo lavoro fornisce un framework robusto di funzioni di perdita che tratta ogni istanza tumorale e ogni classe tumorale con uguale importanza, colmando il divario tra la segmentazione automatizzata e l'accuratezza per lesione richiesta per il trattamento individualizzato delle metastasi cerebrali.