Technical Summary: CoFrGeNet – Continued Fraction Architectures for Language Generation

Problem Statement

Transformers have become the dominant architecture for language generation, relying on Multi-head Attention and Feed-Forward Networks (FFNs) with large hidden layers. However, these components introduce significant computational overhead, particularly in terms of parameter count and training time. While various approximations (e.g., sparse attention, linear attention) have been proposed to improve efficiency, there remains a need for fundamentally different function classes that can replace these core components while maintaining or improving performance.

Previous work introduced CoFrNets (Continued Fraction Networks) for supervised learning, utilizing a ladder-like structure inspired by continued fractions. However, adapting this architecture for generative modeling presents three specific challenges:

1. Multi-dimensional outputs: Generative models require producing high-dimensional outputs, unlike the scalar outputs often seen in supervised settings.
2. Richer function learning: The architecture must learn complex representations suitable for sequence modeling.
3. Causal constraints: The model must learn parameters that depend only on prior tokens to ensure coherent generation.
4. Computational inefficiency: The standard $1/x$ non-linearity in continued fractions requires $d$ divisions for a depth-$d$ ladder. Since division is significantly slower than multiplication on modern hardware, this creates a bottleneck for both training and inference, especially as depth ($d$) and the number of ladders ($L$) increase.

Methodology

1. Continued Fraction Generative Networks (CoFrGeNets)

The authors propose CoFrGeNets, a family of architectures that replace Attention and/or FFN layers in Transformer blocks with components based on continued fractions. The core function class is defined as:
$$w_0x + \frac{1}{w_1x + \frac{1}{w_2x + \dots}}$$
where $x$ is the input and $w_k$ are learnable parameter vectors.

2. Architectural Components

The paper introduces specific designs for replacing Transformer components while maintaining causal constraints:

• Causal Attention Replacement (CAttnU and CAttnM):

  • CAttnU: Uses univariate CoFrNet ladders where inputs are transposed to mix dimensions across tokens. To enforce causality (preventing future token leakage), the architecture employs upper triangular linear layers within the ensembles. Element-wise multiplication of outputs from two ensembles generates cross-terms, enhancing representation power without violating causal flow.
  • CAttnM: Uses $L$ CoFrNet ladders to map inputs to sequence-length embeddings, which serve as attention weights. Causality is maintained by applying a causal softmax (convex combination of only prior tokens) before weighting the value matrix.

• FFN Replacement (Cffn):

  • Since FFNs require feature mixing rather than token mixing, the architecture uses ensembles of p-variate ladders (where $p$ is the embedding dimension) without transposing the input.
  • A linear layer at the end projects the output back to the target dimension.
  • Crucially, the input to these ladders is a gated, non-expanded representation ($\alpha=1$), avoiding the parameter explosion typical of standard Transformer FFNs which often use an expansion factor $\alpha \gg 1$.

3. Efficient Implementation via Continuants

To address the computational inefficiency of repeated divisions, the authors derive a custom implementation based on continuants (polynomials $K_k$).

• Forward Pass: Instead of computing the fraction layer-by-layer with $d$ divisions, the continued fraction is computed as the ratio of two polynomials: $\tilde{f}(a) = K_{d-1} / K_d$. This requires computing the polynomials recursively using only multiplications and additions, followed by a single division ($1/K_d$).
• Backward Pass: The authors derive closed-form expressions for the gradients using the properties of continuants. The gradient $\frac{\partial \tilde{f}}{\partial a_k}$ is expressed as a ratio of continuants squared.
• Optimization: By computing $K_d$ once and reusing $1/K_d$ for all gradient calculations, the number of divisions is reduced from $d$ to 1. This significantly accelerates both training and inference.

4. Custom Training Schedule

The authors propose a dyadic parameter update schedule to stabilize training. Instead of updating all parameters simultaneously:

• Linear components are updated from the start.
• Parameters at higher depths are frozen initially.
• Updates for depth $i$ parameters begin only after $t/2^i$ iterations (where $t$ is the total iterations).
  This incremental approach prevents instability during the early phases of training.

Key Contributions

1. Novel Architectures: Introduction of CoFrGeNet components as plug-in replacements for Attention and FFN layers in Transformers, offering a new function class for generative modeling.
2. Efficient Gradient Formulation: Derivation of custom gradient formulas using continuants that reduce the number of division operations from $d$ to 1 per ladder, drastically improving hardware efficiency.
3. Training Stability: Development of a custom dyadic training schedule that stabilizes the training of deep continued fraction architectures.
4. Empirical Validation: Pre-training and evaluation on large-scale models (GPT2-xl and Llama3) demonstrating competitive performance with reduced parameter counts.

Experimental Results

Setup

• Models: Experiments were conducted on GPT2-xl (1.5B) pre-trained on OpenWebText and GneissWeb, and Llama3 (3.2B) pre-trained on the docling data mix (2T tokens).
• Baselines: Comparisons included the original Transformers, Dense Synthesizer, and Sparse Attention.
• Variants: CoFrGeNet-F (FFN replaced), CoFrGeNet-A (Attention replaced), and CoFrGeNet (both replaced).

Performance

• Downstream Tasks (GLUE): CoFrGeNet models, despite having 2/3 to 1/2 the parameters of the original models, achieved competitive or superior accuracy on classification, Q&A, and reasoning tasks. For instance, CoFrGeNet-F (985M) outperformed the 1.5B GPT2-xl on most GLUE benchmarks.
• Perplexity: On language modeling benchmarks (PTB, Wikitext, Lambada), CoFrGeNet variants generally achieved lower perplexity than the baseline Transformers and other efficient attention baselines.
• Efficiency:
  • Training Time: The continuant-based implementation reduced training time significantly compared to a "literal" implementation (CoFrGeNetB) that did not use continuants. For GPT2-xl, training time dropped from ~203 hours (B) to ~178 hours (CoFrGeNet).
  • Inference Time: Inference speed improved by nearly an order of magnitude for the continuant-based version compared to the literal implementation.
  • Throughput: On Llama3, CoFrGeNet variants showed higher tokens/day throughput and reduced total training time (e.g., CoFrGeNet-F trained in 6.6 days vs. 8.5 days for the baseline).

Significance and Claims

The paper claims that CoFrGeNets represent a viable, parameter-efficient alternative to standard Transformer components. By leveraging the mathematical properties of continued fractions and optimizing their implementation via continuants, the authors demonstrate that:

• High-performance language models can be built with significantly fewer parameters.
• The computational bottleneck of division operations in continued fractions can be effectively mitigated, making the architecture practical for large-scale industrial workflows.
• The approach is "plug-and-play," requiring minimal changes to existing training and inference pipelines.

The authors conclude that this work lays the groundwork for continued fraction-inspired generative architectures, potentially democratizing AI by enabling entities with fewer resources to pre-train high-quality models. They note, however, that these models do not inherently possess safety guards and remain susceptible to hallucinations and adversarial attacks, suggesting future research should focus on addressing these challenges within the specific functional form.