Espacios de estados y mezcla de expertos
S4, Mamba, HiPPO y enrutamiento disperso (MoE)
Dos respuestas al costo de escalar: recurrencias lineales que recuerdan de forma óptima con costo lineal, y redes que solo activan una fracción de sí mismas para cada entrada.
Intuición
Un modelo de espacio de estados (SSM) viene de la teoría de control: un estado evoluciona linealmente según y se lee como . Si la dinámica es lineal, la misma capa puede verse de dos formas: como una recurrencia (barata al generar) o como una convolución gigante (paralelizable al entrenar).
La pregunta clave es qué matriz usar. HiPPO (Gu et al., 2020) la deriva de un principio: el estado debe ser la mejor aproximación polinomial (Legendre) de todo el pasado de la señal. El laboratorio implementa HiPPO-LegS y reconstruye la historia de una señal a partir de solo números.
Una mezcla de expertos (MoE) divide el trabajo: un enrutador decide qué expertos (subredes) procesan cada token, y solo esos se activan. Así, un modelo puede tener billones de parámetros pero usar una fracción por token.
Mecanismo
S4 (Gu, Goel y Ré, 2022) discretiza el SSM con paso ( o bilineal) y calcula el núcleo convolucional de forma eficiente. Mamba (Gu y Dao, 2023) hace que , y dependan de la entrada —selectividad—, lo que rompe la convolución pero permite decidir qué recordar; lo compensa con un scan paralelo consciente del hardware.
MoE clásico (Jacobs, Jordan, Nowlan y Hinton, 1991): con compuerta softmax. MoE disperso (Shazeer et al., 2017; Switch Transformer, 2021) activa solo los top- expertos y añade una pérdida de balance de carga para que ninguno se quede ocioso ni saturado.
En el laboratorio de MoE, cuatro expertos lineales y un enrutador se entrenan en vivo para aproximar una función por tramos; verás cómo el enrutador aprende a especializar a cada experto en una región.
Laboratorio
Memoria HiPPO: el pasado en N números
Dibuja o elige una señal. Un estado de N dimensiones con la matriz HiPPO-LegS la recuerda en línea; la curva roja es la reconstrucción del pasado completo desde el estado.
Mezcla de expertos con enrutador aprendido
Cuatro expertos lineales y una compuerta softmax se entrenan juntos. El color muestra qué experto gana en cada región; activa top-1 para enrutamiento disperso.
Scan selectivo estilo Mamba
Una secuencia entra a un SSM selectivo: las compuertas deciden qué recordar y qué olvidar en cada paso. Compara con un SSM lineal (HiPPO) sobre la misma señal.
Historia
- 1960
Kalman introduce el filtrado en espacios de estados.
Kalman (1960), J. Basic Engineering 82
- 1991
Jacobs, Jordan, Nowlan y Hinton: mezclas adaptativas de expertos locales.
Jacobs et al. (1991), Neural Computation 3(1)
- 2017
Capas MoE con compuertas dispersas y miles de expertos.
Shazeer et al. (2017), ICLR
- 2020
HiPPO: memoria recurrente como proyección óptima sobre polinomios.
Gu, Dao, Ermon, Rudra & Ré (2020), NeurIPS
- 2021
Switch Transformer escala a un billón de parámetros con un experto por token.
Fedus, Zoph & Shazeer (2022), JMLR 23
- 2022
S4 resuelve Long Range Arena, incluida la tarea Path-X de 16 mil pasos.
Gu, Goel & Ré (2022), ICLR
- 2023
Mamba: SSM selectivo con rendimiento de Transformer y costo lineal.
Gu & Dao (2023), arXiv:2312.00752
- 2024
Mixtral 8×7B y Mamba-2 (dualidad espacio de estados–atención).
Jiang et al. (2024), arXiv:2401.04088; Dao & Gu (2024), ICML
Límites
Los SSM comprimen el pasado en un estado fijo: les cuesta copiar o recuperar con exactitud información arbitraria del contexto (Jelassi et al., 2024).
MoE: entrenamiento inestable, balance de carga difícil y alto consumo de memoria (todos los expertos deben residir en algún lugar).
¿Recordar es comprimir?
HiPPO define la memoria como una aproximación óptima: guardar lo suficiente para reconstruir el pasado con error mínimo. La memoria humana, según Bartlett (1932), es reconstructiva: no reproducimos, recreamos.
Y MoE plantea otra pregunta vieja: ¿es la mente un sistema unificado o una sociedad de especialistas, como sugería Minsky en The Society of Mind?