CoFrGeNet: Continued Fraction Architectures for Language Generation

TL;DR

This work introduces CoFrGeNet, a family of continued fraction-inspired architectures that replace selective Transformer components to achieve parameter and training-time efficiency in language generation. By formulating ladders via continuants, the authors derive gradient expressions that reduce divisions to a single division per ladder and implement a custom PyTorch pathway to accelerate training and inference. The combination of attention and FFN replacements, a dyadic training schedule, and pretraining on diverse corpora yields competitive or superior results with substantially fewer parameters and faster runtimes, demonstrated on GPT2-xl and Llama-3.2B. The approach holds promise for hardware-tailored implementations and broader accessibility of high-quality generative models, though safety considerations and further optimizations remain for future work.

Abstract

Transformers are arguably the preferred architecture for language generation. In this paper, inspired by continued fractions, we introduce a new function class for generative modeling. The architecture family implementing this function class is named CoFrGeNets - Continued Fraction Generative Networks. We design novel architectural components based on this function class that can replace Multi-head Attention and Feed-Forward Networks in Transformer blocks while requiring much fewer parameters. We derive custom gradient formulations to optimize the proposed components more accurately and efficiently than using standard PyTorch-based gradients. Our components are a plug-in replacement requiring little change in training or inference procedures that have already been put in place for Transformer-based models thus making our approach easy to incorporate in large industrial workflows. We experiment on two very different transformer architectures GPT2-xl (1.5B) and Llama3 (3.2B), where the former we pre-train on OpenWebText and GneissWeb, while the latter we pre-train on the docling data mix which consists of nine different datasets. Results show that the performance on downstream classification, Q\& A, reasoning and text understanding tasks of our models is competitive and sometimes even superior to the original models with $\frac{2}{3}$ to $\frac{1}{2}$ the parameters and shorter pre-training time. We believe that future implementations customized to hardware will further bring out the true potential of our architectures.

Figures (13)

• Figure 1: Above we see a Transformer block consisting of attention and FFN layers. We propose candidate CoFrNet architectures for Transformer (causal) attention and FFN layers. The circles with the blue curves denote the $\frac{1}{x}$ non-linearity in our architectures. The zoomed out image on the far right shows the mapping between the pictorial representation and the actual equations. Details of the architectures are discussed in section \ref{['sec:meth']}.
• Figure 2: Two CoFrNet architectures to simulate attention a.k.a. causal token-token mixing. For the left architecture (CAttnU) a transpose is taken of the dimension $\times$ sequence length part of the input tensor and the output is transposed back to make it consistent with the later layers. The transpose makes the tokens mix, while upper triangular connections in the second to last layer in the architecture as well as the restricted structure of the ladders make sure information is only shared from previous tokens to following tokens and not bi-directionally (a.k.a. causal sharing). It consists of two ensembles of univariate CoFrNet ladders each of which then have an upper triangular linear layer on top. The representations formed are then element wise multiplied to form the final representation. The element wise multiplication produces interaction terms that otherwise would not occur, significantly enhancing representation power without compromising the causal information flow. The right architecture (CAttnM) we do not transpose the input. We use $L$ CoFrNet ladders that get mapped to a sequence length size embedding which corresponds to attention weights for that token. To maintain causality attention weights are computed only over the prior tokens. These then like in standard attention are used to weight the embeddings in the (value) $V$ matrix.
• Figure 3: CoFrNet architecture to simulate FFNs -- Cffn -- in a transformer block. We create a gated non-expanded (i.e. $\alpha=1$) representation that we pass to the CoFrNet ladders. Here no transpose is taken and hence feature mixing in either direction does not interfere with causal generation which is why we have a linear layer on top. Again the collapsed implementation is described in section \ref{['sec:contimpl']}.
• Figure 4: Architecture for implementing a linear combination of CoFrNet ladders (CF stands for continued fraction).
• Figure 5: GPT2-xl example generation when pre-trained on OWT.

Theorems & Definitions (4)

• Proposition 1
• proof
• proof
• proof : Proof of Lemma 3