Gradient atoms without LoRA

Gradient Atoms (J Rosser, 2026) discovers model behaviors unsupervised: compute one training gradient per fine-tuning document, sparse-code the collection into a dictionary of "atoms", and steer the model by applying atoms as weight-space perturbations. Their pipeline reads the gradients through two lossy windows — a rank-8 LoRA adapter and a truncated EKFAC eigenbasis. This piece asks whether either window is load-bearing: we rebuild the experiment on full-model gradients (3.2B parameters), preconditioned by a 6.9 MB AdaFactor-style second-moment estimate and compressed with CountSketch hashing instead. Everything replicates, several things improve, and the whole run costs about $16.

The LoRA + EKFAC machinery wasn't load-bearing: full-gradient atoms recover the same behavior inventory and steer all five tested behaviors to 0%, on one GPU.

§1Why bother

Two structural worries about the original. First, with standard LoRA (ΔW = BA, B zero-initialised), per-layer gradients reach the adapter only through row-space(A) and col-space(B) — and the LoRA-asymmetry literature says A stays close to its random init, so half the gradient signal is a random sketch whose quality depends on the draw. Second, gradient-kernel's own ablations found that full EKFAC whitening inflates the gradient kernel's stable rank ~4.7 → ~620, making it nearly unsketchable and collapsing judge-rated cluster quality to chance — the diagonal (Adam/AdaFactor) family is the validated operating point for this kind of unsupervised decomposition. So the substitution is: LoRA → full weights + unbiased hash sketch; EKFAC → rank-1 factored second moment.

§2Pipeline

Base model and data match the paper: Gemma-3-4B-IT and their exact 5,000-document general SFT sample (same seed, same shuffle). We full-fine-tune the language stack (attention + MLP; embeddings and vision frozen), then treat that checkpoint's per-document gradients as the object of study:

1. Precondition. Estimate the per-coordinate second moment v from 2,048 token gradients (fp64 on-GPU accumulation, 13 min), then factor it AdaFactor-style: v̂ = r c^T/Σv per weight matrix. The factors are 6.9 MB (vs 12.8 GB for the raw diagonal) and agree with it for 85.5% of coordinates within 2× — nearly identical to the 70M-scale A/B (87.7%).
2. Sketch. One gradient row per document (assistant tokens, sum reduction), scaled by v̂^−1/2 and hashed by CountSketch into 4×8,192 = 32,768 dims. 5,000 rows ≈ 40 min, 330 MB — vs ~77 TB if you materialised the raw gradients.
3. Decompose. The paper's own dictionary-learning recipe, unchanged (MiniBatchDictionaryLearning, K=500, α=0.1), on unit-normalised sketch rows.
4. Steer. For an atom of interest, average the exact recomputed gradients of its top-20 activating documents (never unsketch — S^T would smear hash noise over 3.2B coordinates) and apply θ ± αv directly to the weights.

§3Discovery

The inventory looks like the paper's: the highest-coherence atoms are stereotyped task types, not topics. Sentiment classification (#195), spam detection (#180, #267), true/false answering (#35), refusal on missing input (#292, #7 — top docs all "please provide the document/text/article…"), yes-answering (#374, #469), Python code (#47, #45), fiction/nonfiction genre labels (#343, #382), translation (#224), one-word factual QA (#296). A detector-based matcher (score each atom's top-20 documents with the paper's five behavior regexes) finds an atom for all five of their steering behaviors at hit-fraction 1.00 — no labels or per-query scoring anywhere upstream of that.

§4Steering

All five behaviors are suppressible to exactly 0% (the paper got four of five there; their code atom bottomed at 28%). Amplification matches where headroom exists — yes/no +13pp vs their +12pp — and the two idiosyncrasies of their results reproduce cleanly: dictionary signs are arbitrary (our code atom steers with the "away" sign, exactly like their #64), and numbered lists suppress easily but refuse to amplify. The steered behavior is qualitative, not just regex-deep: asked to "Summarize the article." with no article, the refusal-suppressed model fluently makes up a plausible summary instead of asking for the input. And one finding the detector alone would have mislabelled: the bullets atom pushed "toward" at larger ρ doesn't remove bullets — it compresses answers to a single bullet line. That atom seems to encode list length, not list formatting.

§5Beyond the paper's five: the top-coherence atoms

The five behaviors above were chosen for comparability with the paper. The thirteen highest-coherence atoms mostly lack ready-made regex detectors, so they got the paper's own qualitative protocol instead: steer ± at ρ ∈ {3×10^-6, 10^-5} and diff responses on twenty generic questions. Read "changed" as direction-specific response change on out-of-distribution questions, not steering success:

At the same ρ that swung the five detector-matched behaviors by ±50pp, every one of these atoms leaves 90–95% of generic responses byte-identical, with zero degeneration. That reads as a feature: the high-coherence atoms are narrow — spam-classification or true/false formatting simply isn't invited by "What's the best pet?", so a targeted edit shows no collateral drift off-distribution. The contrast case is instructive: refusal-family atom #292 (coherence 0.542) changes 10% of generic responses, while its low-coherence sibling #7 (0.104) drove the 100%/0% refusal swings — on refusal-inviting questions. Demonstrating steering for the remaining task-type atoms needs behavior-inviting evaluation sets, which is the same lesson the paper's five hand-built detectors embody. What the sweep does establish for all thirteen: fluency-preserving, distribution-local edits.

§6Caveats & what's next

• Coherence numbers aren't strictly commensurable with the paper's (sketched preconditioned full-gradient cosines vs raw LoRA-gradient cosines); the honest claim is "at least as much structure".
• Our fine-tune is full-parameter with prompt masking; theirs was LoRA with packing — so atom-for-atom inventory comparison is qualitative. The clean head-to-head is running this pipeline on J's exact merged adapter checkpoint, so only the gradient representation differs.
• Per-token atoms (~1.25M rows, ~$200–500) are the regime the LoRA+EKFAC pipeline can't reach and the natural next run.
• One seed, one model, regex detectors, five quantitative behaviors.