ShadowPEFT

ShadowPEFT amplía un modelo base decodificador congelado con una red shadow ligera y preentrenable que se ejecuta en paralelo con el backbone. Un pequeño backbone shadow produce un estado shadow inicial s^(0), que luego recorre el bucle decodificador del modelo base: en cada bloque objetivo, la discrepancia entre los estados ocultos base y el estado shadow se inyecta de nuevo en la entrada del bloque (una corrección de bajo rango) y el estado shadow se avanza mediante una actualización residual con compuerta calculada a partir de la salida del bloque. Solo se entrenan los componentes shadow; el modelo base permanece congelado.

Input
  ├──► Shadow backbone (small, trainable) ──► s^(0)
  └──► Base model (frozen, large)
         block_0 ◄── inject(h, s) ─► h_0 ──► update ─► s_1
         block_1 ◄── inject(h, s) ─► h_1 ──► update ─► s_2
         ...                       (the (hidden, shadow) pair rides the loop together)

Dado que la adaptación es una trayectoria dependiente de la entrada en el espacio de capas (el estado sombra evoluciona con los datos) y no una delta estática en el espacio de pesos, ShadowPEFT no puede fusionarse con los pesos base. Llamar a merge, merge_adapter o merge_and_unload genera un error explícito. Para los modelos de lenguaje de Transformers, puedes obtener la red sombra ligera de forma independiente con model.base_model.unload_shadow(), que devuelve un DetachedShadowModel autónomo. La descarga independiente no está soportada para los modelos de Diffusers porque la reconstrucción de un desruidor completo depende de la arquitectura.

Agregar varios adaptadores, alternar entre ellos con set_adapter, eliminarlos y habilitarlos/deshabilitarlos funciona igual que con otros métodos PEFT. Solo puede estar activo un adaptador a la vez, ya que el estado sombra es una única trayectoria a través de la red.

El backbone de sombra se puede construir de dos maneras, controladas por ShadowConfig.shadow_model:

• "mirror" (predeterminado): se crea automáticamente un backbone de sombra más pequeño. Los modelos de lenguaje utilizan una copia reducida de la arquitectura base. Las arquitecturas de Diffusers con un backend registrado utilizan un modelo de arquitectura similar reducido, inicializado a partir de pesos base seleccionados; las arquitecturas compatibles sin backend recurren a un MLP residual por token. Cuando el tamaño oculto de la sombra difiere del base, una proyección entrenada cubre la diferencia.
• un id de modelo o ruta local: el backbone se carga como un modelo preentrenado más pequeño. Los modelos de Transformers usan AutoModel; los backends de Diffusers registrados definen su propia carga de checkpoints compatible.

El soporte de Diffusers con conciencia de la arquitectura se selecciona automáticamente a partir de la clase del modelo. Flux2 actualmente tiene un backend registrado; otros modelos de Diffusers compatibles basados en transformadores utilizan el respaldo genérico de MLP. La función independiente unload_shadow() sigue sin estar soportada para todos los modelos de Diffusers, ya que la reconstrucción de un denoiser completo depende de la arquitectura.

En comparación con los métodos de estilo LoRA, ShadowPEFT añade más parámetros y cómputo (ejecuta una red paralela y envuelve bloques completos del decodificador), pero el adaptador es una red autónoma que se puede entrenar de forma centralizada, reutilizar entre tareas e inicializar a partir de un modelo pequeño preentrenado. Una pérdida auxiliar opcional (auxiliary_loss_weight) aplica la cabeza de la tarea al estado inicial de sombra s^(0) y la suma a la pérdida de la tarea, animando a la ruta de sombra desmontable a resolver la tarea por sí sola. Para LM causal, se reutiliza la cabeza de salida base; incluye "lm_head" en modules_to_save para entrenarla y guardarla a través del mecanismo estándar de PEFT.

Caché KV

ShadowPEFT admite decodificación incremental con una caché KV dual: una para el modelo base congelado y otra para el esqueleto de sombra. La inyección/actualización es local al token, por lo que un token nuevo solo necesita su propio estado de sombra s; la causalidad mantiene válidas las claves/valores del prefijo base (calculadas bajo inyección). El objeto emparejado es una ShadowCache, devuelta como past_key_values cuando use_cache=True. Puedes pasar use_cache=True a generate() como de costumbre.

out = model.generate(input_ids, max_new_tokens=32)  # dual KV cache enabled by default
# or explicitly:
out = model.generate(input_ids, use_cache=True, max_new_tokens=32)

use_cache=False sigue funcionando y reprocesa la secuencia completa en cada paso (útil para depurar).

Uso

from transformers import AutoModelForCausalLM
from peft import ShadowConfig, get_peft_model

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B")
config = ShadowConfig(r=8, shadow_num_hidden_layers=1, task_type="CAUSAL_LM")
model = get_peft_model(model, config)
model.print_healthcheck()

out = model.generate(input_ids, max_new_tokens=32)

Para inicializar el shadow backbone a partir de un modelo preentrenado más pequeño, pase su id o ruta como shadow_model:

config = ShadowConfig(shadow_model="Qwen/Qwen3-0.6B", task_type="CAUSAL_LM")
model = get_peft_model(base_model, config)

Evaluando la ruta de sombra

Por defecto, la salida del modelo (logits) es el modelo base adaptado por sombra: las correcciones de sombra se inyectan en los estados ocultos del modelo base en cada capa, por lo que logits ya refleja ShadowPEFT (use model.disable_adapter() para obtener el modelo base simple para comparar). La pérdida auxiliar entrena adicionalmente la ruta de sombra para resolver la tarea por sí sola.

Para evaluar la red de sombra independiente de un modelo de lenguaje de Transformers (el modelo desmontable y ligero, análogo a ShadowPEFT de merge_and_unload), use unload_shadow(). Devuelve head(projection(backbone(x))) como un modelo de tarea normal que puede evaluar como cualquier otro: para una tarea de LM causal es un LM causal con capacidad de generación (admite generate() y caché de KV), y para una tarea de clasificación de secuencias agrupa el último token y devuelve los logits de clase. Llamar a este método para un modelo de Diffusers lanza NotImplementedError.

shadow = model.base_model.unload_shadow()  # a DetachedShadowModel (a PreTrainedModel)
shadow.eval()
# causal LM:
out = shadow.generate(input_ids, max_new_tokens=32)
# sequence classification:
logits = shadow(input_ids=input_ids, attention_mask=attention_mask).logits  # (batch, num_labels)

Por defecto (copy=False), el modelo devuelto comparte sus módulos con el modelo PEFT, y un backbone de sombra que comparte las incrustaciones de entrada base congeladas los alcanza a través de una referencia que no es un submódulo. Eso está bien para la evaluación, pero significa que save_pretrained escribiría un punto de control sin la tabla de incrustaciones. Pasa copy=True cuando quieras guardar o subir el modelo independiente:

shadow = model.base_model.unload_shadow(copy=True)
shadow.save_pretrained("standalone-shadow")

Resumen de las pruebas de rendimiento

API

ShadowConfig

ShadowModel

ShadowCache