Master Thesis: Long Context Extension for Public Sector Document Retrieval

As Sweden's national center for applied AI, we're on a mission to accelerate the use of AI to benefit our society, our competitiveness, and everyone living in Sweden. We drive impactful initiatives in areas such as healthcare, energy, and public services while pushing the boundaries of AI research in fields such as natural language processing, machine learning and AI security. Join us in harnessing the untapped value of AI to drive innovation and create sustainable value for Sweden. We are now looking for a master thesis student to join our team.

Introduction Extending the context length of dense LLMs beyond their pre-training limits is essential for practical deployment in public administration. Context extension techniques like YaRN or LongRoPE show promise but require rigorous stress-testing on domain-specific, low-resource administrative languages to ensure precision in Retrieval-Augmented Generation (RAG) tasks.

Project Background and Problem Statement AI Sweden is developing context length extensions for the Prelude 9B model. Simultaneously, Skatteverket requires "Labb/experiment" Proof-of-Concepts (POCs) for utilizing LLMs with their internal documents (Väg 1). This thesis investigates:

• How efficiently can Prelude 9B be extended to 32k or 64k (or longer) tokens for Swedish administrative RAG tasks without catastrophic forgetting of its core reasoning and language modeling capabilities?

Outline The goal is to extend and evaluate the context window of Prelude 9B for real-world utility.

• Literature study: Investigate RoPE scaling, YaRN, LongRoPE techniques, and Needle-in-a-Haystack (NIAH) evaluations.
• Implementation: Apply continued pre-training or SFT on a high-quality Swedish long-context dataset to extend Prelude's window to 32k/64k tokens.
• Evaluation: Deploy the model in a localized RAG environment simulating Skatteverket's "Väg 1" lab setting and evaluate precision, recall, and MFU (Model FLOPs Utilization).

Who We’re Looking For We are seeking curious, self-driven MSc students eager to work at the frontier of open-weight European AI research (LLMs). You thrive on empirical discovery, design rigorous experiments, and let data challenge your assumptions.

• Ongoing Master’s studies in Computer Science, Data Science, Machine Learning, Engineering Physics, or a related quantitative field.
• Proficiency in Python and hands-on experience with modern deep learning frameworks (PyTorch, Hugging Face ecosystem).
• Familiarity with LLM post-training alignment (e.g., SFT, DPO, RLHF/RLVR) or context-extension, alongside comfort running distributed GPU training in Linux/HPC environments.

At AI Sweden, we are committed to building diverse and inclusive teams. Some positions may be subject to export control regulations, which means that specific requirements may apply. Why should you do your thesis with AI Sweden? Doing your thesis at AI Sweden means working alongside leading AI scientists and change leaders. AI Sweden is Sweden’s National Center for AI, we drive research questions that have both a long shelf-life and are widely applicable to Swedish industry and the public sector. We aim for publications at the most competitive venues and celebrate a culture of research excellence.

As an organization, we’re uniquely positioned at the sweet spot of governmental influence and startup agility. Small enough to stay adaptive and have fun but backed by and in close contact with both the government, academia and private and public sector.

Practical details Location: Hybrid (Gothenburg / Stockholm) or Remote.

Application Deadline: 2026-11-20 (rolling selection – position may be filled earlier). Start Date : January 2027

Contact If you have any questions or thoughts, don’t hesitate to contact:

Birger Moëll, Senior Research Scientist

Niclas Hertzberg, AI Engineer

AI Sweden does not accept unsolicited support and kindly ask not to be contacted by any advertisement agents, recruitment agencies or manning companies. References

[1] Peng, B., et al., "YaRN: Efficient Context Window Extension of Large Language Models," ICLR 2024. [2] Ding, Y., et al., "LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens," arXiv 2024.