feat(init): coupled component initialization (weight_init: kaiming | coupled)

[Antovigo]

Summary

Adds the coupled component initialization to the JAX trainer as pd.weight_init: kaiming | coupled (kaiming stays the default).

With the "coupled" initialization scheme:

• For matrices with d_in ≤ d_out, the U vectors lie in the column space of the matrix. For matrices with d_in > d_out, the V vectors lie in the row space.
• The two vectors are coupled, such that one is the image of the other according to W

Concretely:

• If a matrix’s d_in ≤ d_out, draw V as random vectors normalized to unit norm, and set U = (W @ V)ᵀ
• If d_in > d_out, draw unit-norm U and set V = (Wᵀ @ Uᵀ).

Testing

This was tested on the targeted decomposition of Llama-8B L18 on arithmetic tasks, with 3 different random seeds per condition.

Black is normal kaiming init, red is the coupled init introduced in this PR. Green initializes the components within the row/column space of the matrices without coupling. It outperforms the baseline, but is not as good as the coupled init, so I didn’t include it in the PR.

In practice this saves about 1000 steps on the arithmetic task, i.e. 5% of training.

Validation on Layer 17 for the same task (evals taken directly at init):

Validation on a full-data decomposition: initialization + 100 steps of faithfulness warmup + 200 steps of training for the "Jose" decomposition (Pile llama 4L):

The coupled initialization is also in the right direction for full-data decomposition, but it’s less clear whether this changes anything in practice. On the other hand, it runs only once and has no practical cost, so it probably doesn’t hurt to use it.

Note that these tests where done using the old Pytorch implementation, which should behave the same as the Jax version.

For the new Jax implementation:

• New param_decomp/tests/test_component_init.py: the defining W-image property on both orientation branches (unit-norm seeds, span membership), and independent per-site draws within a stacked shape group.
• TMS 5-2 engine smokes run end-to-end for both weight_init modes (init → faith warmup → training → eval → checkpoint)