[ARCH-PROP] CoLA: Compute-Efficient Pre-Training of LLMs via Low-Rank Activation

### Architecture Name

CoLA

### Parent issue

#1

### Motivations

Dense Transformer models spend most of their parameters and matrix-multiplication compute in attention projections and feed-forward layers. However, intermediate activations in pretrained language models often exhibit an effective rank substantially below their full feature dimension.

CoLA (Liu et al., EMNLP 2025) addresses this redundancy through an architectural change rather than a parameter- efficient fine-tuning method. It replaces full-size linear projections with nonlinear low-rank bottlenecks. Unlike conventional linear factorization, the nonlinearity between the two factors improves representational capacity while reducing parameters and computation.

### Proposed Architecture

For a standard linear layer \(h=Wx\), CoLA replaces the full-rank matrix \(W\) with two smaller matrices and an intermediate nonlinear activation:

$$C_W(x)=B_W\mathrm{SiLU}(A_Wx),$$

where

$$A_W\in\mathbb{R}^{r\times d_{\mathrm{in}}},\qquad B_W\in\mathbb{R}^{d_{\mathrm{out}}\times r},\qquad r<\min(d_{\mathrm{in}},d_{\mathrm{out}}).$$

#### Attention

The query, key, value, and output projections are replaced with CoLA layers:

$$Q=C_Q(x),\qquad K=C_K(x),\qquad V=C_V(x).$$

#### FFN

The gate, up, and down projections in the gated MLP are also replaced with CoLA layers:

$$g=C_{\mathrm{gate}}(x),\qquad u=C_{\mathrm{up}}(x).$$

The main CoLA MLP is

$$\mathrm{MLP}_{\mathrm{CoLA}}(x)=C_{\mathrm{down}}(g\odot u).$$

Each CoLA layer already contains an internal SiLU activation. Following the original CoLA implementation, the main configuration therefore removes the additional full-width SwiGLU activation.

### Preliminary Results (if any)

https://aclanthology.org/2025.emnlp-main.230.pdf

### Experiments Plan

### Experimental Plan

The primary goal of this experiment is to reproduce CoLA on the OLMo architecture. We will implement CoLA layers in the OLMo training codebase and train the resulting model using the standard OLMo configuration. For a controlled comparison, the training data, tokenizer, model dimensions, token budget, optimizer, learning-rate schedule, batch size, sequence length, and evaluation settings will remain identical to those of the corresponding full-rank OLMo baseline.

The initial CoLA model will use the default bottleneck rank from the original paper:

$$r=\frac{d_{\mathrm{model}}}{4}.$$

The main objective is to determine whether the efficiency and model-quality results reported in the original CoLA paper can be reproduced under the OLMo architecture and training setup.