Verbesserung des Routings in spärlichen Expertenmischungen mit Graphen von Tokens

Sparse Mixture of Experts (SMoE) hat sich als Schlüssel zur Erreichung beispielloser Skalierbarkeit im Deep Learning herauskristallisiert. Indem pro Probe nur ein kleiner Teil der Parameter aktiviert wird, ermöglicht SMoE eine exponentielle Erhöhung der Parameteranzahl bei konstantem Rechenaufwand. Allerdings sind SMoE-Modelle anfällig für Routing-Schwankungen – Veränderungen im Routing eines Eingabewerts zu seinem Ziel-Experten – in der späten Phase des Modelltrainings, was zur Nicht-Robustheit des Modells führt. In dieser Arbeit zeigen wir die Einschränkungen von SMoE aus der Perspektive des probabilistischen graphischen Modells (PGM) auf. Durch dieses PGM-Rahmenwerk heben wir die Unabhängigkeit bei der Expertenauswahl der Tokens hervor, welche das Modell für Routing-Schwankungen und Nicht-Robustheit anfällig macht. Um diese Unabhängigkeit zu verringern, schlagen wir das neuartige Similarity-Aware (S)MoE vor, das Interaktionen zwischen Tokens während der Expertenauswahl berücksichtigt. Anschließend leiten wir ein neues PGM her, das einem (S)MoE-Attention-Block zugrunde liegt und über eine einzelne (S)MoE-Schicht hinausgeht. Unter Nutzung der durch die Attention-Matrix erfassten Token-Ähnlichkeiten schlagen wir das innovative Attention-Aware (S)MoE vor, welches die Attention-Matrix verwendet, um das Routing der Tokens zu den geeigneten Experten in (S)MoE zu steuern. Wir beweisen theoretisch, dass similarity-/attention-aware Routing hilft, die Entropie der Expertenauswahl zu verringern, was zu stabileren Token-Routing-Mechanismen führt. Unsere Modelle validieren wir empirisch bei verschiedenen Aufgaben und Domänen und zeigen signifikante Verbesserungen bei der Reduktion von Routing-Schwankungen, der Steigerung der Genauigkeit und der Erhöhung der Modellrobustheit gegenüber dem Baseline MoE-Transformer mit Token-Routing über Softmax-Gating.