←新闻详情

Together AI Blog更新于 10月10日 11:24

Parcae:用稳定循环模型以更少参数实现更高性能

Parcae: Doing more with fewer parameters using stable looped models

AI 导读

Together AI 发布稳定循环语言模型架构 Parcae,通过反复使用相同层来增加计算量,而不扩大参数规模。实验中,Parcae 的验证困惑度较参数和数据匹配的 RDM 最多降低 6.3%,770M 模型的下游表现接近参数量两倍的 Transformer;训练代码和模型已发布。

推荐理由

精选AI 评分 73

材料给出稳定性分析、同参数与同数据对照及扩展规律,可用于评估层循环在内存受限部署中的参数效率。

主题

标签

正文

当前显示原文,中文全文译文尚未就绪。

Summary

We present Parcae, one of the first stable architectures for looped language models , achieving the quality of a Transformer twice the size with clean, predictable training. Parcae creates a new medium to scale quality by increasing recurrence rather than purely scaling data, opening up an efficient frontier for training memory-constrained on-device models.

Traditional scaling laws tell us that to achieve the best performance, we need to scale FLOPs, often with more parameters or data. But as models move to the edge and inference costs skyrocket, we wonder: Can we scale quality without inflating memory footprint?

To that end, we’ve been exploring looped architectures, models that increase compute by passing activations through the same layers multiple times. While promising, these models have been unstable to train. We tackle this issue directly and introduce Parcae , a stable looped architecture that:

As models move to the edge and inference deployments take on larger portions of compute, there is an increasing interest in scaling model quality without increasing parameters. One mechanism we have been excited about is layer looping, where initial works have trained looped models that match the quality of larger fixed-depth architectures.

To turn a vanilla Transformer into a looped model, we follow prior work and partition its layers into three functional blocks: a prelude ($\mathcal{P}$) , a recurrent ($\mathcal{R}$), and a coda ($\mathcal{C}$). The forward pass works in three stages:

Unfortunately, looped models are a headache to train [2] [3] [4] . We personally found them to suffer from residual state explosion and loss spikes. What makes looped models even trickier is that the recurrent block is composed of several vanilla Transformer blocks, making it difficult to reason about the source of instability.

While instability is a fickle foe, we observed that a simple linear framework captured a significant source of instability. Specifically, we recast looping as a nonlinear time variant dynamical system over the residual, whose update rule is:

$$h_{t+1} = \overline{A} h_t + \overline{B} e + \overline{\mathcal{R}}(h_t, e)$$

where $\overline{A}, \overline{B}$ perform injection and $\overline{\mathcal{R}}$ is the contribution of the Transformer blocks to the residual stream. For the subquadratic sequence mixing fanatics out there, observe that if we ignore the nonlinear term $\overline{\mathcal{R}}$, the resulting system is a discrete linear time-invariant (LTI) dynamical system over the residual state, across model depth.

What's cool is that for discrete LTI systems, their stability and convergence are determined by the eigenvalues of $\overline{A}$. Specifically, stability is categorized using the spectral norm $\rho(\overline{A})$ (i.e., the absolute largest eigenvalue of $\overline{A}$), with stable systems (convergent) being $\rho(\overline{A}) [2] , a prior looped model, we tested against parameter- and data-matched RDMs, observing that Parcae reduces validation perplexity by up to 6.3%.

When retrofitting a very strong Transformer baseline into an RDM, without any hyperparameter tuning, we found Parcae to be robust over RDMs (which just flat-out diverged).

We also took Parcae and used it as a drop-in replacement for a standard fixed-depth Transformer. Using a nanochat-inspired setup, we train a series of language models on FineWeb-Edu, up to 1.3B parameters. We found that Parcae outperformed all parameter- and data-matched Transformers, with our 770M Parcae model almost achieving downstream quality equivalent to a Transformer twice its size!

But is looping actually FLOP efficient? To study this, we explore a setting where, under a fixed parameter count and FLOP budget, we trade off mean recurrence in training with data (e.g., if we increase mean recurrence, we reduce training data to maintain a fixed FLOP budget).

At two scales, we find that increasing the mean recurrence used in training $mu_{\text{rec}}$ while proportionally reducing tokens yields lower validation loss than training with low recurrence and more data. What’s even cooler is that if we use a parabolic fit to extract the optimal $mu_{\text{rec}}$ and token budget at each FLOP level, we find that they both follow power laws with consistent exponents.

Alright, alright. But do we beat a fixed-depth model? Using our optimal recurrence scaling laws, we compare against fixed-depth Parcae models (i.e., those with $\mu_{\text{rec}}=1$) and looped Parcae models following our optimal mean recurrence prediction from our scaling laws. We found that looping creates a stricter Pareto Frontier for validation loss (figure below), which translates into better downstream quality (table below).

We are super excited about how far we can push parameter efficiency. With the growing costs of memory overhead during inference, we think there is a lot to explore in parameter reuse methods such as layer looping. To help accelerate this process, we are releasing training code and models . We aren’t done either; we have tons of new ideas to push looped models further, so stay tuned for what comes next!

If you have any questions or want to work with us on what comes next for Parcae, please reach out to Hayden Prairie at hprairie@ucsd.edu .

The name PaRCae is a homage to the three roman fates: Nona (the Prelude block $\mathcal{P}$), who initializes the computational thread of life , Decima (the Recurrent block $\mathcal{R}$), who measures the thread and evolving through model depth, and Morta (the Coda block $\mathcal{C}$), who finalizes the sequences by cutting the thread to produce the final output.

We would like to thank Together AI for collaborating with us and providing compute for these experiments. We would also like to thank Austin Silveria and Jonah Yi for their helpful feedback on this blog post.

来源

Together AI Blog