🤖 machine learning

Budget Constraints as Riemannian Manifolds

Este trabajo propone la Optimización Restringida Riemanniana (RCO), un marco novedoso que modela las restricciones presupuestarias como variedades riemannianas suaves para permitir una optimización eficiente basada en gradientes de objetivos no descomponibles bajo una aplicación exacta del presupuesto, superando a los métodos existentes de penalización y evolutivos tanto en calidad de solución como en eficiencia computacional para tareas como la cuantización de precisión mixta y el recorte de expertos.

Autores originales: Michael Helcig, Dan Alistarh

Publicado 2026-05-04

📖 5 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Michael Helcig, Dan Alistarh

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que eres el chef principal de un restaurante masivo y de alta gama. Tienes un presupuesto estricto para la noche, pero también tienes un menú con cientos de platos, y cada plato se puede preparar de varias maneras diferentes (por ejemplo, usando ingredientes premium, ingredientes estándar o sustitutos económicos).

Tu objetivo es elegir exactamente una versión de cada plato para servir de modo que el costo total se mantenga exactamente dentro de tu presupuesto, mientras haces que la calidad general de la comida sea lo más deliciosa posible.

¿El problema? La calidad de la comida no es simplemente la suma de los platos individuales. Si eliges un filete elegante, podría combinar mejor con un vino específico, cambiando el "perfil de sabor" de toda la mesa. Esto hace que las matemáticas sean increíblemente difíciles: no puedes mirar cada plato de forma aislada; tienes que resolver un rompecabezas gigante y enredado donde cada elección afecta a todas las demás.

Este es exactamente el problema que enfrentan los ingenieros de aprendizaje automático al intentar comprimir modelos de IA enormes (como los que impulsan los chatbots). Necesitan decidir cuánto "encoger" o "podar" diferentes partes del modelo para ajustarse a un límite de tamaño (el presupuesto) sin arruinar la inteligencia del modelo (la calidad).

Así es como el artículo resuelve esto, utilizando algunas analogías creativas:

1. La Vieja Forma: Adivinar y Penalizar

Anteriormente, los ingenieros probaban dos enfoques principales, ambos torpes:

El Método de la "Penalización": Le decían a la computadora: "Intenta mantenerte bajo el presupuesto, pero si te pasas, te daré una gran 'multa' (una puntuación de penalización)". El problema es que la computadora es mala adivinando la multa correcta. Si la multa es demasiado pequeña, ignora el presupuesto. Si es demasiado grande, la computadora se asusta y deja de aprender. Es como intentar enseñarle a un perro a sentarse gritando "¡No!" a volúmenes aleatorios; el perro nunca aprende la regla exacta.
El Método "Evolutivo": Dejaban que la computadora probara miles de combinaciones aleatorias, guardaba las mejores y repetía el proceso. Esto funciona bien, pero es increíblemente lento. Es como intentar encontrar la mejor receta cocinando cada plato posible en el mundo y probándolos uno por uno. Toma una eternidad.

2. La Nueva Idea: El "Manifold del Presupuesto"

Los autores se dieron cuenta de que si observas el problema a través de una lente matemática específica (usando algo llamado "softmax"), la restricción del presupuesto no es un muro desordenado contra el que tienes que rebotar. En cambio, es una superficie suave y curva (un manifold) sobre la cual puedes caminar.

Piensa en el presupuesto no como una cerca dura, sino como un cuerda floja.

La Superficie: Imagina un trampolín gigante, invisible y curvo que existe solo donde tu costo total es exactamente igual a tu presupuesto.
El Paseo: La computadora no necesita saltar del trampolín y esperar aterrizar de nuevo sobre él. En su lugar, camina a lo largo de la superficie.

3. Cómo Funciona el Nuevo Método (RCO)

El artículo propone un nuevo algoritmo llamado Optimización con Restricciones Riemannianas (RCO). Así es como se mueve a lo largo de esa cuerda floja:

Paso 1: El Paso Tangente (Caminando hacia adelante): La computadora calcula la dirección que hace la comida más sabrosa (el gradiente). Pero en lugar de simplemente caminar en esa dirección, proyecta esa dirección sobre la superficie de la cuerda floja. Esto asegura que nunca se salga accidentalmente de la línea del presupuesto.
Paso 2: La Búsqueda Binaria (El Tobogán Mágico): A veces, incluso al caminar con cuidado, podrías desviarte ligeramente de la línea. En otros métodos, tendrías que realizar un cálculo complejo para volver. Aquí, los autores encontraron un "tobogán mágico". Debido a las matemáticas específicas que usaron, pueden simplemente deslizar todo el plan de comidas hacia arriba o hacia abajo con una sola perilla (una búsqueda binaria) para aterrizar perfectamente de nuevo en la línea del presupuesto. Es como tener un control remoto que arregla instantáneamente tu equilibrio.
Paso 3: El Momento (Manteniendo el Ritmo): Cuando caminas sobre una superficie curva, tu dirección cambia. El algoritmo tiene un truco especial para "transportar" su momento (su memoria de hacia dónde iba) para que no se maree ni pierda el ritmo mientras se mueve a lo largo de la curva.

4. Por Qué es un Gran Logro

El artículo afirma que este método es un cambio de juego por dos razones:

Es Exacto: A diferencia de los antiguos métodos de "penalización" que a menudo terminan ligeramente por encima o por debajo del presupuesto, este método se mantiene exactamente en la línea del presupuesto en cada paso. Es como un equilibrista que nunca tambalea.
Es Rápido: Porque utiliza gradientes (direcciones matemáticas) en lugar de adivinanzas aleatorias, encuentra la mejor solución mucho más rápido.
- El Resultado: En pruebas con rompecabezas sintéticos, los métodos antiguos se estancaron en el 83% de la puntuación posible más alta, mientras que este nuevo método encontró la solución perfecta.
- Mundo Real: Cuando lo probaron comprimiendo modelos de IA enormes (como reducir el tamaño de un "Modelo de Lenguaje Grande"), igualó o superó los resultados de los lentos métodos "evolutivos", pero lo hizo de 3 a 16 veces más rápido.

Resumen

El artículo presenta una nueva forma de resolver problemas de "presupuesto" en la IA. En lugar de tratar el presupuesto como un límite duro que rompe tus cálculos, lo convirtieron en una superficie suave y caminable. Al caminar a lo largo de esta superficie, la computadora puede encontrar el equilibrio perfecto entre costo y calidad mucho más rápido y con mayor precisión que antes, sin necesidad de adivinar o ajustar configuraciones complicadas. Es la diferencia entre tropezar en una habitación oscura tratando de evitar muebles y caminar con confianza por un camino bien iluminado y perfectamente pavimentado.

1. Planteamiento del Problema

El artículo aborda el problema de asignación discreta con restricciones presupuestarias, un desafío recurrente en el aprendizaje automático eficiente. El objetivo es asignar una de $K$ opciones discretas (por ejemplo, anchos de bits de cuantización, decisiones de poda o selección de expertos) a cada uno de los $N$ grupos (por ejemplo, capas de una red neuronal) para minimizar un objetivo global (por ejemplo, la pérdida del modelo) sujeto a un presupuesto total de costos (por ejemplo, tamaño del modelo o latencia).

Desafíos Clave:

Objetivo No Descomponible: La pérdida del modelo depende conjuntamente de todas las asignaciones (por ejemplo, interacciones entre capas), lo que significa que no puede escribirse como una suma de términos independientes por grupo. Esto impide que la Programación Dinámica (PD) estándar optimice la pérdida real directamente; la PD solo puede optimizar objetivos "proxy" (como la sensibilidad por capa).
Limitaciones de los Métodos Existentes:
- Resolvedores Combinatorios (PD): Manejan el presupuesto exactamente pero requieren objetivos descomponibles, forzando el uso de puntuaciones proxy inexactas.
- Búsqueda Evolutiva: Evalúa la pérdida real pero carece de gradientes, requiriendo cientos de pasadas hacia adelante y sufriendo altos costos computacionales.
- Métodos de Penalización/Lagrangianos: Añaden un término de penalización (por ejemplo, $\lambda(C-B)^2$ ) a la pérdida. Estos hacen cumplir el presupuesto solo de manera aproximada, a menudo oscilan y requieren un ajuste sensible de hiperparámetros.

2. Metodología: Optimización con Restricciones Riemannianas (RCO)

Los autores proponen un enfoque geométrico tratando la restricción presupuestaria no como una penalización, sino como una variedad Riemanniana suave en el espacio de los logits.

A. La Variedad Presupuestaria

Los autores parametrizan la distribución de asignación para cada grupo $i$ utilizando logits $\alpha_i \in \mathbb{R}^K$ , con probabilidades $p_i = \text{softmax}(\alpha_i)$ . El costo esperado total se define como:
$C(\alpha) = \sum_{i=1}^N w_i \langle p_i, c \rangle$
donde $w_i$ es el peso del grupo y $c$ es el vector de costos.

Estructura de la Variedad: Se demuestra que el conjunto de nivel $\mathcal{M} = \{ \alpha \mid C(\alpha) = B \}$ es una subvariedad Riemanniana suave de dimensión $(NK-1)$.
Propiedades Geométricas:
1. Normal de Forma Cerrada: El gradiente $\nabla C(\alpha)$ está disponible en forma cerrada y es no nulo (asumiendo costos distintos), definiendo el vector normal a la variedad.
2. Retracción Monótona: Desplazar los logits a lo largo del vector de costos cambia el costo esperado de manera monótona. Esto permite que el algoritmo proyecte cualquier punto de vuelta sobre la variedad $\mathcal{M}$ utilizando una simple búsqueda binaria (o el método de Newton) en lugar de proyecciones iterativas complejas.
3. Transporte Vectorial Económico: Mover el momento de un plano tangente a otro se reduce a una única operación de producto interno.

B. El Algoritmo RCO

El algoritmo propuesto envuelve la optimización de primer orden estándar (Adam) con tres operaciones geométricas para hacer cumplir el presupuesto exactamente en cada paso:

Proyección Tangente: El gradiente de la pérdida se proyecta sobre el espacio tangente de la variedad eliminando la componente a lo largo del vector normal $\nabla C(\alpha)$ . Esto asegura que el optimizador solo vea direcciones que preservan el presupuesto.
Retracción: Después de un paso de Adam, la iteración puede desviarse de la variedad. El algoritmo desplaza todos los logits a lo largo del vector de costos ( $\alpha' = \alpha + t \cdot c$ ) y utiliza búsqueda binaria para encontrar $t$ tal que $C(\alpha') = B$ .
Transporte de Momento: El momento (primer momento) del optimizador se proyecta sobre el nuevo espacio tangente para mantener la consistencia.

C. Manejo de la Factibilidad Discreta (Pasada Adelante)

Dado que la variedad opera en el espacio continuo de logits, el artículo combina RCO con:

Gumbel-STE (Estimador Straight-Through): Para generar asignaciones discretas para la pasada adelante.
Programación Dinámica con Restricción Presupuestaria: Para resolver el problema de asignación discreta (Problema de la Mochila de Elección Múltiple) exactamente durante la pasada adelante, asegurando que la configuración evaluada satisfaga el presupuesto.
Mecanismo: La pasada adelante utiliza PD para encontrar la mejor configuración discreta $z^*$ dados los logits perturbados. La pasada atrás utiliza STE para hacer fluir los gradientes a través de las probabilidades softmax de los logits perturbados, mientras que la proyección de la variedad asegura que estos gradientes no violen la restricción presupuestaria.

3. Contribuciones Clave

Identificación de la Variedad Presupuestaria: Los autores identifican que el conjunto de nivel del costo esperado softmax forma una variedad Riemanniana con propiedades geométricas únicas y computacionalmente económicas (normales de forma cerrada, retracción monótona).
Algoritmo RCO: Un marco de optimización novedoso que hace cumplir restricciones presupuestarias exactas con cero hiperparámetros de restricción (sin $\lambda$ que ajustar). Combina proyección tangente, retracción por búsqueda binaria y transporte de momento alrededor de Adam.
Manejo de Restricciones de Desigualdad: La variedad se extiende a restricciones de desigualdad ( $C(\alpha) \le B$ ) mediante una variable de holgura $s$ , permitiendo que el optimizador utilice naturalmente el presupuesto no utilizado sin desperdiciarlo en opciones inferiores.
Híbrido Discreto-Continuo: La integración de Gumbel-STE y PD con restricción presupuestaria permite la optimización de la pérdida del modelo real no descomponible, manteniendo al mismo tiempo la factibilidad discreta exacta.

4. Resultados Experimentales

El artículo valida RCO en tareas de compresión de Modelos de Lenguaje Grandes (LLM) sintéticas y del mundo real.

Problemas de Mochila Sintéticos (MCKP):
- En instancias difíciles, RCO recupera la solución óptima de la Programación Dinámica (PD).
- Los métodos Lagrangianos se estabilizan en el 83% del valor óptimo y exhiben violaciones significativas de restricciones ( $10^{-1}$ a $10^{-2}$ ), mientras que RCO mantiene violaciones cerca de la precisión de la máquina ( $10^{-9}$ ).
- RCO maneja perfectamente los casos "óptimos baratos" (donde la mejor solución cuesta menos que el presupuesto) utilizando la variable de holgura, mientras que los métodos con restricción de igualdad desperdician presupuesto.
Poda de Expertos MoE:
- Evaluado en OLMoE-1B-7B y Qwen3-30B-A3B.
- RCO iguala o supera la precisión de EvoESAP (una línea base de búsqueda evolutiva) en múltiples puntos de referencia (ARC, MMLU, etc.).
- Eficiencia: RCO logra estos resultados en 3–16× menos tiempo de reloj (por ejemplo, ~10 minutos frente a 5.8 horas para EvoESAP).
Cuantización de Precisión Mixta:
- Evaluado en Qwen3-8B (asignando anchos de bits de 2 a 8 a 252 capas).
- En alta compresión (2.25 bits), RCO reduce la perplejidad en un 36% en comparación con HIGGS (un método de sustituto lineal) y en un 4% en comparación con EvoPress (búsqueda evolutiva).
- RCO converge al mismo óptimo que EvoPress a 2.5 bits pero en ~10× menos tiempo (62 minutos frente a 11–14 horas).

5. Significado e Impacto

Aplicación Exacta de Restricciones: A diferencia de los métodos de penalización, RCO garantiza que el presupuesto se cumpla exactamente en cada iteración, eliminando la necesidad de ajustar hiperparámetros de pesos de penalización.
Optimización de Pérdidas No Descomponibles: Al separar la factibilidad discreta (gestionada por PD en la pasada adelante) de la optimización continua (gestionada por la variedad en la pasada atrás), RCO permite la optimización directa de pérdidas de modelos complejas y no descomponibles (como la precisión de extremo a extremo) bajo estrictas restricciones presupuestarias.
Eficiencia Computacional: Las operaciones geométricas (proyección, retracción, transporte) son computacionalmente económicas (esencialmente productos internos y búsquedas binarias), haciendo que el método sea escalable a modelos grandes donde la búsqueda evolutiva es prohibitivamente costosa.
Generalizabilidad: El marco es aplicable a cualquier problema que involucre parametrización softmax con costos lineales, extendiéndose más allá de la compresión de modelos hacia la búsqueda de arquitecturas neuronales con restricciones de recursos y otros problemas de asignación discreta.

En resumen, el artículo reformula la optimización discreta con restricciones presupuestarias como un problema de navegar una variedad Riemanniana específica, proporcionando una alternativa teóricamente fundamentada, libre de hiperparámetros y altamente eficiente frente a los enfoques basados en penalizaciones y evolutivos.