Can LLMs encode Protected Health Information as structured, extractable concept vectors?
This project answers yes — and shows what that means for interpretability-informed anonymization.
Clinical text contains Protected Health Information (PHI) — 18 identifier types under HIPAA. Anonymizing this data is a prerequisite for safe use of clinical text in AI research.
Existing approaches treat the model as a black box:
- Rule-based regex — brittle, fails on edge cases
- ML-based NER — surface features, misses contextual identifiers
- LLM prompting — unreliable, no mechanistic grounding
This project takes a different approach: operate inside the model's representation space.
We show that LLMs encode PHI as structured directional features in their residual stream, that these features are linearly extractable, and that they can be used for interpretability-informed classification with strong performance.
A logistic regression probe trained on last-token activations achieves AUC 0.979 at Layer 3 of Llama-3.2-1B-Instruct — using only 1200 samples, balanced, 5-fold cross-validated.
Layer AUC F1 Recall MCC
L03 0.9792 0.9352 0.9400 0.8701 ← best single layer
L16 0.9672 0.9051 0.8833 0.8166 ← secondary peak
PHI is not distributed randomly across layers. It exhibits a W-shaped encoding curve: sharp peak at L03, mid-layer dip (L05–L12), partial recovery at L15–L16.
Direction vectors computed at L03 and L16 have cosine similarity of 0.035 — nearly orthogonal.
L03 direction: syntactic/surface name recognition (early, sharp)
L16 direction: generation-oriented PHI re-encoding (late, broader)
These are not redundant. The ensemble of L03 + L16 via weighted vote outperforms any single layer:
L03 baseline AUC=0.9792 Recall=0.9400 MCC=0.8701
L03+L16 ensemble AUC=0.9850 Recall=0.9533 MCC=0.9134
For privacy-critical use cases, the recall gain (0.940 → 0.953) translates to 13 fewer missed PHI instances per 1000 clinical sentences.
PHI direction AUC increases with sentence complexity across both layers:
Simple (<20w) Compound (20-40w) Complex (>40w)
L03: 0.876 0.881 0.891
L16: 0.865 0.874 0.912
The direction becomes more reliable for longer clinical sentences — exactly the regime where clinical notes live.
The cross-layer direction similarity heatmap reveals block structure:
emb–L03 : shallow regime (syntactic, low mutual similarity with middle/late)
L04–L10 : middle regime (semantic entanglement, high within-group similarity)
L11–L16 : late regime (generation-oriented, high within-group similarity)
PHI is re-encoded from scratch at each regime boundary.
Previous experiments extracted only post-MLP hidden states (the accumulated residual stream after each full transformer layer). We extended extraction to four activation types to decompose exactly where and how PHI encoding happens inside each layer.
Per transformer layer i:
h_in ──► [LayerNorm → Attention] ──► attn_out
h_mid = h_in + attn_out ← post_attn (accumulated, after attention)
h_mid ──► [LayerNorm → MLP] ──────► mlp_out
h_out = h_mid + mlp_out ← post_mlp (accumulated, after MLP)
attn_delta = h_mid - h_in ← pure attention contribution (no history)
mlp_delta = h_out - h_mid ← pure MLP contribution (no history)
Accumulated modes (post_attn, post_mlp) carry the full residual stream at a snapshot in time. Delta modes (attn_delta, mlp_delta) carry only what that specific sublayer writes — no accumulated history.
| Mode | Best Layer | Best AUC | Interpretation |
|---|---|---|---|
| post_mlp | L03 | 0.9792 | Full stream after MLP — prior best |
| post_attn | L02 | 0.9892 | Full stream after attention — new best |
| attn_delta | L02 | 0.9842 | Attention's isolated contribution |
| mlp_delta | L02 | 0.9775 | MLP's isolated contribution |
All four modes peak at transformer layer 3. post_attn beats post_mlp at the same layer
because the MLP actively degrades the PHI signal (see Finding 5 below).
Within transformer layer 3, the MLP delta is anti-correlated with the attention direction:
Cross-mode cosine similarity at transformer L3:
post_mlp post_attn attn_delta mlp_delta
post_mlp 1.000 0.630 0.463 0.484
post_attn 0.630 1.000 0.821 -0.374 ← negative
attn_delta 0.463 0.821 1.000 -0.373 ← negative
mlp_delta 0.484 -0.374 -0.373 1.000
Attention pushes activations toward the PHI direction. MLP pushes back against it (-0.37). This is why AUC drops from 0.9892 (post_attn) to 0.9792 (post_mlp) within a single layer — the MLP is not failing to encode PHI, it is encoding other information that competes with the PHI direction.
Cross-layer direction similarity heatmaps for delta modes are nearly all uncorrelated — each layer's attention delta and MLP delta fire into fresh, independent subspaces.
post_mlp / post_attn: clear block structure (adjacent layers correlated)
attn_delta / mlp_delta: near-zero cross-layer correlation throughout
This confirms the core mechanistic interpretability assumption about residual streams: each sublayer writes an independent additive contribution into orthogonal dimensions of the residual stream. The accumulated stream is a genuine superposition of independent feature directions — not an iterated refinement of the same representation.
Implication: visualising the latent space via PCA or direction projections is meaningful for accumulated modes (post_attn, post_mlp) — these carry coherent representations. Delta modes are write operations, not representations — visualising them as clusters is not meaningful.
The cross-layer direction similarity plots reveal a gradient:
Early layers (L1–L3): cosine 0.8–1.0 within group → similar representations
Cross-block transition: cosine drops toward 0 → representation rotates
Late layers (L11–L16): cosine 0.6–0.9 within group → different but coherent
Early → Late cross-block cosine: 0.03–0.10 → nearly orthogonal
Early layers encode what the name IS (syntactic entity, surface identity). Late layers encode what to DO with it (track for coherent generation). These live in perpendicular subspaces because they serve different computational purposes — not more refined versions of the same thing, but different computations entirely.
Extending ensembles to combine different activation modes revealed that cross-mode combinations consistently outperform same-mode combinations.
| Ensemble | Strategy | AUC | Recall | MCC |
|---|---|---|---|---|
| post_mlp L03 (old best single) | — | 0.9792 | 0.9400 | 0.8701 |
| post_attn L02 (new best single) | — | 0.9892 | 0.9434 | 0.8952 |
| attn_delta L02 + post_mlp L16 | concat | 0.9919 | 0.9483 | 0.9152 |
| post_attn L02 + post_mlp L16 | weighted vote | 0.9896 | 0.9567 | 0.9252 |
| post_attn L02 + L15 | concat | 0.9917 | 0.9383 | 0.9106 |
| post_mlp L03 + L16 (old winner) | weighted vote | 0.9856 | 0.9550 | 0.9150 |
AUC winner: attn_delta L02 + post_mlp L16 concat → AUC 0.9919
Recall winner: post_attn L02 + post_mlp L16 weighted vote → Recall 0.9567
For privacy-critical deployment, recall is the primary metric. The new recall winner improves from 0.9533 (old winner) to 0.9567 — 3 fewer missed PHI instances per 1000.
Ensemble Cosine Sim Orthogonal?
attn_delta L02 + post_mlp L16 (AUC winner) 0.019 ✅
post_attn L02 + post_mlp L16 (recall winner) 0.035 ✅
post_attn L02 + L15 (concat winner) 0.052 ✅
post_mlp L03 + L16 (old winner) 0.036 ✅
All top ensembles combine near-perfectly orthogonal components. The ensemble gains are geometrically guaranteed — orthogonal directions capture non-overlapping subspaces, so combining them adds independent PHI signal with zero redundancy.
INPUT SENTENCE
│
[Embedding] — no PHI signal (AUC ≈ 0.52)
│
L1–L3 Attention ──► writes PHI direction sharply into residual stream
L3 MLP ──► writes anti-correlated direction (-0.37) → partially cancels
↑ this is why post_attn > post_mlp at L3
│
L4–L12 (each sublayer writes to orthogonal subspace)
PHI signal dilutes as model builds richer semantic representation
W-shape trough in AUC, bowl in Cohen's d
│
L13–L16 Attention ──► reconstructs PHI for generation
different subspace from L3 (cosine 0.035)
late recovery peak
│
LAST TOKEN ACTIVATION
│
[Best single probe] post_attn L02 AUC = 0.9892
[Best ensemble] attn_delta L02
+ post_mlp L16 concat AUC = 0.9919
(cosine between components = 0.019 → orthogonal)
We tested seven layer combinations across two strategies to determine whether combining probes from multiple layers outperforms the single best layer.
Weighted Vote — one probe per layer, probabilities combined weighted by that layer's prior AUC score. High-AUC layers dominate the decision. Lightweight, composable.
Concatenation — activations from target layers stacked horizontally [n_samples, k × 2048],
single logistic regression on the combined vector. Lets the probe learn which dimensions
from which layer matter most. More expressive but harder to fit with limited data.
| Ensemble | Strategy | AUC | F1 | Recall | MCC |
|---|---|---|---|---|---|
| L03 (baseline) | weighted vote | 0.9792 | 0.9352 | 0.9400 | 0.8701 |
| L03 + L16 (two peaks) | weighted vote | 0.9850 | 0.9565 | 0.9533 | 0.9134 |
| L03 + L16 (two peaks) | concat | 0.9848 | 0.9486 | 0.9383 | 0.8986 |
| Top-3 AUC (L03+L02+L16) | weighted vote | 0.9876 | 0.9480 | 0.9417 | 0.8969 |
| Top-3 AUC (L03+L02+L16) | concat | 0.9863 | 0.9520 | 0.9450 | 0.9054 |
| L02+L03+L04 (peak neighbours) | weighted vote | 0.9831 | 0.9349 | 0.9333 | 0.8702 |
| L02+L03+L04 (peak neighbours) | concat | 0.9844 | 0.9372 | 0.9317 | 0.8752 |
| All early (L01–L04) | weighted vote | 0.9844 | 0.9436 | 0.9333 | 0.8890 |
| All early (L01–L04) | concat | 0.9857 | 0.9385 | 0.9283 | 0.8787 |
| L03+L15+L16 (early+late) | weighted vote | 0.9837 | 0.9385 | 0.9284 | 0.8787 |
| L03+L15+L16 (early+late) | concat | 0.9830 | 0.9423 | 0.9267 | 0.8875 |
| All late (L13–L16) | weighted vote | 0.9657 | 0.8988 | 0.8817 | 0.8025 |
| All late (L13–L16) | concat | 0.9682 | 0.9027 | 0.8833 | 0.8112 |
Weighted vote consistently beats concat for same layer combinations.
Concatenation creates a [1200, k×2048] input — logistic regression is fitting a
harder problem with the same data volume. Weighted vote sidesteps this by composing
already-reliable single-layer probes. Less data pressure, less overfitting risk.
L03 + L16 weighted vote is the clear winner on recall and MCC — the two metrics most relevant to privacy-critical deployment. Despite Top-3 AUC having marginally better raw AUC (0.9876 vs 0.9850), it has lower recall (0.9417 vs 0.9533). Adding L02 (a neighbour of L03) introduces redundancy — correlated layers dilute the diversity gain without contributing new signal.
Peak neighbours (L02+L03+L04) barely improve over L03 alone. Adjacent layers have high mutual direction similarity (L02–L03 cosine sim ≈ 0.49). Ensembling correlated probes = noise averaging, not complementary information. Diversity of mechanism matters more than quantity of layers.
Late-only ensemble is worst performer across the board. Late layers alone (L13–L16) consistently underperform early layers alone. The mid-to-late dip is real — late recovery is partial, not full. Never use late-only in deployment.
The ensemble gain is mechanistically justified, not accidental. L03 and L16 have cosine similarity 0.035 — nearly orthogonal directions. They capture genuinely different aspects of PHI. Their combination is complementary by geometry, which is why the gain is consistent across all metrics.
2D PCA projections (StandardScaler → PCA) show partial PHI/non-PHI separation along PC1 at both layers:
- L03: PC1=8.1%, PC2=6.3%. Elongated diagonal shape — PHI (red) pulls toward high PC1, non-PHI clusters left. Separation is real but PHI information spans more than 2 principal components.
- L16: PC1=6.6%, PC2=4.4%. Fan/cloud shape, less directional. Late-layer representations carry richer information (sentence meaning, clinical context, generation prep) — PHI is a smaller fraction of total variance.
Neither layer shows clean 2D blobs because PCA from 2048D to 2D is brutal compression. The probe operates in full 2048D space — scatter messiness does not contradict probe performance.
Mean-difference direction vectors projected across all samples reveal:
- Both layers: PHI class shifted right relative to non-PHI. Direction is real.
- Direction AUC: L03=0.876, L16=0.882 — well above random (0.5), but below probe AUC.
- Overlap: both distributions overlap substantially. PHI is encoded across multiple dimensions, not a single clean axis. Mean-difference captures the dominant direction but misses complementary subspaces.
This 10pp gap (direction AUC vs probe AUC) is the mechanistic explanation for why single-direction activation steering is imprecise — the signal is real but distributed.
These projection scores are the prototype for future per-sentence PHI scoring. No classifier needed — just a dot product with the direction vector.
Cosine similarity between PHI direction vectors across all 17 layers reveals three distinct representational regimes:
emb–L03 : low similarity with all other layers → shallow syntactic regime
L04–L10 : high within-group similarity (0.4–0.7) → semantic entanglement regime
L11–L16 : high within-group similarity (0.5–0.8) → generation-oriented regime
Key number: L03 vs L16 cosine similarity = 0.035. Nearly orthogonal. PHI is re-encoded from scratch between early and late layers — not refined, fundamentally re-represented. This is the geometric confirmation that the ensemble gain is real and mechanistically justified.
PHI direction AUC by sentence complexity bucket:
Simple (<20w) Compound (20-40w) Complex (>40w)
L03: 0.876 0.881 0.891
L16: 0.865 0.874 0.912
Direction reliability increases with sentence length — counterintuitive but explainable. Longer sentences provide more clinical context, making name-as-PHI more unambiguous in the model's internal representation. Clinical notes are almost always complex (>40w). Direction generalises well exactly where it needs to. ✅
PHI — a privacy/safety concept directly relevant to alignment — is cleanly linearly encoded in the residual stream (probe AUC 0.979, direction AUC 0.876). The mean-difference direction vector is real, consistent, and generalisable across sentence complexity.
This provides empirical grounding for the broader claim that alignment-relevant concepts (harmfulness, deception, refusal, private identity) can be extracted as concept vectors from LLM activations.
The 10pp gap between probe AUC (0.979) and direction AUC (0.876) reveals why.
PHI is not a single clean direction — it is distributed across multiple subspaces. A mean-difference steering vector captures the dominant axis but misses complementary signal. L03 and L16 directions are nearly orthogonal (cosine sim 0.035): at minimum two directions are required to represent PHI adequately for steering purposes.
Single-direction steering will be imprecise by design — the overlapping projection score distributions (histogram) confirm that any threshold on a single direction will have meaningful error on both sides.
Multi-directional steering (weighted combination of L03 + L16 directions) is a more principled approach and represents the next experimental step.
Current experiments localise PHI encoding to specific layers (L03, L16), not specific attention heads or MLP neurons.
The orthogonal directions and distinct representational regimes strongly suggest that separate circuits handle PHI at different abstraction levels — and the search space is now narrowed to two layers rather than all 16.
Full circuit isolation (activation patching, head ablation, neuron-level analysis) within L03 and L16 is a natural and tractable next step — not yet demonstrated.
Sentence-level Anonymization
│
├── Rephrasing (Generational)
│ ├── Activation steering ← multi-direction steering (next step)
│ └── Agentic pipeline
│
└── Redaction
├── Classification
│ ├── Activation alignment scoring ← current work (✅ validated)
│ ├── Layer + classifier ← token-level extension (next)
│ └── Circuit + classifier ← after circuit isolation
│
└── Clinical alignment scoring ← parallel track (clinical direction)
These observations emerged from close reading of results. Some are confirmed findings, some are directionally correct but require larger data to be conclusive.
MLP opposition is stronger than it looks. cosine(-0.37) = arccos(-0.37) ≈ 112 degrees. Roughly 37% of MLP's output magnitude at L3 is actively working against the PHI direction. This is not noise — it is a structural feature of how the model processes this layer.
Context amplifies attention signal. attn_delta (0.9842) < post_attn (0.9892) at the same transformer layer. The gap = contribution of accumulated L1–L2 context feeding into L3's attention. Each attention layer's isolated contribution is meaningful, but the residual stream carrying previous layers' encoding forward provides additional discriminative signal.
MLP encodes PHI in its own direction, not attention's. mlp_delta peaks at L3 (AUC 0.9775) but lower than attn_delta (0.9842). The probe still finds PHI signal in MLP's isolated output — but MLP's PHI direction is weaker and partially conflicted with the attention direction (-0.37 cosine). MLP is not ignoring PHI; it is encoding it in a competing subspace.
L9 is the point of maximum entanglement. The trough in all AUC curves occurs around L9 — roughly the midpoint of the 16 transformer layers. At this point the model has moved furthest from surface syntactic features (early regime) and is not yet in the generation-oriented late regime. PHI signal is most diffuse relative to everything else the model is computing.
Post-attn stability = attention consistently tracks PHI across all layers. post_attn shows the least volatile AUC curve. If this were purely due to accumulated stream persistence, post_mlp (also accumulated) would show equal stability — it does not. The difference: MLP delta contributions are noisy and variable (sometimes anti-correlated with PHI direction), injecting volatility into post_mlp. Attention contributions are consistently PHI-aligned across all layers, producing a smooth post_attn curve. Each attention layer keeps PHI in mind — L3 does it best, but the others maintain it. This is consistent with Elhage et al. (2021): attention acts as an associative relational memory that naturally sustains identity-type features (like person names) across depth. MLP layers are more variable key-value lookups that can reinforce or oppose this signal.
Ensemble differences beyond 0.003 AUC are inconclusive at this dataset size. With 1200 balanced samples and 5-fold CV, the noise floor is approximately 0.002–0.003 AUC. Differences smaller than this between ensemble configurations cannot be attributed to genuine mechanistic differences — they require significantly larger datasets (10k+) to validate. The directional conclusions hold; the specific rankings do not.
Delta modes confirm residual stream superposition empirically. Near-zero cross-layer cosine similarity in both attn_delta and mlp_delta across all layers is direct empirical confirmation that each sublayer writes to fresh orthogonal dimensions. This validates a core assumption in mechanistic interpretability — the residual stream is a genuine superposition of independent additive contributions, not an iterated refinement of a single representation.
Current: Alignment vector + bias vector experiment at sentence level (last token)
Compute PHI direction at L03 and L16. Test two scoring approaches:
- Raw cosine:
score = cosine(activation, direction)- Bias-corrected cosine:
score = cosine(activation - bias_vector, direction)wherebias_vector = mean(non-PHI activations)— strips generic clinical language, leaves PHI-specific residual.Compare both against full probe (AUC 0.979) and raw direction (AUC 0.876).
Next: Token-level alignment — apply the same direction + bias to every token position in a sentence. Validate that name tokens score highest. This is the actual meat: per-token PHI score as prototype for the tradeoff matrix.
Token-level is blocked on sentence-level validation. Do sentence-level first.
Extend sentence-level direction to token-level:
PHI_score(token_i) = activation_at_token_i · PHI_directionEach token in a clinical sentence gets a continuous PHI score. Validate against ground-truth NER labels.
Train a second direction vector on clinically annotated data:
clinical_score(token_i) = activation_at_token_i · clinical_directionThis enables the privacy–utility tradeoff matrix:
PHI_score high + clinical_score low → redact safely
PHI_score high + clinical_score high → flag for human review
PHI_score low + clinical_score high → preserve
Test whether steering along L03 + L16 directions simultaneously produces cleaner PHI suppression in generated output than single-direction steering.
intervention = α * L03_direction + β * L16_directionActivation patching and head ablation within L03 and L16 to identify the minimal circuit responsible for PHI encoding.
Validate direction vectors and probes on Hindi clinical text and English-Hindi code-switched patient messages (Indian clinical deployment context).
meta-llama/Llama-3.2-1B-Instruct — chosen for interpretability (smaller, less RLHF compression).
Base model comparison is a planned ablation.
600 contrastive sentence pairs generated from HealthCareMagic-100k:
- Patient messages extracted and cleaned
- Names inserted naturally using a 120B clinical LLM with varied position/pattern prompts
- Three complexity buckets: simple (<20w), compound (20–40w), complex (>40w)
- 200 pairs per bucket → 1200 labeled sentences total (600 PHI, 600 non-PHI)
Each pair:
{
"original": "My son James has a rash on his arms since last week.",
"redacted": "My son has a rash on his arms since last week.",
"bucket": 1,
"name": "James",
"identifiers": {
"Name": {
"word_level": [{"position": 2, "value": "James"}],
"token_level": [{"position": 4, "value": " James"}]
}
}
}Last-token hidden states extracted at all 17 layers (embedding + 16 transformer layers).
Left-padded tokenization → index [-1] always the last real token.
Output shape: [1200, 17, 2048]
Logistic regression, StandardScaler, 5-fold stratified CV. Metrics: Accuracy, F1, AUC-ROC, Precision, Recall, MCC.
| Configuration | Strategy | AUC | F1 | Recall | MCC |
|---|---|---|---|---|---|
| L03 (baseline) | single probe | 0.9792 | 0.9352 | 0.9400 | 0.8701 |
| L03 + L16 | weighted vote | 0.9850 | 0.9565 | 0.9533 | 0.9134 |
| Top-3 AUC (L03+L02+L16) | weighted vote | 0.9876 | 0.9480 | 0.9417 | 0.8969 |
| Top-3 AUC (L03+L02+L16) | concat | 0.9863 | 0.9520 | 0.9450 | 0.9054 |
| All late (L13–L16) | weighted vote | 0.9657 | 0.8988 | 0.8817 | 0.8025 |
Recommended deployment: L03 + L16 weighted vote — highest recall and MCC.
.
├── data_generation/
│ ├── oneshot.py # GPT-OSS 120B inference wrapper
│ └── generate_data.py # Contrastive pair generation pipeline
│
├── extract_activations.py # Forward pass → last-token hidden states → .pt
├── probe_layers.py # Per-layer logistic regression probe + metrics + plots
├── probe_ensemble.py # Weighted vote + concat ensemble strategies
├── visualize_encoding.py # PCA scatter, direction histograms, heatmap, bucket breakdown
│
└── probe_results/
├── layer_metrics.json
├── layer_metrics.png
├── ensemble_metrics.json
├── ensemble_comparison.png
└── viz/
├── pca_scatter.png
├── phi_direction_histograms.png
├── direction_similarity_heatmap.png
└── bucket_breakdown.png
pip install torch transformers scikit-learn matplotlib faker datasetsRun in order:
python generate_data.py # generate contrastive pairs
python extract_activations.py # extract layer activations
python probe_layers.py # per-layer probe analysis
python probe_ensemble.py # ensemble comparison
python visualize_encoding.py # visualize what layers are encodingThis work is part of a broader research program on interpretability-informed clinical NLP, conducted at IIT Delhi. The long-term goal is a mechanistically grounded, privacy-utility tradeoff-aware anonymization system for Indian clinical settings — supporting English, Hindi, and code-switched text.