Implement causal softmax-free linear attention in parallel and recurrent forms. At position t, qt, kt, and vt are the query, key, and value vectors; St is the state and yt is the output. Superscript T denotes transpose.
St=St−1+ktvtT. yt=qtTSt.The state starts at zero, and each output includes the current token. Compute the sequence outputs both in parallel and by a recurrent scan. Return a dictionary containing parallel_output and recurrent_output, each with shape (B, S, Dv), and final_state with shape (B, Dk, Dv). All three values are PyTorch tensors with the same dtype and device as q.
Input: q = [[[1],[4]]], k = [[[2],[5]]], v = [[[3],[6]]]
Output: {"parallel_output":[[[6],[144]]],"recurrent_output":[[[6],[144]]],"final_state":[[[36]]]}
Explanation: The state is 6 after the first write and 36 after the second. Reading those states with queries 1 and 4 produces 6 and 144 in both formulations.
Input: q = [[[1,0],[0,1]]], k = [[[1,0],[0,1]]], v = [[[2],[3]]]
Output: {"parallel_output":[[[2],[3]]],"recurrent_output":[[[2],[3]]],"final_state":[[[2],[3]]]}
Input: q = [[[1,0]],[[0,1]]], k = [[[1,0]],[[0,1]]], v = [[[2,0]],[[0,5]]]
Output: {"parallel_output":[[[2,0]],[[0,5]]],"recurrent_output":[[[2,0]],[[0,5]]],"final_state":[[[2,0],[0,0]],[[0,0],[0,5]]]}
torch.cumsum computes prefix sums along a chosen dimension.
torch.bmm multiplies a batch of matrices without mixing batch entries.
Accepts: array
Accepts: array
Accepts: array
Implement causal softmax-free linear attention in parallel and recurrent forms. At position t, qt, kt, and vt are the query, key, and value vectors; St is the state and yt is the output. Superscript T denotes transpose.
St=St−1+ktvtT. yt=qtTSt.The state starts at zero, and each output includes the current token. Compute the sequence outputs both in parallel and by a recurrent scan. Return a dictionary containing parallel_output and recurrent_output, each with shape (B, S, Dv), and final_state with shape (B, Dk, Dv). All three values are PyTorch tensors with the same dtype and device as q.
Input: q = [[[1],[4]]], k = [[[2],[5]]], v = [[[3],[6]]]
Output: {"parallel_output":[[[6],[144]]],"recurrent_output":[[[6],[144]]],"final_state":[[[36]]]}
Explanation: The state is 6 after the first write and 36 after the second. Reading those states with queries 1 and 4 produces 6 and 144 in both formulations.
Input: q = [[[1,0],[0,1]]], k = [[[1,0],[0,1]]], v = [[[2],[3]]]
Output: {"parallel_output":[[[2],[3]]],"recurrent_output":[[[2],[3]]],"final_state":[[[2],[3]]]}
Input: q = [[[1,0]],[[0,1]]], k = [[[1,0]],[[0,1]]], v = [[[2,0]],[[0,5]]]
Output: {"parallel_output":[[[2,0]],[[0,5]]],"recurrent_output":[[[2,0]],[[0,5]]],"final_state":[[[2,0],[0,0]],[[0,0],[0,5]]]}
torch.cumsum computes prefix sums along a chosen dimension.
torch.bmm multiplies a batch of matrices without mixing batch entries.
Accepts: array
Accepts: array
Accepts: array
