Когда Muon ортогонализует градиент, он неизбежно отклоняется от направления наискорейшего спуска локального слоя. Это означает, что он жертвует некоторой "точностью градиента" — локальные потери могут уменьшаться не так быстро, как при использовании SGD. Однако результирующая матрица весов, как правило, имеет более плоский спектр сингулярных значений (лучшую обусловленность). Это "сохранение представления", и оно окупается позже, потому что последующие слои сталкиваются с гораздо более простой задачей оптимизации.

Эта работа демонстрирует, что упрощенная версия MCA, реализованная на современном GPU-оборудовании, может эффективно обучать широкий спектр архитектур — включая глубокие многослойные перцептроны, чрезвычайно широкие сети и даже Трансформеры — без какого-либо вычисления градиента.

В статье представлено pure pruning как механизм обучения. В этом подходе, вместо того чтобы возмущать веса на небольшую величину, алгоритм пытается установить вес равным нулю.

import torch
import torch.nn.functional as F

D = 1024

# Векторы концепций
a, b, c, d = F.normalize(torch.randn(4, D))

# Исходные суперпозиции
S1 = a + b
S2 = c + d

# Произведение суперпозиций
S = S1 * S2
S = F.normalize(S.unsqueeze(0)).squeeze(0)

# Пары концепций
ac = a * c
ad = a * d
bc = b * c
bd = b * d

z = torch.vstack([ac, ad, bc, bd])
z = F.normalize(z)

# Находим пары концепций в результирующей суперпозиции
r = S @ z.T
print(r)
# [0.52, 0.54, 0.53, 0.52]

import torch
import torch.nn.functional as F

D = 1024

# Векторы концепций
a, b, c, d = F.normalize(torch.randn(4, D))

# Исходная суперпозиция
S = a + b + c + d

# Квадрат суперпозиции
S **= 2
S = F.normalize(S.unsqueeze(0)).squeeze(0)

# Пары концепций
aa = a * a
bb = b * b
cc = c * c
dd = d * d
ab = a * b
ac = a * c
ad = a * d
bc = b * c
bd = b * d
cd = c * d

z = torch.vstack([aa, bb, cc, dd, ab, ac, ad, bc, bd, cd])
z = F.normalize(z)

# Находим пары концепций в результирующей суперпозиции
r = S @ z.T
print(r)
# [0.50, 0.52, 0.50, 0.53, 0.27, 0.25, 0.30, 0.28, 0.26, 0.29]

import torch
import torch.nn.functional as F

D = 1024

# Векторы концепций
a, b, c, d = F.normalize(torch.randn(4, D))

# Суперпозиция из попарных взаимодействий концепций
S = (a + b) * (c + d)
S = F.normalize(S.unsqueeze(0)).squeeze(0)

# Читаем из памяти по запросу a
Sv = S * a
Sv = F.normalize(Sv.unsqueeze(0)).squeeze(0)

# Смотрим что прочиталось
z = torch.vstack([a, b, c, d])

r = Sv @ z.T
print(r)
# [-0.06, -0.03,  0.37,  0.39]