Learned KV-Cache Eviction

Resumen

Un proyecto final para la materia de Aprendizaje por Refuerzo de la UdeSA. Reformula la evicción de la KV-cache de un LLM congelado (Qwen2.5-1.5B-Instruct) como una secuencia de decisiones discretas y enmascarables, una por paso de decodificación, en vez del ranking one-shot KVP de Apple. Planteado como un entorno de Gymnasium y entrenado con MaskablePPO, la política elige qué tokens cacheados descartar en cada paso bajo un presupuesto de memoria fijo, evaluado en GSM8K, HotpotQA y una tarea sintética de recuperación de passkey.

Destacados

• Una escalera de trece experimentos (features ricas, warm-start, atención cross-token, recompensa densa causal, crédito por capa) que convergen en paridad con la heurística kv_norm sobre GSM8K.
• Un oráculo de atención futura muestra que la paridad es una propiedad del dataset, no del método: el margen aprendible sobre la heurística es +0.07 en GSM8K contra +0.43 en recuperación sintética de passkey y +0.09/+0.19 en HotpotQA, y crece con la agresividad de la compresión.
• En la arena con señal de passkey, una formulación online con recompensa densa causal muestra el primer desplazamiento sostenido por encima de la heurística, acercándose a la política de evicción oráculo.

Informe