PLDR-LLM-v51-DAG-110M

Model Description

PLDR-LLM-v51-DAG-110M is a large language model from power law decoder representations with KV-cache and G-cache support, which is a new foundational language model architecture that utilizes power law graph attention to generate deductive and inductive outputs. This model has a parameter size of 110M. It refers to PLDRv51-DAG-110M whose architecture and training details are provided in Table 1 of the research paper titled PLDR-LLMs Learn A Generalizable Tensor Operator That Can Replace Its Own Deep Neural Net At Inference.

Training data

PLDR-LLM-v51-DAG-110M was pretrained on the RefinedWeb, a publicly available English web dataset with extensive filtering and deduplication.

Training procedure

This model was trained for ~8B tokens on RefinedWeb over 250k steps per rank. It was trained autoregressively with cross-entropy loss.

Intended Use and Limitations

This model is intended to be used for research purposes. Given text as input prompt, it carries out next token prediction to generate continuation text. The context length for this model is 1024 tokens.

How to Use

Via Huggingface Transformers Library

PLDR-LLM has custom model support for Huggingface Transformers library. PLDR-LLM with custom code is evaluated on Transformers 4.56.1 available at the time.

Using pipeline:

from transformers import pipeline

text_generator = pipeline(
                          task="text-generation",
                          model="fromthesky/PLDR-LLM-v51-DAG-110M",
                          device="cuda", # or "cpu"
                          trust_remote_code=True
                         )

prompt="The quick brown fox jumps over the lazy dog."

output=text_generator(prompt, top_p=0.6, top_k=0, temperature=1, do_sample=True,               
                      tokenizer_encode_kwargs={"add_special_tokens":False}, 
                      use_cache=True, max_new_tokens=100)
print(output[0]["generated_text"])

Using AutoModel:

from transformers import AutoModelForCausalLM, AutoTokenizer
device="cuda" # or "cpu"
model=AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path="fromthesky/PLDR-LLM-v51-DAG-110M",
                                           device_map=device,
                                           trust_remote_code=True
                                          )
tokenizer=AutoTokenizer.from_pretrained(pretrained_model_name_or_path="fromthesky/PLDR-LLM-v51-DAG-110M",
                                        add_eos_token=False,
                                        legacy=False,
                                        trust_remote_code=True
                                       )

prompt="The quick brown fox jumps over the lazy dog."
        
inputs = tokenizer([prompt], return_tensors="pt").to(device=device)
generated_ids = model.generate(**inputs,
                               max_new_tokens=100, 
                               top_p=0.6,
                               top_k=0, 
                               temperature=1, 
                               do_sample=True,
                               use_cache=True
                              )
print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))

PLDR-LLM specific configurations:

• custom_G_type: None for learned G values during pretraining, 'identity' for LLM with SDPA equivalent, 'random' for G values from a random normal distribution, 'external' for custom G values that can be assigned after model initialization. This setting is more important for training purposes, for inference it is set in the model config.json file.

• cache_first_G: For batched inference, if set to True, cache G values from the first sample prompt in batch for all samples. If set to False, cache G values separately for each sample prompts in batch. For contrastive generation with custom_G_value=None, this needs to be set to True.

• reference_rope: If set to True, RoPE implementation implemented in the original paper is used. This is the case for model pretrained in this repo. If set to False, RoPE implementation from the Huggingface Transformers library is used.

• output_pldr_attentions=True returns the deductive outputs and learnable parameters of power law graph attention module as tuple containing: the output of the residual metric learner (metric tensor, A), output (A_LM) after application of iSwiGLU on metric tensor, learned exponents of potential tensor, learned weights for energy-curvature tensor, learned bias for energy-curvature tensor, energy-curvature tensor (G_LM), and attention weights.

See config.json for other model configuration details.

Notes:

• This implementation of PLDR-LLM custom code was evaluated on Transformers 4.56.1 and pytorch 2.6.0.
• We also have a fork of transformers library with PLDR-LLM model support for future development. The PLDR-LLM model files are added to the library so custom model files are not necessary.

      git clone https://github.com/burcgokden/transformers
      cd transformers
      git checkout add_PLDR_LLM
      pip install -e ".[dev]"

• Static cache is not supported for models with custom_G_type=None.
• PLDR-LLM uses EOS token "[END]" during pretraining to indicate end of a sequence. For text generation, we do not need to add the EOS token to the prompt. To achieve this, add_eos_token=False can be set in tokenizer_config.json file or while initializing the tokenizer model. For text generation pipeline call method, tokenizer_encode_kwargs={"add_special_tokens":False} can be used.
• When add_bos_token=False and add_eos_token=False are set for the tokenizer model, prompt "" is an invalid input for single batch inference as it doesn't contain any tokens. When padding is enabled, batched inference with prompt "" as one of the samples causes its input_ids to be pad tokens and attention_mask to be all zeros. This edge case is handled differently for _attn_implementation='eager' and 'sdpa', resulting in different generation outputs for this prompt. Setting add_bos_token=True, add_eos_token=True or explicitly providing prompt as "[PAD]", "[START]", or "[END]" gives same output for either implementation. This issue does not affect KV-cache and G-cache.

Via Original Implementation

• The original model implementation files can be found in the folder named paper_saved_model_files/. The model checkpoint and tokenizer can be loaded into the PLDR-LLM framework to generate text as described in the code repository for training this model: PLDR-LLM-with-KVG-cache.

LM Evaluation Harness Support

• The model can be used with a fork of LM-Evaluation-Harness Suite with PLDR-LLM with KV-cache and G-cache support: lm-evaluation-harness-with-PLDR-LLM-kvg-cache.

Limitations and Biases

Large Language Models may generate text that is profane, lewd, socially unacceptable or offensive based on the contents of the dataset it was pretrained. RefinedWeb is a dataset that is as toxic and biased as the Pile. Please see the papers for RefinedWeb and the Pile for more information. Moreover, large language models are also susceptible to hallucinations and may generate text that contains incorrect, irrelevant or misleading information. Since it is very hard to expect the contents of generated text ahead of time, the output of the large language models need to be heavily moderated and curated to avoid undesired content to appear without warning.

Eval results

• The evaluation results on benchmarks with zero-shot setting and their comparison to LLM models of similar size reported in the literature can be found in Tables 3-5 and 7 of the research paper.
• For implementation via huggingface transformers library, evaluating on the same benchmark suite gives same results as in the paper for all benchmarks.

BibTeX entry and citation info

Please cite this model as:

@misc{gokden2025pldrllmkvgcache,
      title={PLDR-LLMs Learn A Generalizable Tensor Operator That Can Replace Its Own Deep Neural Net At Inference}, 
      author={Burc Gokden},
      year={2025},
      eprint={2502.13502},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2502.13502}, 
}