LASP-2: Rethinking Sequence Parallelism for Linear Attention and its Hybrid

Several key parts need to be clarified.

1. A contribution of this paper is extending LASP-2 to a hybrid model that integrates linear and standard attention layers. However, the description On Standard Attention Module on page 7 is unclear. How is the proposed all-gather collective communication used on standard attention that includes softmax attention? It seems that the same communication is not valid since the softmax attention hinders the right-product-first trick. Could the authors explicitly write down equations describing this part?

2. How are the 1/4 hybrid models constructed? Are the first 1/4 layers linear modules? Could the authors provide an ablation study on the ratio (1/4 here) as their claim lies in the performance of the hybrid model?

3. In LASP-2 with masking, there is a discrepancy in attention weights between inter-chunks and intra-chunks. Is there any possibility that the intra-chunks are pronounced more than inter-chunks as the weights in intra-chunks are always positive?

4. The LASP paper claims that LASP scales sequence length up to 4096K on 128 GPUs. Does LASP-2 give better scalability than SASP in that setting?

5. To further emphasize the novelty of LASP-2, I recommend describing LASP and the superior properties of LASP-2 in detail with exact formulations.

I list the main weaknesses below:

• In my opinion, the structure of the paper can be reworked. Namely, the section «Related work» should be in the beginning.
• The differences of LASP-1 and LASP-2 (section 3.3) look rather marginal. Could you please elaborate more on this?
• The proposed method is validated only by experiments on one dataset. There are no theoretical guarantees, thus it is difficult to judge how good the method is. I would suggest either add some theoretical results or extend the experimental part.

Overall, I'm reasonably happy with the paper but I think the paper would be stronger with two additions (adding a theoretical analysis and accuracy on other benchmarks) and by addressing bunch of other smaller comments.

Theoretical Analysis

While the experiments do show an improvement over LASP (and other parallel implementation like Ringed Attention), I do not have a sense of what is the theoretically maximum gain that could be achieved with LASP-2? I think the paper would be much stronger if there were a theoretical analysis to show what would be the ideal improvement (in the limit of $N\to\infty$) of LASP-2 over LASP and Ringed Attention?

The reason I as this is because I'm not sure if the proposed method's improvement is a constant factor (like $2\times$) or is is super-constant.

I understand that matching the theoretical gains in experiments is hard (and I'm not expecting this) but I want to get a sense of what is the best-case scenario improvement with the proposed method.

Moving beyond one benchmark

The experimental results on performance was only done on the SlimPajama benchmark. It is hard to calibrate how good a system is if comparison is made on only one benchmark. It would make the paper stronger if there were comparisons done on other benchmarks, ideally with different language modeling tasks.

(Relatively) Minor Weaknesses

All the comments below are more along the lines of asking clarifications on some of the experimental choices made in the paper and the rest are presentation related. I'm listing these in order of when they appear in the paper and not in the order of their importance:

1. [Line 135] This is a nitpick: the $k_i$ and $v_i$ should be $\mathbf{k}_i$ and $\mathbf{v}_i$.
2. [Lines 142-143] Could you please expand on what it means that ignoring the normalizing denominator "works effectively in practice"? My understanding was that the normalization was important even for linear attention models.
3. [Line 377] Combing smaller documents into one seems like would lead to worse results accuracy-wise?
4. [Line 430] Could you please elaborate exactly which parts of Based were used? Specifically Based a Taylor series expansion that approximates Attention as well as an exact sliding window standard Attention. I'm guessing the experiments in the paper used both of these parts of Based but I just wanted to double-check.
5. [Fig 3+4] Do you get similar results when replacing basic linear attention with other linear attention variants considered in the paper?

The claim here is that relative to LASP of Sun et al. 2024, the proposed method, coined LASP-2, offers better communication efficiency since it is the state matrix that gets transferred whose dimension does not depend on the sequence or chunk length, as opposed to K and V matrices that depend on the chunk length. But this is not true. The communication requirement is the same with the original LASP algorithm that communicates the KV activation.

Likewise, the claim above and below eq (5) about concurrent computations for each chunk being unique to LASP-2 is unfounded. It is easy to see that identical operations are already done in LASP.