Maxwell-TCS-v0.2

Maxwell Task Complexity Scorer-v0.2

This repository contains the model from the PiSSA: Principal Singular Values and Singular Vectors Adaptation of Large Language Models paper.

Maxwell-TCS-v0.2 is a task/instruction complexity score annotator based on the ModernBERT-Large and trained on a large dataset of English instructions. It is designed to be fast, efficient and accurate for scoring/annotation related tasks.

Model Details

Model Description

Maxwell-TCS-v0.2 is an experimental SOTA task complexity scorer based on the state-of-the-art ModernBERT-Large backbone. It is trained under a data heavy setting (TTP ratio of ~3.5 toks/trainable-param) over 66.5K diverse instruction-score pairs, and performs at par with other complexity scorers 34 times larger in size. For a given user instruction, the model predicts normalized scores between 0 and 1 across a single complexity dimension.

Maxwell-TCS can be used in a variety of downstreaming tasks such as prompt difficulty prediction, dataset annotation, dataset augmentation and more.

• Developed by: Shreyan C | BUD Ecosystem Inc.
• Model type: Bidirectional Encoder Representations from Transformers, based on the ModernBERT-Large architecture.
• Language(s) (NLP): English (en)
• License: Apache License, Version 2.0
• Code: https://github.com/thethinkmachine/Maxwell-Task-Complexity-Scorer
• Finetuned from model: ModernBERT-Large

Recommendations

To reproduce Evol-Complexity scores:

To reproduce the original Evol-Complexity scores, it is recommended to use the following transformation over the predicted scores, $$\text{S}_{predicted} \times (6 - 1) + 1$$ This formula effectively converts normalized scores back to the continuous range [1,6]. Now all that remains is to ensure that the model's predictions correspond to the original discrete scoring scale, which can be achieved simply by rounding the transformed scores to the nearest integer.

To use a different scaling factor:

If you wish to use a different scaling factor, simply replace the values in the formula above with the desired range. For example, to scale the scores to the range [max,min], use the denormalization formula, $$\text{S}_{predicted} \times (\text{max} - \text{min}) + \text{min}$$ Note: Scaling factors are arbitrary and can be adjusted as needed. Ordering of the scores is preserved during min-max scaling, so the model should still be able to predict the relative complexity of instructions accurately regardless of the scaling factor used.

Get Started

Use the code below to get started with the model.

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_name = "thethinkmachine/Maxwell-Task-Complexity-Scorer-v0.2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

def get_deita_complexity_score(question: str) -> int:
    inputs = tokenizer(question, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    normalized_pred = outputs.logits.squeeze()
    final_score = normalized_pred * (6 - 1) + 1
    final_score = torch.clamp(final_score, min=1.0, max=6.0)
    final_score = torch.round(final_score)    
    return final_score.item()

def get_scaled_complexity_score(question: str) -> float:
    inputs = tokenizer(question, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    normalized_pred = outputs.logits.squeeze()
    final_score = normalized_pred * (max_score - min_score) + min_score
    final_score = torch.clamp(final_score, min=min_score, max=max_score)
    final_score = final_score.item()
    return round(final_score, 2)

query = "Is learning equivalent to decreasing local entropy?"
max_score = 100
min_score = 0

print("DEITA Evol-Complexity Score:", get_deita_complexity_score(query))
print("Scaled Complexity Score:", get_scaled_complexity_score(query))

Training Details

Training Data

We use BhabhaAI/DEITA-Complexity 'deita'set for training the model. The dataset contains 66.5K diverse English instructions along with their complexity scores computed using the DEITA-Evol-Complexity scoring scheme which uses an LLM-judge to rank a sextuple containing 1 seed + 5 progressively complexified (evolved) instructions based on their complexity & difficulty. The scheme assigns scores within [1, 6] range, with 1 being the least complex and 6 being the most complex.

However, the training dataset used was observed to have instruction-score pairs across a diversity of scores within the range [0,9]. We suspect that this range includes scoring errors, as anomalous scores (0, 7, 8, 9) account for less than 1% of the total instructions.

The distribution of scores within the dataset is as follows:

To mitigate the influence of these outliers, instructions with scores > 6 and < 1 were pruned from the dataset. It may also be observed that the dataset is not perfectly balanced, as approximately 37.4% of the instructions have a mode score of 5, indicating a higher concentration around this score. Min-Max scaling was applied to normalize the scores within the range [0,1] and preserve the ordering of complexity relationships.

Bias, Risks, and Limitations

From a qualitative standpoint, the training dataset is strongly biased towards general English instructions, which may not help the model generalize well to other languages. Furthermore, the dataset suffers from the scarcity of instructions from technical tasks such as coding, mathematics, etc. As such, the model is presently speculated to not perform well on instructions from these technical domains, further analysis is still required.

You are advised to use the model keeping these factors in mind.

Training Procedure

Training Hyperparameters

• Base Model: ModernBERT-Large
• Task: Sequence Classification
• Training regime: FP32 Non-Mixed Precision
• # Training tokens: 50.3 million tokens
• Tokens-Per-Parameter Ratio: ~3.5 (on 14.4 million trainable parameters)
• Batch size: 8
• Max epochs: 3
• Learning rate: 5e-5
• Optimizer: AdamW
• Warmup steps ratio: 0.1
• Weight decay: 0.01
• Max sequence length: 512

LoRA Hyperparameters

• LoRA Target Modules: "attn.Wo", "attn.Wqkv", "mlp.Wi", "mlp.Wo"
• LoRA Rank: 32
• LoRA Alpha: 64
• LoRA Dropout: 0.1
• LoRA Initialization: PISSA

Environmental Impact

CO2 Emissions

Experiments were conducted using Google Cloud Platform in region asia-south1, which has a carbon efficiency of 0.92 kgCO2eq/kWh. A cumulative of 13.24 hours of computation was performed on hardware of type L4 (TDP of 72W).

Total emissions are estimated to be 0.87 kgCO2eq of which 100% was directly offset by the cloud provider.

Carbon emissions were estimated using the Machine Learning Impact calculator presented in Lacoste et al. (2019).

• Hardware Type: 1xL4 GPU
• Hours used: 16 hours
• Cloud Provider: Google Cloud Platform
• Compute Region: South Asia
• Carbon Emitted: 0.87 kgCO2eq (fully offset by provider)

Author

Shreyan C @thethinkmachine/ @thinkingmachines

For any queries, suggestions or feedback, please contact Shreyan C at shreyan(at)bud(dot)studio / shreyan(dot)chaubey(at)gmail(dot)com.

References

• [2312.15685] What Makes Good Data for Alignment? A Comprehensive Study of Automatic Data Selection in Instruction Tuning
• [2404.02948] PiSSA: Principal Singular Values and Singular Vectors Adaptation of Large Language Models
• DEITA-Complexity
• ModernBERT-Large