Per-Head Muon prevents attention heads with different momentum scales from sharing one coupled orthogonalization. For parameter matrix Θ, gradient Gt, previous momentum Mt−1, and coefficient μ, form
Mt=μMt−1+Gt.Partition the output rows of Mt into equal attention-head blocks. For compact singular value decomposition Mt(h)=UhΣhVhT, use the polar factor
Oh=UhVhT.Concatenate the head factors in their original row order and update with learning rate η as Θt+1=Θt−ηO. Return a dictionary with updated_parameter, updated_momentum, and orthogonalized_update. Every tensor preserves the parameter shape, dtype, and device.
Examples
Input: parameter = [[2]], gradient = [[3]], previous_momentum = [[1]], num_heads = 1, momentum_coefficient = 0.5, learning_rate = 0.1
Output: {"updated_parameter": tensor([[1.9]]), "updated_momentum": tensor([[3.5]]), "orthogonalized_update": tensor([[1.0]])}
Explanation: The new momentum is 3.5, whose scalar polar factor is 1. A learning-rate step of 0.1 changes the parameter from 2 to 1.9.
Input: parameter.shape = gradient.shape = previous_momentum.shape = (4, 2), num_heads = 2, momentum_coefficient = 0.8, learning_rate = 0.05
Output: {"updated_parameter": tensor of shape (4, 2), "updated_momentum": tensor of shape (4, 2), "orthogonalized_update": tensor of shape (4, 2)}
Input: parameter.shape = gradient.shape = previous_momentum.shape = (2, 4), num_heads = 2, momentum_coefficient = 0.9, learning_rate = 0.02
Output: {"updated_parameter": tensor of shape (2, 4), "updated_momentum": tensor of shape (2, 4), "orthogonalized_update": tensor of shape (2, 4)}
Hint 1
Split the momentum matrix into equal contiguous row blocks, one block per attention head.
Hint 2
For each block, multiply the left and right singular-vector matrices from a compact singular value decomposition.
Hint 3
Concatenate the block updates in their original row order before updating the parameter.
Requirements
- Return a dictionary with exactly updated_parameter, updated_momentum, and orthogonalized_update.
- Every returned tensor has the same shape, dtype, and device as the parameter.
- Orthogonalization is independent across output-row head blocks.
- Do not mutate the parameter, gradient, or previous momentum.
Constraints
- Parameter, gradient, and momentum are finite matrices with identical shapes.
- The output-row count is divisible by the positive head count.
- Momentum coefficient lies in [0, 1).
- Learning rate is positive.
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Accepts: number
Accepts: number
Per-Head Muon
Per-Head Muon prevents attention heads with different momentum scales from sharing one coupled orthogonalization. For parameter matrix Θ, gradient Gt, previous momentum Mt−1, and coefficient μ, form
Mt=μMt−1+Gt.Partition the output rows of Mt into equal attention-head blocks. For compact singular value decomposition Mt(h)=UhΣhVhT, use the polar factor
Oh=UhVhT.Concatenate the head factors in their original row order and update with learning rate η as Θt+1=Θt−ηO. Return a dictionary with updated_parameter, updated_momentum, and orthogonalized_update. Every tensor preserves the parameter shape, dtype, and device.
Examples
Input: parameter = [[2]], gradient = [[3]], previous_momentum = [[1]], num_heads = 1, momentum_coefficient = 0.5, learning_rate = 0.1
Output: {"updated_parameter": tensor([[1.9]]), "updated_momentum": tensor([[3.5]]), "orthogonalized_update": tensor([[1.0]])}
Explanation: The new momentum is 3.5, whose scalar polar factor is 1. A learning-rate step of 0.1 changes the parameter from 2 to 1.9.
Input: parameter.shape = gradient.shape = previous_momentum.shape = (4, 2), num_heads = 2, momentum_coefficient = 0.8, learning_rate = 0.05
Output: {"updated_parameter": tensor of shape (4, 2), "updated_momentum": tensor of shape (4, 2), "orthogonalized_update": tensor of shape (4, 2)}
Input: parameter.shape = gradient.shape = previous_momentum.shape = (2, 4), num_heads = 2, momentum_coefficient = 0.9, learning_rate = 0.02
Output: {"updated_parameter": tensor of shape (2, 4), "updated_momentum": tensor of shape (2, 4), "orthogonalized_update": tensor of shape (2, 4)}
Hint 1
Split the momentum matrix into equal contiguous row blocks, one block per attention head.
Hint 2
For each block, multiply the left and right singular-vector matrices from a compact singular value decomposition.
Hint 3
Concatenate the block updates in their original row order before updating the parameter.
Requirements
- Return a dictionary with exactly updated_parameter, updated_momentum, and orthogonalized_update.
- Every returned tensor has the same shape, dtype, and device as the parameter.
- Orthogonalization is independent across output-row head blocks.
- Do not mutate the parameter, gradient, or previous momentum.
Constraints
- Parameter, gradient, and momentum are finite matrices with identical shapes.
- The output-row count is divisible by the positive head count.
- Momentum coefficient lies in [0, 1).
- Learning rate is positive.
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Accepts: number
Accepts: number