Fine-tuning de Needle 3 en Español con LoRA

Este artículo está basado en un experimento local con Needle 3: datos sintéticos en inglés y español, entrenamiento LoRA en mi GTX 1080 Ti y evaluación corriendo el motor real del modelo.

Introducción

Hace una semana Cactus publicó Needle 3, un modelo pequeño especializado en function calling y extracción estructurada. No es un modelo para platicar: le das una lista de herramientas y una petición en lenguaje natural, y te regresa qué función llamar y con qué argumentos. Va de 25 a 121 millones de parámetros, pesa entre 8 y 29 MB en su cuantización CQ2 (2.125 bits por peso) y corre en CPU con un motor en C de menos de 1 MB por plataforma.

La arquitectura es lo que más me llamó la atención. Es una Simple Attention Network, o sea, no tiene el bloque FFN que tienen todos los transformers. El mezclado de canales lo hace un MLP tipo Monarch-Hadamard, y la memoria de hechos vive en un engram, que son tablas de n-gramas hasheados (70.8 millones de los 121 millones de parámetros del modelo de 20 capas). Además usan algo que llaman intelligence laddering: cada profundidad de 2 a 20 capas es un modelo que se puede desplegar por sí solo, porque durante el entrenamiento van muestreando profundidades y destilando hacia el modelo completo.

Lo que promete Cactus es que con fine-tuning sobre DroidCall el modelo sube entre 18 y 36 puntos en cada peldaño, y que desde 4 capas (29M) el modelo ajustado le gana a DeepSeek V4 Flash. Yo quería comprobar dos cosas:

1. Si el fine-tuning local con LoRA, que es gratis, realmente mejora el modelo.
2. Si funciona en español.

La misma documentación ya anticipa problemas: la cabeza de confianza no se actualiza con un LoRA local, y en idiomas distintos al inglés la confianza no es confiable (reportan llamadas correctas en español con confianza 0.0).

El Problema: Falsos Positivos con Confianza Alta

Antes de entrenar nada, probé el modelo base con seis herramientas de domótica. Estas son respuestas reales del set de prueba en español:

> "dame una receta de pasta"
play_music(query="pasta")                    # confidence 1.0

> "resúmeme este artículo"
play_music(query="resúmeme artículo")        # confidence 1.0

> "baja el volumen de la radio"
control_lights(room="radio", state="on")     # confidence 1.0

Que se equivoque en algo que está fuera de sus herramientas es esperable. Lo que preocupa es que lo haga con confianza de 1.0. La idea de Needle es que uses la confianza para decidir si ejecutas la llamada, si le pides confirmación al usuario o si la rechazas, y si la confianza no está calibrada ese enrutado no sirve. En español el modelo base solo rechazó 2 de 12 peticiones fuera de tema; en inglés, 7 de 12.

El Experimento

La tarea es un asistente de casa con 6 herramientas: luces, termostato, puertas, alarma, mensajes y música.

• Datos: 300 ejemplos de entrenamiento, 40 de validación y 120 de prueba por idioma. Las plantillas y los valores (habitaciones, puertas, nombres) no se repiten entre los splits, así que el modelo no puede memorizar el set de prueba. El 10% son peticiones que se deben rechazar (answers: []) y el 15% piden dos llamadas.
• Generación: los datos los generó otro modelo y los validé con un script: que cumplan el esquema, que cada argumento aparezca en la petición y que no haya duplicados entre splits.
• Métrica: exact match ordenado, es decir, mismo nombre de función, mismo orden y todos los argumentos iguales. También calculo dos variantes más permisivas: loose, que ignora argumentos opcionales que el modelo agrega sin que se los pidan, y loose-ci, que además compara sin distinguir mayúsculas.
• Inferencia: todo corre en el motor real de Needle, con reset() entre cada petición. Sumo function_calls y suppressed_calls porque los modelos ajustados localmente no tienen cabeza de confianza.

Cada línea del dataset se ve así:

{"query": "baja las luces de la terraza al 60", "tools": [...], "answers": [{"name": "control_lights", "arguments": {"room": "terraza", "brightness": 60}}], "reasoning": "'terraza' -> room; 'al 60' -> brightness"}

Entrenamiento

La configuración fue LoRA de rango 16 y alpha 32, sobre las proyecciones q/k/v/gate/out de las 20 capas. El modelo base queda congelado en float32 y el entrenamiento simula la cuantización W4 del export (STE). Batch de 16, learning rate de 1e-4 con warmup y cosine, 25 épocas (425 pasos) y 10% de los datos para validación.

pip install "cactus-needle[train]==3.0.4"
needle finetune es_train.jsonl --epochs 25 --out adapter_es.safetensors
needle build --lora adapter_es.safetensors --out tuned_es_20l.cact
needle build --lora adapter_es.safetensors --layers 8 --out tuned_es_8l.cact

Aquí me encontré con el primer problema. Mi tarjeta de video sigue siendo la GTX 1080 Ti, que es arquitectura Pascal (sm_61), y JAX 0.11 ya no genera código para ella. Las operaciones básicas de JAX sí corren, pero el train_step truena al cargar. La solución fue crear otro entorno virtual con una versión anterior de JAX que todavía soporta sm_61:

python3 -m venv .venv-gpu
.venv-gpu/bin/pip install "cactus-needle[train]==3.0.4" "jax[cuda12]==0.4.38"

Con eso cada paso tarda como 2.5 segundos en la GPU, unos 18 minutos por idioma. En CPU también funciona, pero es alrededor de un minuto por paso, casi 7 horas por idioma.

Para usar el modelo ajustado solo hay que pasarle el archivo:

agent = needle.Needle(tools=[...], weights="tuned_es_20l.cact")
agent.complete("baja las luces del salón al 20")
# confidence: None, el LoRA local no entrena la cabeza de confianza y needle build la descarta

Resultados

Exact match estricto sobre los 120 casos de prueba de cada idioma:

Y el detalle entre el modelo base y el de 20 capas:

¿Qué fue lo que mejoró?

En estricto la mejora es de +15 en inglés y +18 en español, que se ve muy bien. Pero con loose-ci baja a +3 y +5, y la columna de nombres casi no se mueve (110 a 110, y 104 a 105). Eso quiere decir que el modelo ya escogía la herramienta correcta; lo que aprendió fue a llenar los argumentos como yo los espero. Por ejemplo:

> "pon las luces del estudio a 65 por ciento"
base:   control_lights(room="estudio", brightness=65, state="on")
tuned:  control_lights(room="estudio", brightness=65)

> "cierra con llave la puerta del desván ya"
base:   lock_door(door="desván")
tuned:  lock_door(door="puerta del desván")

La primera corrección es buena: el modelo base agrega state="on" en cualquier comando de brillo. La segunda es más bien que el modelo aprendió las convenciones de mi dataset. No está mal, para eso es el fine-tuning, pero hay que tenerlo en cuenta antes de celebrar los 18 puntos.

Por herramienta, las luces pasan de 7 a 15 aciertos en inglés y de 8 a 15 en español, y las puertas de 14 a 17 y de 9 a 15. set_alarm sigue siendo la peor (de 11 a 13 y de 7 a 9), y el termostato empeora en uno en ambos idiomas.

Lo que no mejoró

Los rechazos quedaron igual: 7 de 12 en inglés y 2 de 12 en español. Aunque el 10% de los datos de entrenamiento eran peticiones para rechazar, el modelo ajustado sigue llamando herramientas cuando no debe:

> "dame una receta de pasta"
play_music(query="receta de pasta")

> "convierte cincuenta euros en dólares"
set_thermostat(room="cocina", temperature=50, mode="heat")

En español los acentos son frágiles. El modelo ajustado escribe aceo en lugar de aseo en varios casos, y el tokenizador gasta más tokens en español que en inglés. Fuera de eso, el comportamiento en español es muy parecido al de inglés, que era lo que más me preocupaba.

La pérdida de validación bajó de manera constante (inglés 0.585 a 0.144, español 0.530 a 0.137), así que no es falta de entrenamiento ni sobreajuste: los errores que quedan son comportamientos que los datos no cubren.

¿Por qué Fallan los Modelos de 4 y 8 Capas?

Los recortes de 8 y 4 capas colapsan: 42 y 24 aciertos en inglés, 25 y 13 en español, y ninguno rechaza nada (0 de 12).

El LoRA local solo entrena el modelo completo de 20 capas. La escalera original se entrena con self-distillation: el 80% de los pasos usa el modelo completo y el 20% una profundidad aleatoria entre 2 y 19 capas, con una pérdida que la acerca al modelo completo. Así cada submodelo aprende a hacer la tarea con las capas que tiene. Cuando recortas después con --layers 8, el submodelo hereda los pesos ajustados, pero nunca lo entrenaron para resolver la tarea con solo 8 capas.

Los números de 4 capas que publica Cactus vienen de su plataforma, que entrena el modelo completo en cada profundidad y además calibra la cabeza de confianza con tus herramientas. Ese camino no lo probé; cuesta 19 dólares (3 entrenamientos) y por lo que documentan es la única forma de obtener peldaños pequeños ajustados.

Conclusión

El fine-tuning local sí funciona, y funciona en español. Te da un modelo de 20 capas ajustado a una tarea estrecha, con +15 y +18 puntos en estricto, que corre en CPU y pesa 63 MB (con W4, contra los 35 MB del modelo original a 2 bits).

Lo que no te da es lo más interesante de Needle: rechazos confiables, modelos de 4 u 8 capas y la confianza para decidir si ejecutar una llamada. Si tu caso de uso necesita alguna de esas tres cosas, lo siguiente que probaría es subir el rango del LoRA a 32 con más ejemplos negativos, o pagar la plataforma. Para un asistente sencillo donde siempre hay una herramienta que llamar, el modelo de 20 capas ajustado localmente ya es usable.