Split-VIB Reconstructor for Feature-Level Intervention in LLM Hidden States

This document describes a model that:

Takes a single hidden representation extracted from a frozen LLM
Disentangles a specific linguistic feature (e.g. verb tense) from the rest of the information
Reconstructs the original hidden representation from the disentangled parts
Enables single-example intervention by modifying the feature representation and reinjecting it into the LLM

The model is trained using single-input forward passes, with pairwise supervision used only in the training loss.

1. Overall structure

For a single hidden state h:

h → shared_trunk → (z_feature , z_rest) → reconstructor → h_reconstructed

Where:

z_feature is a bottlenecked latent meant to encode only the target feature
z_rest is a bottlenecked latent meant to encode everything else
h_reconstructed approximates the original h

The frozen LLM is never updated.

h ∈ R^{d_h}
- A single hidden representation from a frozen LLM
- Typically: one token (e.g. the verb) at a chosen layer

s = f_{η} (h)

Shape: $h : (B, d_{h}) \to s : (B, d_{s})$
Purpose:
- Produce a shared intermediate representation
- Feed both latent branches
Trained by:
- Feature prediction loss
- Reconstruction loss
- KL losses
- Pairwise loss (during training)

$(μ_{f e a t u re}, lo g σ_{f e a t u re}^{2}) = q_{φ} (s)$
$z_{f e a t u re} N (μ_{f e a t u re}, d ia g (σ_{f e a t u re}^{2}))$

Shape: (B, d_s) → (B, d_feature)
Purpose:
- Encode only the target feature (e.g. tense)
- Be as small and information-limited as possible
Supervised by:
- Feature classification loss
- KL divergence to N(0, I)

(μ_{rest}, \logσ²_{rest}) = q_ψ(s) $$

z_{rest} ~ N(μ_{rest}, diag(σ²_{rest}))

- Shape: `(B, d_s) \to (B, d_rest)` - Purpose: - Encode all information needed to reconstruct `h` - Explicitly **exclude the target feature** - Supervised by: - Reconstruction loss - KL divergence - Adversarial loss (via GRL) - Pairwise rest-invariance loss (training only) --- ### 2.5 Reconstructor

h_{reconstructed} = g_θ([z_{feature} ; z_{rest}])

- Shape: `(B, d_feature + d_rest) \to (B, d_h)` - Purpose: - Ensure no information is lost by the split - Make `z_feature` and `z_rest` jointly sufficient - Trained by: - Reconstruction loss --- ### 2.6 Feature classifier

ŷ_{feature} = C_ω(z_{feature})

- Shape: `(B, d_feature) \to (B, num_classes)` - Purpose: - Force the feature information into `z_feature` - Trained by: - Cross-entropy loss on the feature label --- ### 2.7 Adversary with Gradient Reversal Layer (GRL)

ŷ_{rest} = A_ξ(GRL(z_{rest}))

- Shape: `(B, d_rest) \to (B, num_classes)` - Purpose: - Try to predict the feature from `z_rest` - Via GRL, force the rest encoder to **remove feature information** - Trained by: - Cross-entropy loss - GRL reverses gradients flowing into the rest encoder --- ## 3. Losses and what they enforce ### 3.1 Feature prediction loss

L_{feature} = CE(C_ω(z_{feature}), y)

- Enforces: - The target feature is encoded in `z_feature` - Updates: - Shared trunk - Feature latent head - Feature classifier --- ### 3.2 Feature KL loss (VIB)

L_KL_{feature} = KL(q(z_{feature} | h) || N(0, I))

- Enforces: - `z_feature` is minimal and compressed - Updates: - Shared trunk - Feature latent head --- ### 3.3 Reconstruction loss

L_{recon} = || h - h_{reconstructed} ||²

- Enforces: - `(z_feature, z_rest)` preserve all information in `h` - Updates: - Shared trunk - Both latent heads - Reconstructor --- ### 3.4 Rest KL loss (VIB)

L_KL_{rest} = KL(q(z_{rest} | h) || N(0, I))

- Enforces: - `z_rest` does not trivially copy `h` - Updates: - Shared trunk - Rest latent head --- ### 3.5 Adversarial loss (with GRL)

L_{adv} = CE(A_ξ(GRL(z_{rest})), y)

- Enforces: - Feature information is **removed** from `z_rest` - Updates: - Adversary learns to predict the feature - Rest encoder learns to *hide* the feature (via reversed gradient) --- ### 3.6 Pairwise rest-invariance loss (training only) For a minimal pair `(h⁰, h¹)` differing only in the feature:

L_pair = || μ_{rest}(h⁰) - μ_{rest}(h¹) ||²

- Enforces: - `z_rest` is invariant to the feature - Updates: - Shared trunk - Rest latent head - **Not used at inference** --- All losses except `L_pair` are computed **per example**. `L_pair` is computed **only during training**, using the two forward passes. --- ## 5. Training process (step-by-step) 1. **Freeze the LLM** - Choose a layer and token position (e.g. the verb) 2. **Extract hidden states** - For training: extract two hidden states per example (minimal pair) - For inference: extract one hidden state 3. **Forward pass (single-input model)** - Run the split-VIB model independently on each hidden state 4. **Compute losses** - Single-example losses for each forward pass - Pairwise rest-invariance loss using both outputs 5. **Backpropagation** - Update only the split-VIB model parameters - LLM remains frozen 6. **Inference / intervention** - Run model on a single hidden state - Modify `z_feature` - Reconstruct a new hidden state - Inject it back into the LLM and continue generation --- ## 6. Key property > **The model always operates on a single input.** > Pairwise supervision is a *training-time regularizer*, not part of the model interface. This guarantees: - Clean single-example interventions - Strong disentanglement from minimal-pair supervision - No mismatch between training and inference usage