ShadowPEFT

ShadowPEFT augments a frozen base decoder-only model with a lightweight, pretrainable shadow network that runs in parallel with the backbone. A small shadow backbone produces an initial shadow state s^(0), which then rides the base model’s decoder loop: at every targeted block the discrepancy between the base hidden states and the shadow state is injected back into the block input (a low-rank correction), and the shadow state is advanced by a gated residual update computed from the block output. Only the shadow components are trained; the base model stays frozen.

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)

Because the adaptation is an input-dependent trajectory in layer space (the shadow state evolves with the data) rather than a static weight-space delta, ShadowPEFT cannot be merged into the base weights. Calling merge, merge_adapter, or merge_and_unload raises an explicit error. For Transformers language models, you can obtain the lightweight shadow network on its own with model.base_model.unload_shadow(), which returns a standalone DetachedShadowModel. Standalone unloading is not supported for Diffusers models because reconstructing a complete denoiser is architecture-specific.

Adding multiple adapters, switching between them with set_adapter, deleting them, and enabling/disabling them all work as with other PEFT methods. Only one adapter can be active at a time, because the shadow state is a single trajectory through the network.

The shadow backbone can be built in two ways, controlled by ShadowConfig.shadow_model:

• "mirror" (default): a smaller shadow backbone is created automatically. Language models use a reduced copy of the base architecture. Diffusers architectures with a registered backend use a reduced architecture-alike model initialized from selected base weights; compatible architectures without a backend fall back to a token-wise residual MLP. When the shadow hidden size differs from the base, a trained projection bridges the gap.
• a model id or local path: the backbone is loaded as a smaller pre-trained model. Transformers models use AutoModel; registered Diffusers backends define their own compatible checkpoint loading.

Architecture-aware Diffusers support is selected automatically from the model class. Flux2 currently has a registered backend; other compatible transformer-based Diffusers models use the generic MLP fallback. Standalone unload_shadow() remains unsupported for all Diffusers models because reconstructing a complete denoiser is architecture-specific.

Compared to LoRA-style methods, ShadowPEFT adds more parameters and compute (it runs a parallel network and wraps whole decoder blocks), but the adapter is a self-contained network that can be trained centrally, reused across tasks, and initialized from a pre-trained small model. An optional auxiliary loss (auxiliary_loss_weight) applies the task head to the initial shadow state s^(0) and adds it to the task loss, encouraging the detachable shadow path to solve the task on its own. For causal LM, the base output head is reused; include "lm_head" in modules_to_save to train and save it through the standard PEFT mechanism.

KV cache

ShadowPEFT supports incremental decoding with a dual KV cache: one for the frozen base model and one for the shadow backbone. Inject/update are token-local, so a new token only needs its own shadow state s; causality keeps prefix base keys/values (computed under injection) valid. The paired object is a ShadowCache, returned as past_key_values when use_cache=True. You can pass use_cache=True to generate() as usual.

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 still works and reprocesses the full sequence each step (useful for debugging).

Usage

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_trainable_parameters()

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

To initialize the shadow backbone from a smaller pre-trained model, pass its id or path as shadow_model:

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

Evaluating the shadow path

By default the model output (logits) is the shadow-adapted base model: the shadow corrections are injected into the base model’s hidden states at every layer, so logits already reflects ShadowPEFT (use model.disable_adapter() to get the plain base model for comparison). The auxiliary loss additionally trains the shadow path to solve the task on its own.

To evaluate the standalone shadow network for a Transformers language model (the detachable, lightweight model — the ShadowPEFT analogue of merge_and_unload), use unload_shadow(). It returns head(projection(backbone(x))) as a normal task model that you can evaluate like any other: for a causal-LM task it is a generation-capable causal LM (supports generate() and KV caching), and for a sequence-classification task it pools the last token and returns class logits. Calling this method for a Diffusers model raises 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)

By default (copy=False) the returned model shares its modules with the PEFT model, and a shadow backbone that shares the frozen base input embeddings reaches them through a reference that is not a submodule. That is fine for evaluation, but it means save_pretrained would write a checkpoint without the embedding table. Pass copy=True when you want to save or push the standalone model:

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

API

ShadowConfig

ShadowModel

ShadowCache