DeepSeek mHC: Spiegazione del paper e codice pytorch

In mHC: Manifold-Constrained Hyper-Connections, DeepSeek perfeziona le Hyper-Connections (HC). A differenza delle classiche connessioni residue (che sommano semplicemente l'input all'output), le HC gestiscono il segnale attraverso svariati flussi latenti che si "mescolano" dinamicamente tra loro tramite matrici di mixing, aumentando drasticamente la capacità espressiva della rete.

Tuttavia, questa flessibilità rompe la stabilità garantita dalla "identity mapping" classica, causando spesso l'esplosione o la sparizione dei gradienti in profondità. mHC risolve questo problema vincolando le matrici di connessione a risiedere sul politopo di Birkhoff.

Il politopo di Birkhoff è uno spazio geometrico composto da tutte le matrici "doppiamente stocastiche", ovvero matrici quadrate con valori non negativi dove la somma di ogni riga e di ogni colonna è esattamente pari a 1. Proiettando i pesi su questo spazio tramite l'algoritmo Sinkhorn-Knopp, mHC garantisce che l'energia del segnale rimanga costante (senza amplificarsi o ridursi) attraverso i layer, ripristinando la stabilità matematica necessaria per scalare modelli massivi con un costo computazionale trascurabile.

• Code: Unofficial pytorch implementation based on lucidrain's HC
• Paper: mHC: Manifold-Constrained Hyper-Connections
• Dataset: /

mHC: Obiettivi e domande di ricerca

• Perché le Hyper-Connections standard (HC) diventano instabili su larga scala? Le HC introducono matrici di mixing non vincolate tra i vari flussi paralleli del residuo. Questo rompe la proprietà di "identity mapping" (mappatura identità) fondamentale per le reti profonde (come le ResNet), portando i gradienti ad esplodere o svanire quando la rete diventa molto profonda, impedendo la convergenza o degradando le performance.

• Come garantisce mHC la stabilità del segnale durante il training? mHC costringe la matrice di connessione residua a risiedere su uno spazio geometrico specifico: quella delle matrici doppiamente stocastiche (dove la somma di ogni riga e di ogni colonna è 1). Questo assicura che la norma del segnale venga preservata (non-expansive mapping) attraverso i layer, ripristinando la stabilità simile a quella dell'identity mapping.

• Qual è il costo computazionale di questa stabilità aggiuntiva? L'introduzione dell'algoritmo Sinkhorn-Knopp e delle proiezioni aggiunge un overhead teorico, ma grazie a ottimizzazioni infrastrutturali (kernel fusion e sovrapposizione della comunicazione), l'impatto reale è contenuto (circa il 6.7% di FLOPs in più rispetto a un modello base), risultando in un trade-off favorevole vista la migliore performance. Se la tecnica risulta vincente, vedremo implementazioni ancora più ottimizzate nei mesi a venire.

mHC: Cosa ne pensa Mauro?

• Novelty: 4/5. L'approccio è decisamente innovativo. Non si limita a proporre una nuova architettura, ma applica rigorosi principi di geometria (teoria delle varietà) al design delle connessioni neurali. L'idea di vincolare le matrici di mixing delle hyper-connections affinché siano doppiamente stocastiche (doubly stochastic) per preservare la norma del segnale è un contributo teorico significativo.
• Clarity: 3/5. Il paper è denso e richiede una solida base matematica (algebra lineare e ottimizzazione) per essere compreso appieno, specialmente nelle sezioni riguardanti il politopo di Birkhoff e l'algoritmo Sinkhorn-Knopp. Tuttavia, la narrazione del problema (rottura dell'identity mapping nelle HC) è chiara.
• Reproducibility: 2/5. Al momento della stesura, il codice di training specifico non sembra essere stato rilasciato pubblicamente in un repository dedicato, sebbene DeepSeek sia nota per l'open source. L'implementazione efficiente richiede kernel personalizzati (scritti in TileLang) per la fusione delle operazioni, il che rende la riproduzione da zero complessa per un team indipendente. Tuttavia, iniziano a comparire le prime riproduzioni non ufficiali.
• Affordability: 1/5. Bassa per la riproduzione completa. Gli esperimenti sono condotti su modelli MoE (Mixture of Experts) fino a 27 miliardi di parametri (27B). Sebbene l'overhead computazionale introdotto da mHC sia ridotto (~6.7%), il pre-training di modelli di questa scala richiede cluster di GPU di classe enterprise (es. H100/A100) e budget elevati.

mHC: Metodi e Tecniche

Manifold Constraint & Birkhoff Polytope

Il cuore di mHC è la proiezione della matrice di connessione residua $H_{res}$ sul politopo di Birkhoff, l'insieme delle matrici doppiamente stocastiche.

In una rete con Hyper-Connections, il segnale si propaga attraverso $n$ "corsie" parallele. La formula di aggiornamento per un layer è:

$X_{l+1} = H_{res}^l X_l + \text{FeedForward}(X_l)$

Se $H_{res}$ è una matrice arbitraria, dopo molti layer il segnale $X$ può diventare instabile. mHC impone che $H_{res}$ soddisfi:

1. Tutti gli elementi sono non-negativi ($H_{ij} \ge 0$).
2. La somma di ogni riga è 1 ($\sum_j H_{ij} = 1$).
3. La somma di ogni colonna è 1 ($\sum_i H_{ij} = 1$).

Sinkhorn-Knopp Algorithm

Per ottenere una matrice che soddisfi i vincoli sopra partendo da pesi apprendibili, il paper utilizza l'algoritmo iterativo di Sinkhorn-Knopp. Questo algoritmo normalizza in maniera alternata righe e colonne fino alla convergenza.

Ecco una rappresentazione in pseudo-codice PyTorch di come potrebbe essere implementato il blocco di proiezione:

import torch
import torch.nn as nn

class ManifoldResLayer(nn.Module):
    def __init__(self, num_streams, iterations=3):
        super().__init__()
        self.iterations = iterations
        # Pesi grezzi inizializzati (log-space)
        self.raw_weights = nn.Parameter(torch.randn(num_streams, num_streams))

    def get_constrained_matrix(self):
        # 1. Rendere positivi i pesi (exp)
        W = torch.exp(self.raw_weights)
        
        # 2. Sinkhorn-Knopp Iterations
        for _ in range(self.iterations):
            # Normalizza colonne
            W = W / (W.sum(dim=0, keepdim=True) + 1e-6)
            # Normalizza righe
            W = W / (W.sum(dim=1, keepdim=True) + 1e-6)
            
        return W

    def forward(self, x):
        # x shape: [Batch, Length, Num_Streams, Dim]
        H_res = self.get_constrained_matrix()
        
        # Applica il mixing tra gli stream
        # Einsum: mixa la dimensione Num_Streams (s) con i pesi H
        out = torch.einsum('bsd, st -> btd', x, H_res) 
        return out

Infrastructure Optimization

Per mitigare il costo di calcolo di queste operazioni aggiuntive, DeepSeek utilizza diverse tecniche:

• Kernel Fusion: Implementazione di kernel custom (CUDA o Triton/TileLang) che fondono l'operazione di Sinkhorn con le altre operazioni del layer per ridurre l'accesso alla memoria (memory-bound operations).
• DualPipe: Ottimizzazione della pipeline parallela per sovrapporre (overlap) il calcolo delle proiezioni con le comunicazioni di rete tra GPU.

mHC: Dataset

I modelli (fino a 27B parametri) sono stati addestrati su corpus di testo proprietari o standard per LLM, tokenizzati con un vocabolario di 128k token.

mHC: Licenze

Il paper non specifica una licenza per il codice in quanto non è stato fornito un link diretto al repository nel PDF. Tuttavia, i rilasci precedenti di DeepSeek (es. DeepSeek-V3) sono stati spesso sotto licenze permissive (MIT o Apache 2.0) o licenze custom che permettono l'uso commerciale, ma per mHC è necessario attendere il rilascio ufficiale degli artefatti.