La tesis central de Dust es que, a medida que la disponibilidad de cómputo a escala de hyperscaler aumenta, los algoritmos de aprendizaje basados en búsqueda y fuerza bruta, con menos dependencias de sesgos inductivos como la diferenciabilidad, pueden superar a los métodos tradicionales como Backprop. Este trabajo aborda el problema fundamental de la asignación de crédito en redes neuronales profundas, tradicionalmente resuelto por Backprop, proponiendo una alternativa de orden cero que opera en el espacio de activaciones en lugar del espacio de pesos.
Históricamente, Backprop ha sido el pilar del entrenamiento de redes neuronales debido a su eficiencia en el cálculo de gradientes de primer orden. Sin embargo, impone restricciones arquitectónicas y puede no explorar óptimamente el espacio de pérdida. Dust se alinea con la 'Bitter Lesson' de Sutton (2019), sugiriendo que los métodos generales que escalan con el cómputo eventualmente prevalecen, como se vio en AlphaGo Zero (Silver et al., 2017). La propuesta es que la diferenciabilidad y Backprop, si bien eficientes en regímenes de bajo cómputo, podrían limitar la exploración en regímenes de alto cómputo, abriendo la puerta a algoritmos más flexibles basados en búsqueda para una mejor generalización.
Arquitectura del Sistema
Dust es un algoritmo de optimización de orden cero que perturba las activaciones de las capas, recompensa cada perturbación por su impacto en la reducción de la pérdida y promedia las perturbaciones ponderadas por recompensa sobre una población para estimar el gradiente. A diferencia de las Estrategias de Evolución (ES) tradicionales que perturban los pesos (Salimans et al., 2017), Dust opera en el espacio de activaciones, una técnica conocida como 'node perturbation' (Werfel et al., 2003; Widrow y Lehr, 1990).
La innovación clave es la 'población virtual': en lugar de materializar cada miembro de la población con una copia perturbada de los pesos y una pasada hacia adelante individual, Dust añade ruido gaussiano a la salida de cada capa lineal, de forma independiente en cada token. Esto significa que cada token actúa como un miembro de la población virtual, y una única pasada hacia adelante evalúa miles de miembros en paralelo. Para una capa lineal, el ruido se añade como , donde es el ruido gaussiano y es la escala de ruido. La reducción de pérdida centrada se calcula como , donde es la pérdida perturbada y es la pérdida perturbada media. La recompensa del jitter en el token se basa en la reducción de pérdida en y, con un decaimiento , en las reducciones de pérdida de los tokens posteriores a los que el jitter también afecta a través de la atención. El ruido ponderado por recompensa, promediado sobre múltiples 'draws' (perturbaciones), se convierte en el error estimado en la salida de la capa. El gradiente de los pesos se calcula como el producto externo de este error estimado con la entrada de la capa, sumado sobre los tokens.
Para mitigar la interferencia entre múltiples perturbaciones en una misma pasada hacia adelante, Dust emplea tres estrategias: (1) Diferentes tipos de capas se perturban en pasadas hacia adelante separadas, cada una con su propia escala de ruido. (2) Los componentes internos de la atención (query, key, value, gate, value embedding) se perturban por separado y se recompensan a través de la salida de atención estimada, no directamente por la pérdida del token. (3) La capa de salida (language modeling head) se perturba directamente sobre los logits cacheados, reevaluando solo la entropía cruzada y una porción del vocabulario por 'draw', lo que reduce significativamente el costo computacional. Los hiperparámetros se ajustan mediante búsqueda en cuadrícula para maximizar la similitud de coseno con el gradiente de Backprop en un solo batch.
Flujo de Estimación de Gradiente con Dust
- 1 Entrada de Red Secuencia de tokens para el modelo Transformer.
- 2 Pasada Adelante (Clean) Cálculo de activaciones y pérdida sin perturbaciones, se cachean.
- 3 Perturbación de Activaciones Se añade ruido gaussiano a la salida de capas lineales, independientemente po...
- 4 Pasada Adelante (Perturbada) Se ejecuta la pasada hacia adelante con activaciones perturbadas.
- 5 Cálculo de Reducción de Pérdida Se calcula la reducción de pérdida centrada por token y su recompensa, con de...
- 6 Estimación de Error de Salida Se promedia el ruido ponderado por recompensa sobre 'draws' para estimar el e...
- 7 Cálculo de Gradiente de Peso Producto externo del error estimado con la entrada de la capa, sumado sobre t...
- 8 Actualización de Pesos Los pesos del modelo se actualizan usando el gradiente estimado.
| Capa | Tecnología | Justificación |
|---|---|---|
| compute | Transformer Language Models | Arquitectura de red neuronal objetivo para el preentrenamiento. Modelos GPT-style, 8 capas, ancho 512, hasta 243M parámetros. |
| data-processing | FineWeb dataset | Conjunto de datos de preentrenamiento para modelos de lenguaje. 4096-token BPE tokenizer, batch de 16k tokens (8 secuencias de 2048 tokens). |
| compute | SGD with momentum | Optimizador principal utilizado para el entrenamiento de Dust y Backprop. vs Adam Tasa de aprendizaje constante, un epoch. |
| compute | Adam optimizer | Optimizador alternativo probado para evaluar la compatibilidad de Dust con optimizadores modernos. vs SGD with momentum |
Trade-offs
Ganancias
- ▲▲ Eficiencia de población
- ▲ Exploración del espacio de búsqueda
- ▲ Flexibilidad arquitectónica (no requiere diferenciabilidad end-to-end)
- △ Rendimiento en preentrenamiento (supera a Backprop con suficiente cómputo)
Costes
- ▲▲ Eficiencia computacional (actualmente más costoso que Backprop para el mismo rendimiento)
- △ Complejidad de ajuste de hiperparámetros (múltiples escalas de ruido, decaimientos de crédito)
Fundamentos Teóricos
La idea de Dust se conecta directamente con el campo de la optimización de orden cero y las estrategias de evolución, que han sido estudiadas en la academia durante décadas. El concepto de 'node perturbation' (perturbación de nodos) fue introducido por Widrow y Lehr en 1990 y explorado por Werfel et al. en 2003, prediciendo la posibilidad de entrenar redes neuronales sin gradientes explícitos. Este trabajo extiende esas ideas a la escala de los Transformers modernos.
La comparación con Backprop, formalizada por Rumelhart, Hinton y Williams en 1986, es fundamental. Dust busca una alternativa a este algoritmo que ha dominado el aprendizaje profundo, explorando un camino que se alinea con la 'Bitter Lesson' de Richard S. Sutton (2019), que postula que los métodos generales que escalan con el cómputo eventualmente superan a los métodos con fuertes sesgos inductivos. Además, el trabajo se relaciona con la investigación sobre la exploración del espacio de pérdida, donde Liu et al. (2020) han demostrado que los métodos basados en gradientes pueden converger a mínimos globales subóptimos, sugiriendo la necesidad de algoritmos de búsqueda más flexibles.